214 lines
7.3 KiB
Python
214 lines
7.3 KiB
Python
"""
|
||
文心一言 (百度千帆) 平台适配器 - 使用百度千帆 API 获取真实AI回答
|
||
|
||
API文档: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/flfmc9do2
|
||
使用 ERNIE-Bot (completions) chat completion API
|
||
认证方式: API Key + Secret Key 换取 access_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秒间隔
|
||
|
||
# access_token 缓存
|
||
_cached_token: Optional[str] = None
|
||
_token_expires_at: float = 0.0
|
||
|
||
|
||
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 WenxinAdapter(BasePlatformAdapter):
|
||
"""文心一言平台适配器 - 使用百度千帆 API 获取真实AI回答"""
|
||
|
||
platform_name = "wenxin"
|
||
platform_url = "https://yiyan.baidu.com"
|
||
|
||
# 百度千帆 API 端点
|
||
_token_url = "https://aip.baidubce.com/oauth/2.0/token"
|
||
# ERNIE-Bot 4.0 (completions_pro) 或 ERNIE-Bot (completions)
|
||
_chat_url_template = (
|
||
"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/{model}"
|
||
"?access_token={token}"
|
||
)
|
||
_default_model = "completions_pro" # ERNIE-Bot 4.0
|
||
|
||
def __init__(self):
|
||
self._api_key: Optional[str] = None
|
||
self._secret_key: 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.BAIDU_QIANFAN_API_KEY
|
||
return self._api_key
|
||
|
||
@property
|
||
def secret_key(self) -> Optional[str]:
|
||
if self._secret_key is None:
|
||
self._secret_key = settings.BAIDU_QIANFAN_SECRET_KEY
|
||
return self._secret_key
|
||
|
||
@property
|
||
def is_configured(self) -> bool:
|
||
return bool(
|
||
self.api_key and self.api_key.strip()
|
||
and self.secret_key and self.secret_key.strip()
|
||
)
|
||
|
||
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),
|
||
)
|
||
return self._client
|
||
|
||
async def _get_access_token(self) -> str:
|
||
"""获取百度千帆 access_token，带缓存"""
|
||
global _cached_token, _token_expires_at
|
||
|
||
now = time.monotonic()
|
||
if _cached_token and now < _token_expires_at:
|
||
return _cached_token
|
||
|
||
client = self._get_client()
|
||
response = await client.post(
|
||
self._token_url,
|
||
params={
|
||
"grant_type": "client_credentials",
|
||
"client_id": self.api_key,
|
||
"client_secret": self.secret_key,
|
||
},
|
||
)
|
||
|
||
if response.status_code != 200:
|
||
raise RuntimeError(
|
||
f"百度千帆获取 access_token 失败: {response.status_code} {response.text[:300]}"
|
||
)
|
||
|
||
data = response.json()
|
||
token = data.get("access_token")
|
||
if not token:
|
||
error_desc = data.get("error_description", "未知错误")
|
||
raise RuntimeError(f"百度千帆获取 access_token 失败: {error_desc}")
|
||
|
||
# 缓存 token，提前5分钟过期
|
||
expires_in = data.get("expires_in", 2592000) # 默认30天
|
||
_cached_token = token
|
||
_token_expires_at = now + expires_in - 300
|
||
|
||
logger.info("百度千帆 access_token 获取成功")
|
||
return token
|
||
|
||
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()
|
||
|
||
access_token = await self._get_access_token()
|
||
chat_url = self._chat_url_template.format(
|
||
model=self._default_model,
|
||
token=access_token,
|
||
)
|
||
|
||
client = self._get_client()
|
||
payload = {
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": keyword,
|
||
},
|
||
],
|
||
"system": (
|
||
"你是一个专业的AI搜索助手。请基于你的知识，"
|
||
"详细回答用户的问题。如果引用了外部来源，"
|
||
"请在回答中标注来源URL或出处名称。"
|
||
),
|
||
"temperature": 0.7,
|
||
"max_output_tokens": 2000,
|
||
}
|
||
|
||
response = await client.post(chat_url, 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()
|
||
|
||
# 检查API错误码
|
||
error_code = data.get("error_code")
|
||
if error_code:
|
||
error_msg = data.get("error_msg", "未知错误")
|
||
raise RuntimeError(f"百度千帆 API 错误 {error_code}: {error_msg}")
|
||
|
||
content = data.get("result", "")
|
||
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
|