Browse Source

Update chat_service: add new openrouter model

StrayWarrior 4 weeks ago
parent
commit
7934e166ba
1 changed files with 9 additions and 1 deletions
  1. 9 1
      pqai_agent/chat_service.py

+ 9 - 1
pqai_agent/chat_service.py

@@ -15,7 +15,7 @@ from pqai_agent.logging_service import logger
 import cozepy
 import cozepy
 from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
 from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
 import time
 import time
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI, http_client
 
 
 COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
 COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
 COZE_CN_BASE_URL = 'https://api.coze.cn'
 COZE_CN_BASE_URL = 'https://api.coze.cn'
@@ -34,6 +34,9 @@ OPENAI_API_TOKEN = 'sk-proj-6LsybsZSinbMIUzqttDt8LxmNbi-i6lEq-AUMzBhCr3jS8sme9AG
 OPENAI_BASE_URL = 'https://api.openai.com/v1'
 OPENAI_BASE_URL = 'https://api.openai.com/v1'
 OPENAI_MODEL_GPT_4o = 'gpt-4o'
 OPENAI_MODEL_GPT_4o = 'gpt-4o'
 OPENAI_MODEL_GPT_4o_mini = 'gpt-4o-mini'
 OPENAI_MODEL_GPT_4o_mini = 'gpt-4o-mini'
+OPENROUTER_API_TOKEN = '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'
+OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1/'
+OPENROUTER_MODEL_CLAUDE_3_7_SONNET = 'anthropic/claude-3.7-sonnet'
 
 
 class ChatServiceType(Enum):
 class ChatServiceType(Enum):
     OPENAI_COMPATIBLE = auto()
     OPENAI_COMPATIBLE = auto()
@@ -55,6 +58,9 @@ class OpenAICompatible:
             OPENAI_MODEL_GPT_4o_mini,
             OPENAI_MODEL_GPT_4o_mini,
             OPENAI_MODEL_GPT_4o
             OPENAI_MODEL_GPT_4o
         ]
         ]
+        openrouter_models = [
+            OPENROUTER_MODEL_CLAUDE_3_7_SONNET,
+        ]
         if model_name in volcengine_models:
         if model_name in volcengine_models:
             llm_client = OpenAI(api_key=VOLCENGINE_API_TOKEN, base_url=VOLCENGINE_BASE_URL, **kwargs)
             llm_client = OpenAI(api_key=VOLCENGINE_API_TOKEN, base_url=VOLCENGINE_BASE_URL, **kwargs)
         elif model_name in deepseek_models:
         elif model_name in deepseek_models:
@@ -68,6 +74,8 @@ class OpenAICompatible:
                 )
                 )
                 kwargs['http_client'] = http_client
                 kwargs['http_client'] = http_client
             llm_client = OpenAI(api_key=OPENAI_API_TOKEN, base_url=OPENAI_BASE_URL, **kwargs)
             llm_client = OpenAI(api_key=OPENAI_API_TOKEN, base_url=OPENAI_BASE_URL, **kwargs)
+        elif model_name in openrouter_models:
+            llm_client = OpenAI(api_key=OPENROUTER_API_TOKEN, base_url=OPENROUTER_BASE_URL, **kwargs)
         else:
         else:
             raise Exception("Unsupported model: %s" % model_name)
             raise Exception("Unsupported model: %s" % model_name)
         return llm_client
         return llm_client