浏览代码

Update chat_service: add new openrouter model

StrayWarrior 4 周之前
父节点
当前提交
7934e166ba
共有 1 个文件被更改,包括 9 次插入1 次删除
  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
 from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
 import time
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI, http_client
 
 COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
 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_MODEL_GPT_4o = 'gpt-4o'
 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):
     OPENAI_COMPATIBLE = auto()
@@ -55,6 +58,9 @@ class OpenAICompatible:
             OPENAI_MODEL_GPT_4o_mini,
             OPENAI_MODEL_GPT_4o
         ]
+        openrouter_models = [
+            OPENROUTER_MODEL_CLAUDE_3_7_SONNET,
+        ]
         if model_name in volcengine_models:
             llm_client = OpenAI(api_key=VOLCENGINE_API_TOKEN, base_url=VOLCENGINE_BASE_URL, **kwargs)
         elif model_name in deepseek_models:
@@ -68,6 +74,8 @@ class OpenAICompatible:
                 )
                 kwargs['http_client'] = http_client
             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:
             raise Exception("Unsupported model: %s" % model_name)
         return llm_client