feat: Add openrouter backend (#55) (#56)

* Update router to determine api provider
* Add backend for openrouter
* Update webui app
* Update README
* Update requirements for new openai python package

Author: dexhunter

Diff for: README.md

@@ -99,6 +99,8 @@ Set up your OpenAI (or Anthropic) API key:
 export OPENAI_API_KEY=<your API key>
 # or
 export ANTHROPIC_API_KEY=<your API key>
+# or
+export OPENROUTER_API_KEY=<your API key>
 ```
 
 To run AIDE:

Diff for: aide/backend/__init__.py

@@ -1,5 +1,26 @@
-from . import backend_anthropic, backend_openai
+from . import backend_anthropic, backend_openai, backend_openrouter
 from .utils import FunctionSpec, OutputType, PromptType, compile_prompt_to_md
+import re
+import logging
+
+logger = logging.getLogger("aide")
+
+
+def determine_provider(model: str) -> str:
+    if model.startswith("gpt-") or re.match(r"^o\d", model):
+        return "openai"
+    elif model.startswith("claude-"):
+        return "anthropic"
+    # all other models are handle by openrouter
+    else:
+        return "openrouter"
+
+
+provider_to_query_func = {
+    "openai": backend_openai.query,
+    "anthropic": backend_anthropic.query,
+    "openrouter": backend_openrouter.query,
+}
 
 
 def query(
@@ -35,13 +56,14 @@ def query(
 
     # Handle models with beta limitations
     # ref: https://platform.openai.com/docs/guides/reasoning/beta-limitations
-    if model.startswith("o1"):
+    if re.match(r"^o\d", model):
         if system_message:
             user_message = system_message
         system_message = None
         model_kwargs["temperature"] = 1
 
-    query_func = backend_anthropic.query if "claude-" in model else backend_openai.query
+    provider = determine_provider(model)
+    query_func = provider_to_query_func[provider]
     output, req_time, in_tok_count, out_tok_count, info = query_func(
         system_message=compile_prompt_to_md(system_message) if system_message else None,
         user_message=compile_prompt_to_md(user_message) if user_message else None,

Diff for: aide/backend/backend_openrouter.py

@@ -0,0 +1,81 @@
+"""Backend for OpenRouter API"""
+
+import logging
+import os
+import time
+
+from funcy import notnone, once, select_values
+import openai
+
+from .utils import FunctionSpec, OutputType, backoff_create
+
+logger = logging.getLogger("aide")
+
+_client: openai.OpenAI = None  # type: ignore
+
+OPENAI_TIMEOUT_EXCEPTIONS = (
+    openai.RateLimitError,
+    openai.APIConnectionError,
+    openai.APITimeoutError,
+    openai.InternalServerError,
+)
+
+
+@once
+def _setup_openrouter_client():
+    global _client
+    _client = openai.OpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=os.getenv("OPENROUTER_API_KEY"),
+        max_retries=0,
+    )
+
+
+def query(
+    system_message: str | None,
+    user_message: str | None,
+    func_spec: FunctionSpec | None = None,
+    **model_kwargs,
+) -> tuple[OutputType, float, int, int, dict]:
+    _setup_openrouter_client()
+    filtered_kwargs: dict = select_values(notnone, model_kwargs)  # type: ignore
+
+    if func_spec is not None:
+        raise NotImplementedError(
+            "We are not supporting function calling in OpenRouter for now."
+        )
+
+    # in case some backends dont support system roles, just convert everything to user
+    messages = [
+        {"role": "user", "content": message}
+        for message in [system_message, user_message]
+        if message
+    ]
+
+    t0 = time.time()
+    completion = backoff_create(
+        _client.chat.completions.create,
+        OPENAI_TIMEOUT_EXCEPTIONS,
+        messages=messages,
+        extra_body={
+            "provider": {
+                "order": ["Fireworks"],
+                "ignore": ["Together", "DeepInfra", "Hyperbolic"],
+            },
+        },
+        **filtered_kwargs,
+    )
+    req_time = time.time() - t0
+
+    output = completion.choices[0].message.content
+
+    in_tokens = completion.usage.prompt_tokens
+    out_tokens = completion.usage.completion_tokens
+
+    info = {
+        "system_fingerprint": completion.system_fingerprint,
+        "model": completion.model,
+        "created": completion.created,
+    }
+
+    return output, req_time, in_tokens, out_tokens, info

Diff for: aide/webui/app.py

@@ -51,6 +51,7 @@ def load_env_variables():
         return {
             "openai_key": os.getenv("OPENAI_API_KEY", ""),
             "anthropic_key": os.getenv("ANTHROPIC_API_KEY", ""),
+            "openrouter_key": os.getenv("OPENROUTER_API_KEY", ""),
         }
 
     @staticmethod
@@ -127,9 +128,20 @@ def render_sidebar(self):
                 type="password",
                 label_visibility="collapsed",
             )
+            st.markdown(
+                "<p style='text-align: center;'>OpenRouter API Key</p>",
+                unsafe_allow_html=True,
+            )
+            openrouter_key = st.text_input(
+                "OpenRouter API Key",
+                value=self.env_vars["openrouter_key"],
+                type="password",
+                label_visibility="collapsed",
+            )
             if st.button("Save API Keys", use_container_width=True):
                 st.session_state.openai_key = openai_key
                 st.session_state.anthropic_key = anthropic_key
+                st.session_state.openrouter_key = openrouter_key
                 st.success("API keys saved!")
 
     def render_input_section(self, results_col):
@@ -340,6 +352,8 @@ def set_api_keys():
             os.environ["OPENAI_API_KEY"] = st.session_state.openai_key
         if st.session_state.get("anthropic_key"):
             os.environ["ANTHROPIC_API_KEY"] = st.session_state.anthropic_key
+        if st.session_state.get("openrouter_key"):
+            os.environ["OPENROUTER_API_KEY"] = st.session_state.openrouter_key
 
     def prepare_input_directory(self, files):
         """

Diff for: requirements.txt

@@ -4,7 +4,7 @@ funcy==2.0
 humanize==4.8.0
 jsonschema==4.19.2
 numpy==1.26.2
-openai>=1.3.5
+openai>=1.69.0
 anthropic>=0.20.0
 pandas==2.1.4
 pytest==7.4.3