175 lines
6.0 KiB
Python
175 lines
6.0 KiB
Python
"""
|
||
豆包 (字节跳动/火山引擎) 平台适配器 - 使用火山方舟 API 获取真实AI回答
|
||
|
||
API文档: https://www.volcengine.com/docs/82379/1298454
|
||
使用火山方舟推理接入点 (Endpoint) API
|
||
认证方式: API Key (Bearer Token)
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from typing import Optional
|
||
|
||
import httpx
|
||
|
||
from app.config import settings
|
||
from app.workers.platforms.base import BasePlatformAdapter
|
||
from app.workers.platforms.search_engine import fetch_search_content
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 模块级频率限制器
|
||
_last_request_time: float = 0.0
|
||
_min_interval: float = 6.0 # 每分钟10次 -> 6秒间隔
|
||
|
||
|
||
async def _rate_limit_wait():
|
||
"""确保请求间隔不低于最小间隔"""
|
||
global _last_request_time
|
||
now = time.monotonic()
|
||
elapsed = now - _last_request_time
|
||
if elapsed < _min_interval:
|
||
wait_time = _min_interval - elapsed
|
||
logger.debug(f"豆包 频率限制等待 {wait_time:.1f}s")
|
||
await asyncio.sleep(wait_time)
|
||
_last_request_time = time.monotonic()
|
||
|
||
|
||
class DoubaoAdapter(BasePlatformAdapter):
|
||
"""豆包平台适配器 - 使用火山方舟 API 获取真实AI回答"""
|
||
|
||
platform_name = "doubao"
|
||
platform_url = "https://www.doubao.com/"
|
||
|
||
# 火山方舟 API 端点 (OpenAI兼容格式)
|
||
_api_base = "https://ark.cn-beijing.volces.com/api/v3"
|
||
_default_model = "doubao-pro-4k" # 默认模型，可被 endpoint_id 覆盖
|
||
|
||
def __init__(self):
|
||
self._api_key: Optional[str] = None
|
||
self._endpoint_id: Optional[str] = None
|
||
self._client: Optional[httpx.AsyncClient] = None
|
||
|
||
@property
|
||
def api_key(self) -> Optional[str]:
|
||
if self._api_key is None:
|
||
self._api_key = settings.DOUBAO_API_KEY
|
||
return self._api_key
|
||
|
||
@property
|
||
def endpoint_id(self) -> Optional[str]:
|
||
if self._endpoint_id is None:
|
||
self._endpoint_id = settings.DOUBAO_ENDPOINT_ID
|
||
return self._endpoint_id
|
||
|
||
@property
|
||
def is_configured(self) -> bool:
|
||
return bool(self.api_key and self.api_key.strip())
|
||
|
||
def _get_model_id(self) -> str:
|
||
"""获取实际使用的模型ID（优先使用 endpoint_id）"""
|
||
if self.endpoint_id and self.endpoint_id.strip():
|
||
return f"ep-{self.endpoint_id}" if not self.endpoint_id.startswith("ep-") else self.endpoint_id
|
||
return self._default_model
|
||
|
||
def _get_client(self) -> httpx.AsyncClient:
|
||
if self._client is None or self._client.is_closed:
|
||
self._client = httpx.AsyncClient(
|
||
timeout=httpx.Timeout(60.0, connect=10.0),
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
},
|
||
)
|
||
return self._client
|
||
|
||
async def query(self, keyword: str) -> str:
|
||
"""在豆包查询关键词，返回原始响应文本"""
|
||
last_error = None
|
||
for attempt in range(3):
|
||
try:
|
||
return await self._do_query(keyword)
|
||
except Exception as e:
|
||
last_error = e
|
||
logger.warning(f"豆包查询第 {attempt + 1} 次尝试失败: {e}")
|
||
if attempt < 2:
|
||
await asyncio.sleep(2 ** attempt)
|
||
|
||
# 所有尝试失败后回退到搜索引擎
|
||
logger.warning(f"豆包 API 调用全部失败，回退到搜索引擎: {last_error}")
|
||
return await self._fallback_search(keyword)
|
||
|
||
async def _do_query(self, keyword: str) -> str:
|
||
"""通过火山方舟 API 获取真实AI回答"""
|
||
if not self.is_configured:
|
||
logger.warning("豆包 API Key 未配置，回退到搜索引擎")
|
||
return await self._fallback_search(keyword)
|
||
|
||
await _rate_limit_wait()
|
||
|
||
model_id = self._get_model_id()
|
||
client = self._get_client()
|
||
|
||
payload = {
|
||
"model": model_id,
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"你是一个专业的AI搜索助手。请基于你的知识，"
|
||
"详细回答用户的问题。如果引用了外部来源，"
|
||
"请在回答中标注来源URL或出处名称。"
|
||
),
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": keyword,
|
||
},
|
||
],
|
||
"temperature": 0.7,
|
||
"max_tokens": 2000,
|
||
}
|
||
|
||
response = await client.post(
|
||
f"{self._api_base}/chat/completions",
|
||
json=payload,
|
||
)
|
||
|
||
if response.status_code == 429:
|
||
retry_after = int(response.headers.get("Retry-After", "10"))
|
||
logger.warning(f"豆包 API 限流，等待 {retry_after}s 后重试")
|
||
await asyncio.sleep(retry_after)
|
||
raise RuntimeError(f"豆包 API 限流")
|
||
|
||
if response.status_code != 200:
|
||
error_body = response.text[:500]
|
||
raise RuntimeError(
|
||
f"豆包 API 返回错误 {response.status_code}: {error_body}"
|
||
)
|
||
|
||
data = response.json()
|
||
choices = data.get("choices", [])
|
||
if not choices:
|
||
raise RuntimeError("豆包 API 返回空 choices")
|
||
|
||
content = choices[0].get("message", {}).get("content", "")
|
||
if not content:
|
||
raise RuntimeError("豆包 API 返回空内容")
|
||
|
||
# 标记数据来源为AI平台
|
||
result = f"[data_source: ai_platform]\n{content}"
|
||
logger.info(f"豆包 API 调用成功，返回 {len(content)} 字符")
|
||
return result
|
||
|
||
async def _fallback_search(self, keyword: str) -> str:
|
||
"""回退到搜索引擎模式，标记数据来源"""
|
||
content = await fetch_search_content(self.platform_name, keyword)
|
||
return f"[data_source: search_engine]\n{content}"
|
||
|
||
async def close(self):
|
||
"""清理资源"""
|
||
if self._client and not self._client.is_closed:
|
||
await self._client.aclose()
|
||
self._client = None
|