# Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Requirement: Any integration or derivative work must explicitly attribute # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its # documentation or user interface, as detailed in the NOTICE file. """ LLM管理器 - 管理多个专用LLM实例 """ from mcp_scan.utils import config from mcp_scan.utils.llm import LLM from mcp_scan.utils.loging import logger class LLMManager: """管理多个专用LLM实例，支持不同用途的模型配置""" # 预定义的模型配置（从环境变量读取） DEFAULT_CONFIGS = { "default": { "model": config.DEFAULT_MODEL, "base_url": config.DEFAULT_BASE_URL, # context_window 用于 agent 的上下文压缩判定。 "context_window": config.DEFAULT_MODEL_CONTEXT_WINDOW, "description": "默认模型", }, "thinking": { "model": config.THINKING_MODEL, "base_url": config.THINKING_BASE_URL, "api_key": config.THINKING_API_KEY, # 可选，为 None 时使用主 API Key "context_window": config.THINKING_MODEL_CONTEXT_WINDOW, "description": "专门用于深度思考和推理的模型", }, "coding": { "model": config.CODING_MODEL, "base_url": config.CODING_BASE_URL, "api_key": config.CODING_API_KEY, # 可选，为 None 时使用主 API Key "context_window": config.CODING_MODEL_CONTEXT_WINDOW, "description": "专门用于代码生成和分析的模型", }, "fast": { "model": config.FAST_MODEL, "base_url": config.FAST_BASE_URL, "api_key": config.FAST_API_KEY, # 可选，为 None 时使用主 API Key "context_window": config.FAST_MODEL_CONTEXT_WINDOW, "description": "用于快速响应的轻量级模型", }, } def __init__(self, api_key: str, base_url: str = None): """ 初始化LLM管理器 Args: api_key: 主 API 密钥（作为默认值） base_url: 主 API 基础URL（作为默认值，如果不提供则使用 config.DEFAULT_BASE_URL） """ self.default_api_key = api_key self.default_base_url = base_url or config.DEFAULT_BASE_URL self._llm_instances: dict[str, LLM] = {} self._custom_configs: dict[str, dict] = {} def configure( self, purpose: str, model: str, base_url: str | None = None, api_key: str | None = None, ) -> None: """ 配置特定用途的模型 Args: purpose: 用途标识（如 "thinking", "coding"） model: 模型名称 base_url: API 基础 URL（可选，不提供则使用默认） api_key: API 密钥（可选，不提供则使用默认） """ self._custom_configs[purpose] = { "model": model, "base_url": base_url, "api_key": api_key, } # 清除已有实例，下次获取时重新创建 if purpose in self._llm_instances: del self._llm_instances[purpose] logger.info(f"Configured LLM for purpose '{purpose}': {model}") def get_llm(self, purpose: str = "default") -> LLM | None: """ 获取指定用途的LLM实例 Args: purpose: 用途标识 Returns: LLM实例，如果未配置返回None """ # 如果已有实例，直接返回 if purpose in self._llm_instances: return self._llm_instances[purpose] # 查找配置 llm_config = None if purpose in self._custom_configs: llm_config = self._custom_configs[purpose].copy() elif purpose in self.DEFAULT_CONFIGS: llm_config = self.DEFAULT_CONFIGS[purpose].copy() if llm_config is None: logger.warning(f"No configuration found for purpose '{purpose}'") return None # 确定使用的 API Key 和 Base URL # 优先级：配置中指定的 > 默认配置中的 > 主配置 api_key = llm_config.get("api_key") or self.default_api_key base_url = llm_config.get("base_url") or self.default_base_url context_window = llm_config.get("context_window") # 创建新实例 try: llm = LLM( model=llm_config["model"], api_key=api_key, base_url=base_url, context_window=context_window, ) self._llm_instances[purpose] = llm logger.info( f"Created LLM instance for purpose '{purpose}': {llm_config['model']} @ {base_url}" ) return llm except Exception as e: logger.error(f"Failed to create LLM for purpose '{purpose}': {e}") return None def get_specialized_llms(self, purposes: list[str] = None) -> dict[str, LLM]: """ 批量获取多个专用LLM实例 Args: purposes: 用途列表，如果为None则返回所有已配置的 Returns: 用途到LLM实例的字典 """ if purposes is None: # 获取所有已配置的 purposes = list( set(list(self._custom_configs.keys()) + list(self.DEFAULT_CONFIGS.keys())) ) result = {} for purpose in purposes: llm = self.get_llm(purpose) if llm: result[purpose] = llm return result def list_available_purposes(self) -> list[tuple[str, str]]: """ 列出所有可用的用途及其描述 Returns: (用途, 描述) 的列表 """ purposes = [] # 自定义配置 for purpose, llm_config in self._custom_configs.items(): model = llm_config["model"] base_url = llm_config.get("base_url") or self.default_base_url purposes.append((purpose, f"Custom: {model} @ {base_url}")) # 默认配置 for purpose, llm_config in self.DEFAULT_CONFIGS.items(): if purpose not in self._custom_configs: model = llm_config["model"] base_url = llm_config.get("base_url") or self.default_base_url desc = llm_config["description"] purposes.append((purpose, f"{desc} ({model} @ {base_url})")) return purposes # 全局单例（可选） _global_manager: LLMManager | None = None def get_global_manager() -> LLMManager | None: """获取全局LLM管理器实例""" return _global_manager def init_global_manager(api_key: str, base_url: str = "https://openrouter.ai/api/v1") -> LLMManager: """ 初始化全局LLM管理器 Args: api_key: API密钥 base_url: API基础URL Returns: LLM管理器实例 """ global _global_manager _global_manager = LLMManager(api_key, base_url) return _global_manager if __name__ == "__main__": import os # 使用示例 - 需要设置环境变量 OPENROUTER_API_KEY api_key = os.environ.get("OPENROUTER_API_KEY") if not api_key: print("请设置环境变量 OPENROUTER_API_KEY") exit(1) manager = LLMManager(api_key=api_key, base_url="https://openrouter.ai/api/v1") # 自定义配置（可以指定独立的 base_url） manager.configure( "thinking", "deepseek/deepseek-reasoner", base_url="https://api.deepseek.com/v1", # 可选：使用专门的 base_url ) manager.configure( "coding", "anthropic/claude-3.5-sonnet", # 不指定 base_url，使用默认值 ) # 获取专用LLM thinking_llm = manager.get_llm("thinking") coding_llm = manager.get_llm("coding") # 批量获取 specialized_llms = manager.get_specialized_llms(["thinking", "coding"]) print("Available purposes:") for purpose, desc in manager.list_available_purposes(): print(f" - {purpose}: {desc}")