Add OpenRouter API Support

ai_lab_repo.py

@@ -705,6 +705,7 @@ def __init__(self): pass
     else: parser.construct_agentRxiv = False
     if 'agentrxiv-papers' in agentlab_data: parser.agentrxiv_papers = agentlab_data["agentrxiv-papers"]
     else:  parser.agentrxiv_papers = 5
+    if 'openrouter-api-key' in agentlab_data: parser.openrouter_api_key = agentlab_data['openrouter-api-key']
 
     if 'lab-index' in agentlab_data: parser.lab_index = agentlab_data["lab-index"]
     else: parser.lab_index = 0
@@ -741,6 +742,9 @@ def __init__(self): pass
 
     api_key = (os.getenv('OPENAI_API_KEY') or args.api_key) if (hasattr(args, 'api_key') or os.getenv('OPENAI_API_KEY')) else None
     deepseek_api_key = (os.getenv('DEEPSEEK_API_KEY') or args.deepseek_api_key) if (hasattr(args, 'deepseek_api_key') or os.getenv('DEEPSEEK_API_KEY')) else None
+    openrouter_api_key = (os.getenv('OPENROUTER_API_KEY') or getattr(args, 'openrouter_api_key', None))
+    if openrouter_api_key is not None and os.getenv('OPENROUTER_API_KEY') is None:
+        os.environ['OPENROUTER_API_KEY'] = openrouter_api_key
     if api_key is not None and os.getenv('OPENAI_API_KEY') is None: os.environ["OPENAI_API_KEY"] = args.api_key
     if deepseek_api_key is not None and os.getenv('DEEPSEEK_API_KEY') is None: os.environ["DEEPSEEK_API_KEY"] = args.deepseek_api_key
 

experiment_configs/MATH_agentlab.yaml

@@ -7,10 +7,12 @@ research-topic: "Your goal is to design reasoning and prompt engineering techniq
 # Here you can put your OpenAI API key--if you don't have one or OpenAI doesn't work for you, you can also instead use `deepseek-api-key`
 api-key: "OPENAI-API-KEY-HERE"
 # or deepseek-api-key: "DEEPSEEK-API-KEY-HERE"
+openrouter-api-key: "OPENROUTER_API_KEY"
+
 # Agent Laboratory backend
-llm-backend: "o3-mini"
+llm-backend: "openrouter/deepseek/deepseek-chat-v3-0324:free"
 # Literature review backend
-lit-review-backend: "o3-mini"
+lit-review-backend: "openrouter/deepseek/deepseek-chat-v3-0324:free"
 
 # Base language
 language: "English"

inference.py

@@ -3,6 +3,8 @@
 from openai import OpenAI
 import os, anthropic, json
 import google.generativeai as genai
+import requests  # For completeness, in case it's used elsewhere
+from openrouter_inference import query_openrouter_model  # Import the new function
 
 TOKENS_IN = dict()
 TOKENS_OUT = dict()
@@ -34,6 +36,7 @@ def curr_cost_est():
 
 def query_model(model_str, prompt, system_prompt, openai_api_key=None, gemini_api_key=None,  anthropic_api_key=None, tries=5, timeout=5.0, temp=None, print_cost=True, version="1.5"):
     preloaded_api = os.getenv('OPENAI_API_KEY')
+    preloaded_openrouter_api = os.getenv('OPENROUTER_API_KEY')
     if openai_api_key is None and preloaded_api is not None:
         openai_api_key = preloaded_api
     if openai_api_key is None and anthropic_api_key is None:
@@ -47,6 +50,7 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, gemini_ap
         os.environ["GEMINI_API_KEY"] = gemini_api_key
     for _ in range(tries):
         try:
+            answer = None  # ensure answer is always defined
             if model_str == "gpt-4o-mini" or model_str == "gpt4omini" or model_str == "gpt-4omini" or model_str == "gpt4o-mini":
                 model_str = "gpt-4o-mini"
                 messages = [
@@ -186,24 +190,40 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, gemini_ap
                     completion = client.chat.completions.create(
                         model="o1-preview", messages=messages)
                 answer = completion.choices[0].message.content
+            elif model_str.startswith("openrouter/"):
+                openrouter_model = model_str[len("openrouter/"):]  # Remove the prefix
+                answer = query_openrouter_model(
+                    model_str=openrouter_model,
+                    prompt=prompt,
+                    system_prompt=system_prompt,
+                    openrouter_api_key=preloaded_openrouter_api,
+                    temp=temp
+                )
 
             try:
-                if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet", "o1", "o3-mini"]:
-                    encoding = tiktoken.encoding_for_model("gpt-4o")
-                elif model_str in ["deepseek-chat"]:
-                    encoding = tiktoken.encoding_for_model("cl100k_base")
+                # Use a default tokenizer for unknown or OpenRouter models
+                if (
+                    model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet", "o1", "o3-mini"]
+                    or model_str.startswith("openrouter/")
+                    or "deepseek" in model_str
+                ):
+                    encoding = tiktoken.get_encoding("cl100k_base")
                 else:
                     encoding = tiktoken.encoding_for_model(model_str)
                 if model_str not in TOKENS_IN:
                     TOKENS_IN[model_str] = 0
                     TOKENS_OUT[model_str] = 0
                 TOKENS_IN[model_str] += len(encoding.encode(system_prompt + prompt))
-                TOKENS_OUT[model_str] += len(encoding.encode(answer))
+                if answer is not None:
+                    TOKENS_OUT[model_str] += len(encoding.encode(answer))
                 if print_cost:
                     print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost")
             except Exception as e:
                 if print_cost: print(f"Cost approximation has an error? {e}")
-            return answer
+            if answer is not None:
+                return answer
+            else:
+                raise Exception("No answer returned from model.")
         except Exception as e:
             print("Inference Exception:", e)
             time.sleep(timeout)

openrouter_inference.py

@@ -0,0 +1,32 @@
+import os
+import requests
+
+def query_openrouter_model(model_str, prompt, system_prompt, openrouter_api_key=None, temp=None):
+    if openrouter_api_key is None:
+        openrouter_api_key = os.getenv('OPENROUTER_API_KEY')
+    if openrouter_api_key is None:
+        raise Exception("No OpenRouter API key provided")
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt}
+    ]
+    headers = {
+        "Authorization": f"Bearer {openrouter_api_key}",
+        "Content-Type": "application/json"
+    }
+    data = {
+        "model": model_str,  # e.g. "openrouter/phi-3-mini-128k-instruct"
+        "messages": messages,
+    }
+    if temp is not None:
+        data["temperature"] = temp
+    response = requests.post(
+        "https://openrouter.ai/api/v1/chat/completions",
+        headers=headers,
+        json=data,
+        timeout=60
+    )
+    if response.status_code != 200:
+        raise Exception(f"OpenRouter API error: {response.text}")
+    completion = response.json()
+    return completion["choices"][0]["message"]["content"]