""" AI接口提供商服务 管理AI接口的配置和调用 """ import uuid from datetime import datetime from typing import List, Dict, Any, Optional from loguru import logger from models.ai_provider import AIProvider from adapters import get_adapter, BaseAdapter, ChatMessage, AdapterResponse from utils.encryption import encrypt_api_key, decrypt_api_key from utils.rate_limiter import rate_limiter class AIProviderService: """ AI接口提供商服务类 负责AI接口的CRUD操作和调用 """ # 缓存适配器实例 _adapter_cache: Dict[str, BaseAdapter] = {} @classmethod async def create_provider( cls, provider_type: str, name: str, model: str, api_key: str = "", base_url: str = "", use_proxy: bool = False, proxy_config: Optional[Dict[str, Any]] = None, rate_limit: Optional[Dict[str, int]] = None, timeout: int = 60, extra_params: Optional[Dict[str, Any]] = None ) -> AIProvider: """ 创建新的AI接口配置 Args: provider_type: 提供商类型 name: 自定义名称 model: 模型名称 api_key: API密钥 base_url: API基础URL use_proxy: 是否使用代理 proxy_config: 代理配置 rate_limit: 速率限制配置 timeout: 超时时间 extra_params: 额外参数 Returns: 创建的AIProvider文档 """ # 验证提供商类型 try: get_adapter(provider_type) except ValueError as e: raise ValueError(f"不支持的提供商类型: {provider_type}") # 生成唯一ID provider_id = f"{provider_type}-{uuid.uuid4().hex[:8]}" # 加密API密钥 encrypted_key = encrypt_api_key(api_key) if api_key else "" # 创建文档 provider = AIProvider( provider_id=provider_id, provider_type=provider_type, name=name, api_key=encrypted_key, base_url=base_url, model=model, use_proxy=use_proxy, proxy_config=proxy_config or {}, rate_limit=rate_limit or {"requests_per_minute": 60, "tokens_per_minute": 100000}, timeout=timeout, extra_params=extra_params or {}, enabled=True, created_at=datetime.utcnow(), updated_at=datetime.utcnow() ) await provider.insert() # 注册速率限制 rate_limiter.register( provider_id, provider.rate_limit.get("requests_per_minute", 60), provider.rate_limit.get("tokens_per_minute", 100000) ) logger.info(f"创建AI接口配置: {provider_id} ({name})") return provider @classmethod async def get_provider(cls, provider_id: str) -> Optional[AIProvider]: """ 获取指定AI接口配置 Args: provider_id: 接口ID Returns: AIProvider文档或None """ return await AIProvider.find_one(AIProvider.provider_id == provider_id) @classmethod async def get_all_providers( cls, enabled_only: bool = False ) -> List[AIProvider]: """ 获取所有AI接口配置 Args: enabled_only: 是否只返回启用的接口 Returns: AIProvider列表 """ if enabled_only: return await AIProvider.find(AIProvider.enabled == True).to_list() return await AIProvider.find_all().to_list() @classmethod async def update_provider( cls, provider_id: str, **kwargs ) -> Optional[AIProvider]: """ 更新AI接口配置 Args: provider_id: 接口ID **kwargs: 要更新的字段 Returns: 更新后的AIProvider或None """ provider = await cls.get_provider(provider_id) if not provider: return None # 如果更新了API密钥，需要加密 if "api_key" in kwargs and kwargs["api_key"]: kwargs["api_key"] = encrypt_api_key(kwargs["api_key"]) # 更新字段 kwargs["updated_at"] = datetime.utcnow() for key, value in kwargs.items(): if hasattr(provider, key): setattr(provider, key, value) await provider.save() # 清除适配器缓存 cls._adapter_cache.pop(provider_id, None) # 更新速率限制 if "rate_limit" in kwargs: rate_limiter.unregister(provider_id) rate_limiter.register( provider_id, provider.rate_limit.get("requests_per_minute", 60), provider.rate_limit.get("tokens_per_minute", 100000) ) logger.info(f"更新AI接口配置: {provider_id}") return provider @classmethod async def delete_provider(cls, provider_id: str) -> bool: """ 删除AI接口配置 Args: provider_id: 接口ID Returns: 是否删除成功 """ provider = await cls.get_provider(provider_id) if not provider: return False await provider.delete() # 清除缓存和速率限制 cls._adapter_cache.pop(provider_id, None) rate_limiter.unregister(provider_id) logger.info(f"删除AI接口配置: {provider_id}") return True @classmethod async def get_adapter(cls, provider_id: str) -> Optional[BaseAdapter]: """ 获取AI接口的适配器实例 Args: provider_id: 接口ID Returns: 适配器实例或None """ # 检查缓存 if provider_id in cls._adapter_cache: return cls._adapter_cache[provider_id] provider = await cls.get_provider(provider_id) if not provider or not provider.enabled: return None # 解密API密钥 api_key = decrypt_api_key(provider.api_key) if provider.api_key else "" # 创建适配器 adapter_class = get_adapter(provider.provider_type) adapter = adapter_class( api_key=api_key, base_url=provider.base_url, model=provider.model, use_proxy=provider.use_proxy, proxy_config=provider.proxy_config, timeout=provider.timeout, **provider.extra_params ) # 缓存适配器 cls._adapter_cache[provider_id] = adapter return adapter @classmethod async def chat( cls, provider_id: str, messages: List[Dict[str, str]], temperature: float = 0.7, max_tokens: int = 2000, **kwargs ) -> AdapterResponse: """ 调用AI接口进行对话 Args: provider_id: 接口ID messages: 消息列表 [{"role": "user", "content": "..."}] temperature: 温度参数 max_tokens: 最大token数 **kwargs: 额外参数 Returns: 适配器响应 """ adapter = await cls.get_adapter(provider_id) if not adapter: return AdapterResponse( success=False, error=f"AI接口不存在或未启用: {provider_id}" ) # 检查速率限制 estimated_tokens = sum(len(m.get("content", "")) for m in messages) // 4 if not await rate_limiter.acquire_wait(provider_id, estimated_tokens): return AdapterResponse( success=False, error="请求频率超限，请稍后重试" ) # 转换消息格式 chat_messages = [ ChatMessage( role=m.get("role", "user"), content=m.get("content", ""), name=m.get("name") ) for m in messages ] # 调用适配器 response = await adapter.chat( messages=chat_messages, temperature=temperature, max_tokens=max_tokens, **kwargs ) return response @classmethod async def test_provider(cls, provider_id: str) -> Dict[str, Any]: """ 测试AI接口连接 Args: provider_id: 接口ID Returns: 测试结果 """ adapter = await cls.get_adapter(provider_id) if not adapter: return { "success": False, "message": f"AI接口不存在或未启用: {provider_id}" } return await adapter.test_connection() @classmethod async def test_provider_config( cls, provider_type: str, api_key: str, base_url: str = "", model: str = "", use_proxy: bool = False, proxy_config: Optional[Dict[str, Any]] = None, timeout: int = 30, **kwargs ) -> Dict[str, Any]: """ 测试AI接口配置（不保存） Args: provider_type: 提供商类型 api_key: API密钥 base_url: API基础URL model: 模型名称 use_proxy: 是否使用代理 proxy_config: 代理配置 timeout: 超时时间 **kwargs: 额外参数 Returns: 测试结果 """ try: adapter_class = get_adapter(provider_type) except ValueError: return { "success": False, "message": f"不支持的提供商类型: {provider_type}" } adapter = adapter_class( api_key=api_key, base_url=base_url, model=model, use_proxy=use_proxy, proxy_config=proxy_config, timeout=timeout, **kwargs ) return await adapter.test_connection()