Update agentverse/llms/openai.py

agentverse/llms/openai.py

@@ -17,6 +17,9 @@ from . import llm_registry
 from .base import BaseChatModel, BaseCompletionModel, BaseModelArgs
 from .utils.jsonrepair import JsonRepair
 
+openai.api_base = "https://openrouter.ai/api/v1"
+openai.api_key = os.getenv("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")
+
 try:
     import openai
     from openai.error import OpenAIError
@@ -44,7 +47,7 @@ else:
 
 
 class OpenAIChatArgs(BaseModelArgs):
-    model: str = Field(default="gpt-3.5-turbo")
+    model: str = Field(default="openai/gpt-3.5-turbo")
     deployment_id: str = Field(default=None)
     max_tokens: int = Field(default=2048)
     temperature: float = Field(default=1.0)
@@ -93,9 +96,9 @@ class OpenAIChatArgs(BaseModelArgs):
 #         )
 
 
-@llm_registry.register("gpt-35-turbo")
-@llm_registry.register("gpt-3.5-turbo")
-@llm_registry.register("gpt-4")
+@llm_registry.register("openai/gpt-3.5-turbo")
+@llm_registry.register("openai/gpt-3.5-turbo-16k")
+@llm_registry.register("openai/gpt-4")
 class OpenAIChat(BaseChatModel):
     args: OpenAIChatArgs = Field(default_factory=OpenAIChatArgs)
 
@@ -294,23 +297,23 @@ class OpenAIChat(BaseChatModel):
 
     def get_spend(self) -> int:
         input_cost_map = {
-            "gpt-3.5-turbo": 0.0015,
-            "gpt-3.5-turbo-16k": 0.003,
-            "gpt-3.5-turbo-0613": 0.0015,
-            "gpt-3.5-turbo-16k-0613": 0.003,
-            "gpt-4": 0.03,
-            "gpt-4-0613": 0.03,
-            "gpt-4-32k": 0.06,
+            "openai/gpt-3.5-turbo": 0.0015,
+            "openai/gpt-3.5-turbo": 0.003,
+            "openai/gpt-4": 0.0015,
+            "openai/gpt-3.5-turbo-16k": 0.003,
+            "openai/gpt-4": 0.03,
+            "openai/gpt-4": 0.03,
+            "openai/gpt-4": 0.06,
         }
 
         output_cost_map = {
-            "gpt-3.5-turbo": 0.002,
-            "gpt-3.5-turbo-16k": 0.004,
-            "gpt-3.5-turbo-0613": 0.002,
-            "gpt-3.5-turbo-16k-0613": 0.004,
-            "gpt-4": 0.06,
-            "gpt-4-0613": 0.06,
-            "gpt-4-32k": 0.12,
+            "openai/gpt-3.5-turbo": 0.002,
+            "openai/gpt-3.5-turbo-16k": 0.004,
+            "openai/gpt-3.5-turbo-16k": 0.002,
+            "openai/gpt-3.5-turbo-16k": 0.004,
+            "openai/gpt-4": 0.06,
+            "openai/gpt-4": 0.06,
+            "openai/gpt-4": 0.12,
         }
 
         model = self.args.model