""" 火山方舟官方 SDK Provider ========================== 基于火山方舟官方 Python SDK 实现，支持： - 文本生成 (Chat Completions) - 流式输出 - 向量化 - 深度思考 - 工具调用 安装依赖: pip install 'volcengine-python-sdk[ark]' 文档: https://www.volcengine.com/docs/82379 """ from __future__ import annotations import logging import time from app.ai.providers.base import ( GenerationResult, LLMProvider, ModelHealth, ProviderError, ) from app.core.exceptions import PlatformError, PlatformErrorType logger = logging.getLogger(__name__) # 尝试导入火山方舟 SDK try: from volcenginesdkarkruntime import AsyncArk VOLCENGINE_SDK_AVAILABLE = True except ImportError: VOLCENGINE_SDK_AVAILABLE = False logger.warning("火山方舟 SDK 未安装，请运行: pip install 'volcengine-python-sdk[ark]'") class VolcengineProvider(LLMProvider): """ 火山方舟官方 SDK Provider 支持能力： - 文本对话 (Chat Completions) - 向量化 (Embeddings) - 深度思考 (Reasoning) """ provider_id = "volcengine" provider_name = "火山方舟" # 默认配置 DEFAULT_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3" DEFAULT_TIMEOUT = 1800 # 秒，方舟推荐 1800 秒以上 # 模型 ID 映射（从配置文件加载） PRESET_MODELS: dict[str, str] = {} @classmethod def load_models_from_config(cls, models: list[dict]): """ 从配置文件加载模型列表 Args: models: 模型列表，每个模型包含 model_name 字段 """ cls.PRESET_MODELS = {} for model in models: model_id = model.get("model_name") model_alias = model.get("id") if model_id and model_alias: cls.PRESET_MODELS[model_alias] = model_id # 确保至少有一个默认模型 if not cls.PRESET_MODELS: cls.PRESET_MODELS = { "doubao-seed-2-0-pro": "doubao-seed-2-0-pro-260215", } logger.info(f"已加载 {len(cls.PRESET_MODELS)} 个模型: {list(cls.PRESET_MODELS.keys())}") def __init__( self, api_key: str | None = None, base_url: str | None = None, timeout: int = DEFAULT_TIMEOUT, default_model: str | None = None, **kwargs, ): """ 初始化火山方舟 Provider Args: api_key: API Key，从火山方舟控制台获取 base_url: API 基础地址，默认北京节点 timeout: 请求超时时间（秒） default_model: 默认模型（Model ID） """ from app.config import get_settings from app.core.platform_config import get_platform_config_loader settings = get_settings() # API Key 从环境变量读取 resolved_api_key = api_key or settings.VOLCENGINE_API_KEY # base_url 从 platform-config.yaml 读取，fallback 到代码常量 loader = get_platform_config_loader() platform = loader.get_platform("volcengine_ark") yaml_base_url = platform.base_url if platform else None resolved_base_url = base_url or yaml_base_url or self.DEFAULT_BASE_URL super().__init__(resolved_api_key, resolved_base_url, **kwargs) if not VOLCENGINE_SDK_AVAILABLE: raise ProviderError( "火山方舟 SDK 未安装，请运行: pip install 'volcengine-python-sdk[ark]'", provider_id=self.provider_id, ) if not self.api_key: raise ProviderError("火山方舟 API Key 未配置", provider_id=self.provider_id) self.timeout = timeout # 使用模型 ID 映射（自动映射） if default_model: self.default_model = self.PRESET_MODELS.get(default_model, default_model) elif self.PRESET_MODELS: self.default_model = list(self.PRESET_MODELS.values())[0] else: # 兜底：使用一个默认模型ID（如果用户未配置任何模型） self.default_model = "doubao-seed-2-0-pro-260215" self.client = self._create_client() def _create_client(self) -> AsyncArk: """创建火山方舟异步客户端""" return AsyncArk( api_key=self.api_key, base_url=self.base_url or self.DEFAULT_BASE_URL, timeout=self.timeout, ) async def generate( self, prompt: str, model: str | None = None, temperature: float = 0.7, max_tokens: int | None = None, system_prompt: str | None = None, **kwargs, ) -> GenerationResult: """ 同步生成文本 Args: prompt: 用户提示词 model: 模型 ID（如 doubao-seed-2-0-pro-260215） temperature: 随机性 (0-2) max_tokens: 最大生成 token 数 system_prompt: 系统提示词（可选） **kwargs: 额外参数（如 reasoning_effort 控制思考程度） Returns: GenerationResult: 生成结果 """ try: # 构建消息 messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": prompt}) # 映射模型名称到模型 ID model_id = self.PRESET_MODELS.get(model, model) if model else self.default_model # 构建请求参数 request_params = { "model": model_id, "messages": messages, "temperature": temperature, } if max_tokens: request_params["max_tokens"] = max_tokens # 额外参数（如深度思考） if "reasoning_effort" in kwargs: request_params["reasoning_effort"] = kwargs["reasoning_effort"] if kwargs.get("response_format") == "json_object": request_params["response_format"] = {"type": "json_object"} # 调用 API completion = await self.client.chat.completions.create(**request_params) # 解析结果 content = completion.choices[0].message.content or "" usage = None if completion.usage: usage = { "prompt_tokens": completion.usage.prompt_tokens, "completion_tokens": completion.usage.completion_tokens, "total_tokens": completion.usage.total_tokens, } return GenerationResult( content=content, usage=usage, model=completion.model or model or self.default_model, ) except Exception as e: raise self._wrap_error(e) async def create_embeddings(self, texts: list[str], model: str | None = None, **kwargs) -> dict: """ 文本向量化 Args: texts: 文本列表 model: 向量化模型 Returns: dict: 包含向量化结果 """ try: response = await self.client.embeddings.create( model=model or "doubao-embedding-1.5", input=texts, **kwargs ) embeddings = [] for item in response.data: embeddings.append( { "index": item.index, "embedding": item.embedding, } ) return { "embeddings": embeddings, "model": response.model, "usage": { "prompt_tokens": response.usage.prompt_tokens, "total_tokens": response.usage.total_tokens, }, } except Exception as e: raise self._wrap_error(e) async def health_check(self, model: str | None = None) -> ModelHealth: """健康检查""" start_time = time.time() test_model = model or self.default_model try: await self.client.chat.completions.create( model=test_model, messages=[{"role": "user", "content": "Hi"}], max_tokens=5, ) response_time = (time.time() - start_time) * 1000 return ModelHealth( id=test_model, name=f"火山方舟 {test_model}", is_available=True, response_time=response_time, last_error=None, ) except Exception as e: return ModelHealth( id=test_model, name=f"火山方舟 {test_model}", is_available=False, response_time=(time.time() - start_time) * 1000, last_error=str(e), ) def _wrap_error(self, e: Exception) -> PlatformError: """把 SDK 异常翻译为标准 PlatformError""" message = str(e) status = getattr(e, "status_code", None) or getattr(e, "code", None) if status == 429 or "rate limit" in message.lower(): return PlatformError( message, platform="volcengine_ark", retryable=True, error_type=PlatformErrorType.RATE_LIMIT, status_code=status, ) elif status in (401, 403) or "authentication" in message.lower(): return PlatformError( message, platform="volcengine_ark", retryable=False, error_type=PlatformErrorType.AUTH_FAILED, status_code=status, ) elif status and status >= 500: return PlatformError( message, platform="volcengine_ark", retryable=True, error_type=PlatformErrorType.SERVER_ERROR, status_code=status, ) elif "timeout" in message.lower() or isinstance(e, TimeoutError): return PlatformError( message, platform="volcengine_ark", retryable=True, error_type=PlatformErrorType.TIMEOUT, ) else: return PlatformError( message, platform="volcengine_ark", retryable=False, error_type=PlatformErrorType.UNKNOWN, ) @property def available_models(self) -> list[str]: """返回可用模型列表（与 platform-config.yaml 配置保持一致）""" return [ "doubao-seed-2-0-pro", ]