diff --git a/mcp_server.py b/mcp_server.py index 9f95d8c6..a04bb43d 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -55,19 +55,40 @@ HTTP_CLIENT = httpx.AsyncClient(timeout=60.0) # ==================== Agent系统配置 ==================== -# Kimi 配置 - 用于计划制定和深度推理 -KIMI_CONFIG = { - "api_key": "sk-TzB4VYJfCoXGcGrGMiewukVRzjuDsbVCkaZXi2LvkS8s60E5", - "base_url": "https://api.moonshot.cn/v1", - "model": "kimi-k2-thinking", # 思考模型 +# ==================== 多模型配置 ==================== + +# 模型配置字典(支持动态切换) +MODEL_CONFIGS = { + "kimi-k2": { + "api_key": "sk-TzB4VYJfCoXGcGrGMiewukVRzjuDsbVCkaZXi2LvkS8s60E5", + "base_url": "https://api.moonshot.cn/v1", + "model": "moonshot-v1-8k", # 快速模型 + }, + "kimi-k2-thinking": { + "api_key": "sk-TzB4VYJfCoXGcGrGMiewukVRzjuDsbVCkaZXi2LvkS8s60E5", + "base_url": "https://api.moonshot.cn/v1", + "model": "kimi-k2-thinking", # 深度思考模型 + }, + "glm-4.6": { + "api_key": "", # 需要配置智谱AI密钥 + "base_url": "https://open.bigmodel.cn/api/paas/v4", + "model": "glm-4", + }, + "deepmoney": { + "api_key": "", # 空值 + "base_url": "http://111.62.35.50:8000/v1", + "model": "deepmoney", + }, + "gemini-3": { + "api_key": "", # 需要配置Google API密钥 + "base_url": "https://generativelanguage.googleapis.com/v1", + "model": "gemini-pro", + }, } -# DeepMoney 配置 - 用于新闻总结 -DEEPMONEY_CONFIG = { - "api_key": "", # 空值 - "base_url": "http://111.62.35.50:8000/v1", - "model": "deepmoney", -} +# 保持向后兼容的配置(默认使用 kimi-k2-thinking) +KIMI_CONFIG = MODEL_CONFIGS["kimi-k2-thinking"] +DEEPMONEY_CONFIG = MODEL_CONFIGS["deepmoney"] # ==================== MCP协议数据模型 ==================== @@ -143,6 +164,8 @@ class AgentChatRequest(BaseModel): user_avatar: Optional[str] = None # 用户头像URL subscription_type: Optional[str] = None # 用户订阅类型(free/pro/max) session_id: Optional[str] = None # 会话ID(如果为空则创建新会话) + model: Optional[str] = "kimi-k2-thinking" # 选择的模型(kimi-k2, kimi-k2-thinking, glm-4.6, deepmoney, gemini-3) + tools: Optional[List[str]] = None # 选择的工具列表(工具名称数组,如果为None则使用全部工具) # ==================== MCP工具定义 ==================== @@ -1579,6 +1602,7 @@ class MCPAgentIntegrated: user_nickname: str = None, user_avatar: str = None, cookies: dict = None, + model_config: dict = None, # 新增:动态模型配置 ) -> AsyncGenerator[str, None]: """主流程(流式输出)- 逐步返回执行结果""" logger.info(f"[Agent Stream] 处理查询: {user_query}") @@ -1586,11 +1610,24 @@ class MCPAgentIntegrated: # 将 cookies 存储为实例属性,供工具调用时使用 self.cookies = cookies or {} + # 如果传入了自定义模型配置,使用自定义配置,否则使用默认的 Kimi + if model_config: + planning_client = OpenAI( + api_key=model_config["api_key"], + base_url=model_config["base_url"], + ) + planning_model = model_config["model"] + logger.info(f"[Agent Stream] 使用自定义模型: {planning_model}") + else: + planning_client = self.kimi_client + planning_model = self.kimi_model + logger.info(f"[Agent Stream] 使用默认模型: {planning_model}") + try: # 发送开始事件 yield self._format_sse("status", {"stage": "start", "message": "开始处理查询"}) - # 阶段1: Kimi 制定计划(流式,带 DeepMoney 备选) + # 阶段1: 使用选中的模型制定计划(流式,带 DeepMoney 备选) yield self._format_sse("status", {"stage": "planning", "message": "正在制定执行计划..."}) messages = [ @@ -1603,9 +1640,9 @@ class MCPAgentIntegrated: use_fallback = False try: - # 尝试使用 Kimi 流式 API - stream = self.kimi_client.chat.completions.create( - model=self.kimi_model, + # 尝试使用选中的模型流式 API + stream = planning_client.chat.completions.create( + model=planning_model, messages=messages, temperature=1.0, max_tokens=16000, @@ -2165,11 +2202,12 @@ async def agent_chat_stream(chat_request: AgentChatRequest, request: Request): except Exception as e: logger.error(f"[ES] 保存用户消息失败: {e}") - # 获取工具列表 - tools = [tool.dict() for tool in TOOLS] + # ==================== 动态工具过滤 ==================== + # 获取所有可用工具 + all_tools = [tool.dict() for tool in TOOLS] # 添加特殊工具:summarize_news - tools.append({ + all_tools.append({ "name": "summarize_news", "description": "使用 DeepMoney 模型总结新闻数据,提取关键信息", "parameters": { @@ -2188,6 +2226,21 @@ async def agent_chat_stream(chat_request: AgentChatRequest, request: Request): } }) + # 如果用户指定了工具列表,则进行过滤 + if chat_request.tools is not None and len(chat_request.tools) > 0: + selected_tool_names = set(chat_request.tools) + tools = [tool for tool in all_tools if tool["name"] in selected_tool_names] + logger.info(f"[工具过滤] 用户选择了 {len(tools)}/{len(all_tools)} 个工具: {selected_tool_names}") + else: + # 默认使用全部工具 + tools = all_tools + logger.info(f"[工具过滤] 使用全部 {len(tools)} 个工具") + + # ==================== 动态模型选择 ==================== + selected_model = chat_request.model or "kimi-k2-thinking" + model_config = MODEL_CONFIGS.get(selected_model, MODEL_CONFIGS["kimi-k2-thinking"]) + logger.info(f"[模型选择] 使用模型: {selected_model} ({model_config['model']})") + # 返回流式响应 return StreamingResponse( agent.process_query_stream( @@ -2199,6 +2252,7 @@ async def agent_chat_stream(chat_request: AgentChatRequest, request: Request): user_nickname=chat_request.user_nickname, user_avatar=chat_request.user_avatar, cookies=cookies, # 传递 cookies 用于认证 API 调用 + model_config=model_config, # 传递选中的模型配置 ), media_type="text/event-stream", headers={