Merge pull request #1 from akshaypardhanani/feat/add-support-for-openrouter

feat: Add support for OpenRouter

Author: akshaypardhanani

.gitignore

@@ -46,3 +46,11 @@ MANIFEST
 *.pyc
 __pycache__/
 *~
+
+digest.txt
+uv.lock
+.vscode/launch.json
+amadeusgpt/modules_embedding.pickle
+temp_answer.json
+
+logs/

Makefile

@@ -1,4 +1,4 @@
 export streamlit_app=True
 app:
 
-	streamlit run amadeusgpt/app.py --server.fileWatcherType none --server.maxUploadSize 1000
+	uv run --with streamlit streamlit run amadeusgpt/app.py --server.fileWatcherType none --server.maxUploadSize 1000

amadeusgpt/__init__.py

@@ -4,6 +4,7 @@
 SOURCE CODE: https://github.com/AdaptiveMotorControlLab/AmadeusGPT
 Apache-2.0 license
 """
+import streamlit as st
 
 from matplotlib import pyplot as plt
 

amadeusgpt/analysis_objects/llm.py

@@ -13,6 +13,7 @@
 
 from amadeusgpt.programs.sandbox import Sandbox
 from amadeusgpt.utils import AmadeusLogger, QA_Message, create_qa_message
+from amadeusgpt.utils.openai_adapter import OpenAIAdapter
 
 from .base import AnalysisObject
 
@@ -22,6 +23,7 @@ class LLM(AnalysisObject):
     prices = {
         "gpt-4o": {"input": 5 / 10**6, "output": 15 / 10**6},
         "gpt-4o-mini": {"input": 0.15 / 10**6, "output": 0.6 / 10**6},
+        "thudm/glm-z1-32b:free": {"input": 0, "output": 0},
     }
     total_cost = 0
 
@@ -65,17 +67,17 @@ def connect_gpt_oai_1(self, messages, **kwargs):
         This is routed to openai > 1.0 interfaces
         """
 
-        if self.config.get("use_streamlit", False):
-            if "OPENAI_API_KEY" in os.environ:
-                openai.api_key = os.environ["OPENAI_API_KEY"]
-        else:
-            openai.api_key = os.environ["OPENAI_API_KEY"]
+        # if self.config.get("use_streamlit", False):
+        #     if "OPENAI_API_KEY" in os.environ:
+        #         openai.api_key = os.environ["OPENAI_API_KEY"]
+        # else:
+        #     openai.api_key = os.environ["OPENAI_API_KEY"]
         response = None
         # gpt_model is default to be the cls.gpt_model, which can be easily set
-        gpt_model = self.gpt_model
+        # gpt_model = self.gpt_model
         # in streamlit app, "gpt_model" is set by the text box
 
-        client = OpenAI()
+        client = OpenAIAdapter().get_client()
 
         if self.config.get("use_streamlit", False):
             if "gpt_model" in st.session_state:
@@ -89,7 +91,7 @@ def connect_gpt_oai_1(self, messages, **kwargs):
 
         # the usage was recorded from the last run. However, since we have many LLMs that
         # share the call of this function, we will need to store usage and retrieve them from the database class
-        num_retries = 3
+        num_retries = 1
         for _ in range(num_retries):
             try:
                 json_data = {
@@ -259,7 +261,7 @@ def speak(self, sandbox: Sandbox, image: np.ndarray):
             multi_image_content=multi_image_content,
             in_place=True,
         )
-        response = self.connect_gpt(self.context_window, max_tokens=2000)
+        response = self.connect_gpt(self.context_window, max_tokens=20000)
         text = response.choices[0].message.content.strip()
 
         print("description of the image frame provided")
@@ -319,7 +321,7 @@ def speak(
 
         self.update_history("user", query)
 
-        response = self.connect_gpt(self.context_window, max_tokens=2000)
+        response = self.connect_gpt(self.context_window, max_tokens=20000)
         text = response.choices[0].message.content.strip()
         # need to keep the memory of the answers from LLM
         self.update_history("assistant", text)
@@ -374,7 +376,7 @@ def speak(self, qa_message):
 Can you correct the code? Make sure you only write one function which is the updated function.
 """
         self.update_history("user", query)
-        response = self.connect_gpt(self.context_window, max_tokens=4096)
+        response = self.connect_gpt(self.context_window, max_tokens=20000)
         text = response.choices[0].message.content.strip()
         print(text)
         pattern = r"```python(.*?)```"

amadeusgpt/app.py

@@ -1,9 +1,7 @@
 import os
 import traceback
 
-import streamlit as st
-
-from amadeusgpt import app_utils
+from amadeusgpt import app_utils, st
 from amadeusgpt.utils import validate_openai_api_key
 
 # Set page configuration
@@ -21,6 +19,8 @@ def main():
     if "exist_valid_openai_api_key" not in st.session_state:
         if "OPENAI_API_KEY" in os.environ:
             st.session_state["exist_valid_openai_api_key"] = True
+        elif "OPENROUTER_API_KEY" in os.environ:
+            st.session_state["exist_valid_openai_api_key"] = True
         else:
             st.session_state["exist_valid_openai_api_key"] = False
 
@@ -30,6 +30,8 @@ def valid_api_key():
         print("inside valid api key function")
         if "OPENAI_API_KEY" in os.environ:
             api_token = os.environ["OPENAI_API_KEY"]
+        elif "OPENROUTER_API_KEY" in os.environ:
+            api_token = os.environ["OPENROUTER_API_KEY"]
         else:
             api_token = st.session_state["openAI_token"]
         check_valid = validate_openai_api_key(api_token)

amadeusgpt/configs/Horse_template.yaml

@@ -7,6 +7,7 @@ keypoint_info:
     nose: "nose"
     neck: "neck" 
 llm_info:
+  gpt_model: "thudm/glm-z1-32b:free"
   keep_last_n_messages: 2 
 object_info:
   load_objects_from_disk: false

amadeusgpt/integration_module_hub.py

@@ -1,17 +1,18 @@
 import os
 import pickle
 
-from openai import OpenAI
+from amadeusgpt.utils.openai_adapter import OpenAIAdapter
 from sklearn.metrics.pairwise import cosine_similarity
 
+from amadeusgpt import st
 from amadeusgpt.programs.api_registry import INTEGRATION_API_REGISTRY
 
-client = OpenAI()
-
 
 class IntegrationModuleHub:
     def __init__(self):
         self.amadeus_root = os.path.dirname(os.path.realpath(__file__))
+        self.client = OpenAIAdapter(
+            api_key=st.session_state.get("OPENAI_API_KEY") or st.session_state.get("OPENROUTER_API_KEY") or st.session_state.get("openAI_token")).get_client()
 
     def save_embeddings(self):
         result = {}
@@ -20,7 +21,7 @@ def save_embeddings(self):
             docstring = module_info["description"]
             text = docstring.replace("\n", " ")
             embedding = (
-                client.embeddings.create(input=[text], model=model).data[0].embedding
+                self.client.embeddings.create(input=[text], model=model).data[0].embedding
             )
             result[module_name] = embedding
         if len(result) > 0:
@@ -35,7 +36,7 @@ def match_module(self, query):
         model = "text-embedding-3-small"
 
         query_embedding = (
-            client.embeddings.create(input=[query], model=model).data[0].embedding
+            self.client.embeddings.create(input=[query], model=model).data[0].embedding
         )
 
         if not os.path.exists(

amadeusgpt/utils/__init__.py

@@ -9,6 +9,7 @@
 from amadeusgpt.analysis_objects.event import Event
 from amadeusgpt.logger import AmadeusLogger
 from IPython.display import Markdown, Video, display, HTML
+from amadeusgpt.utils.openai_adapter import OpenAIAdapter
 
 def filter_kwargs_for_function(func, kwargs):
     sig = inspect.signature(func)
@@ -36,13 +37,8 @@ def parse_error_message_from_python():
     return traceback_str
 
 def validate_openai_api_key(key):
-    import openai
-    openai.api_key = key
-    try:
-        openai.models.list()
-        return True
-    except openai.AuthenticationError:
-        return False
+    client = OpenAIAdapter(key)
+    return client.validate()
 
 def flatten_tuple(t):
     """

amadeusgpt/utils/embedding_adapter.py

@@ -0,0 +1,36 @@
+from typing import Union
+from sentence_transformers import SentenceTransformer
+
+class EmbeddingModel:
+    def __init__(self, provider_type: str, model_name: str = "BAAI/bge-m3", openai_model: str = "text-embedding-3-small"):
+        self.provider_type = provider_type
+        if self.provider_type == "openrouter":
+            self.embedding_model = SentenceTransformer(model_name)
+        self.embeddings = self._Embeddings(self.provider_type, self.embedding_model or openai_model)
+
+    class _Embeddings:
+        def __init__(self, provider: str, model: Union[str, SentenceTransformer]):
+            self.provider = provider
+            self.model = model
+        
+        def create(self, input: list[str], model: Union[str, SentenceTransformer]):
+            if self.provider == "openrouter":
+                vecs = self.model.encode(input, normalize_embeddings=True, convert_to_tensor=False)
+                if isinstance(input, str):
+                    vecs = [vecs]  # (1, D)
+
+                # Convert each embedding to a plain Python list
+                vecs = [v.tolist() for v in vecs]
+
+                # Mimic OpenAI's Embedding.create() response
+                return type("FakeResponse", (), {
+                    "data": [
+                        type("EmbeddingData", (), {
+                            "embedding": emb,
+                            "index": idx
+                        })()
+                        for idx, emb in enumerate(vecs)
+                    ]
+                })()
+            else:
+                return self.model.embeddings.create(input=input, model=model)

amadeusgpt/utils/openai_adapter.py

@@ -0,0 +1,65 @@
+import openai
+from amadeusgpt.utils.embedding_adapter import EmbeddingModel
+
+class OpenAIAdapter:
+    _instance = None
+    _initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(OpenAIAdapter, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self, api_key: str = None):
+        if self._initialized:
+            return
+
+        if api_key is None:
+            raise ValueError("API Key is required")
+        
+        self.api_key = api_key
+        self.base_url = "https://openrouter.ai/api/v1"
+        self.client = None
+        self.is_openrouter = False
+        self._initialized = True
+        self._initialize_client()
+        self.client.embeddings = EmbeddingModel(provider_type=self.get_provider_type()).embeddings
+
+    def __call__(self, *args, **kwargs):
+        if self.client is None:
+            raise RuntimeError("Client is not initialized")
+        return self.client
+
+    def _initialize_client(self):
+        try:
+            client = openai.OpenAI(
+                api_key=self.api_key,
+            )
+            client.models.list()
+            self.client = client
+        except openai.AuthenticationError:
+            try:
+                client = openai.OpenAI(
+                    api_key=self.api_key,
+                    base_url=self.base_url,
+                )
+                client.models.list()
+                self.client = client
+                self.is_openrouter = True
+            except openai.AuthenticationError:
+                raise openai.AuthenticationError("Invalid API key")
+
+    def validate(self):
+        return self.client is not None
+
+    def get_provider_type(self):
+        return "openrouter" if self.is_openrouter else "openai"
+
+    def get_client(self):
+        return self.client
+        
+    @classmethod
+    def reset(cls):
+        cls._instance = None
+        cls._initialized = False
+