feat(ai-client): 实现多 AI 提供商支持与 fallback 机制
- 支持 MiniMax、智谱、通义千问等多家 AI 提供商 - 实现自动 fallback 机制(主失败则自动切换备用) - 修复 MiniMax Anthropic 格式响应解析 - 修复智谱 API 端点(自动添加 /chat/completions) - 新增 _call_api_with_fallback 方法统一处理 fallback 逻辑 - 更新 _load_ai_config 加载主配置和 fallbacks 配置 配置文件位置: ~/.config/obsidian-kb/config.yaml
This commit is contained in:
@@ -1,9 +1,10 @@
|
||||
"""Obsidian KB - AI客户端
|
||||
C5.3: 实现长视频分段分析逻辑
|
||||
支持多个 AI 提供商和 fallback 机制
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, Any, List, Optional
|
||||
import requests
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
from src.models import ContentType
|
||||
from src.config import get_config
|
||||
from src.content_chunker import ContentChunker, get_chunking_strategy, chunk_content
|
||||
@@ -11,21 +12,217 @@ from src.content_chunker import ContentChunker, get_chunking_strategy, chunk_con
|
||||
|
||||
class AIClient:
|
||||
"""
|
||||
AI客户端 - 支持分段分析长视频内容
|
||||
AI客户端 - 支持分段分析长视频内容和 fallback 机制
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.config = get_config()
|
||||
self.chunker = ContentChunker()
|
||||
# 获取 AI 配置
|
||||
self._load_ai_config()
|
||||
|
||||
def _load_ai_config(self) -> None:
|
||||
"""加载 AI 配置(主 provider + fallbacks)"""
|
||||
self.ai_configs = []
|
||||
|
||||
# 加载主配置
|
||||
main_provider = self.config.get('ai.provider', 'openai')
|
||||
main_config = {
|
||||
'provider': main_provider,
|
||||
'api_key': self.config.get('ai.api_key'),
|
||||
'model': self.config.get('ai.model', 'gpt-4'),
|
||||
'base_url': self.config.get('ai.base_url'),
|
||||
'is_primary': True
|
||||
}
|
||||
self.ai_configs.append(main_config)
|
||||
|
||||
# 加载 fallback 配置
|
||||
fallbacks = self.config.get('ai.fallbacks', [])
|
||||
for fb in fallbacks:
|
||||
self.ai_configs.append({
|
||||
'provider': fb.get('provider'),
|
||||
'api_key': fb.get('api_key'),
|
||||
'model': fb.get('model'),
|
||||
'base_url': fb.get('base_url'),
|
||||
'is_primary': False
|
||||
})
|
||||
|
||||
def _call_api_with_fallback(self, prompt: str, max_tokens: int = 1000, response_format: str = 'text') -> Tuple[bool, str]:
|
||||
"""
|
||||
使用 fallback 机制调用 API
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
max_tokens: 最大 token 数
|
||||
response_format: 响应格式 ('text' 或 'json')
|
||||
|
||||
Returns:
|
||||
(success, response) 元组
|
||||
"""
|
||||
last_error = None
|
||||
|
||||
for i, ai_config in enumerate(self.ai_configs):
|
||||
try:
|
||||
if not ai_config.get('api_key'):
|
||||
print(f"[AIClient] 跳过 {ai_config.get('provider')} - API Key 为空")
|
||||
continue
|
||||
|
||||
provider_type = ai_config.get('provider')
|
||||
print(f"[AIClient] 尝试使用 {provider_type} (优先级 {i+1}/{len(self.ai_configs)})...")
|
||||
|
||||
response = self._call_single_provider(ai_config, prompt, max_tokens, response_format)
|
||||
return True, response
|
||||
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
print(f"[AIClient] {provider_type} 调用失败: {e}")
|
||||
continue
|
||||
|
||||
# 所有 provider 都失败了
|
||||
error_msg = f"所有 AI 提供商均失败,最后错误: {last_error}" if last_error else "没有可用的 AI 配置"
|
||||
print(f"[AIClient] {error_msg}")
|
||||
return False, self._mock_generate(prompt, max_tokens) # 返回 mock 响应作为兜底
|
||||
|
||||
def _call_single_provider(self, ai_config: Dict, prompt: str, max_tokens: int, response_format: str) -> str:
|
||||
"""
|
||||
调用单个 AI 提供商
|
||||
|
||||
Args:
|
||||
ai_config: AI 配置
|
||||
prompt: 提示词
|
||||
max_tokens: 最大 token 数
|
||||
response_format: 响应格式
|
||||
|
||||
Returns:
|
||||
API 响应文本
|
||||
"""
|
||||
provider = ai_config.get('provider')
|
||||
api_key = ai_config.get('api_key')
|
||||
model = ai_config.get('model')
|
||||
base_url = ai_config.get('base_url')
|
||||
|
||||
if provider == 'minimax':
|
||||
return self._call_minimax(api_key, model, base_url, prompt, max_tokens, response_format)
|
||||
elif provider == 'zhipu':
|
||||
return self._call_zhipu(api_key, model, base_url, prompt, max_tokens, response_format)
|
||||
elif provider == 'alibaba':
|
||||
return self._call_alibaba(api_key, model, base_url, prompt, max_tokens, response_format)
|
||||
elif provider == 'openai':
|
||||
return self._call_openai(api_key, model, base_url, prompt, max_tokens, response_format)
|
||||
else:
|
||||
raise ValueError(f"不支持的 provider: {provider}")
|
||||
|
||||
def _call_minimax(self, api_key: str, model: str, base_url: str, prompt: str, max_tokens: int, response_format: str) -> str:
|
||||
"""调用 MiniMax API (Anthropic Messages 兼容接口)"""
|
||||
url = base_url or "https://api.minimaxi.com/anthropic/v1/messages"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
}
|
||||
|
||||
if response_format == 'json':
|
||||
payload['messages'][0]['content'] = prompt + "\n\n请直接返回 JSON 格式,不要有其他文字。"
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# MiniMax 返回的是 Anthropic 格式,content 是一个数组
|
||||
# 需要找到 type='text' 的元素
|
||||
for item in data.get('content', []):
|
||||
if item.get('type') == 'text':
|
||||
return item.get('text', '')
|
||||
|
||||
# 如果没有找到 text 类型,返回第一个元素的值(兼容旧格式)
|
||||
return data.get('content', [{}])[0].get('text', '')
|
||||
|
||||
def _call_zhipu(self, api_key: str, model: str, base_url: str, prompt: str, max_tokens: int, response_format: str) -> str:
|
||||
"""调用智谱 AI API (OpenAI 兼容接口)"""
|
||||
# 确保 base_url 包含 /chat/completions
|
||||
if base_url and not base_url.endswith('/chat/completions'):
|
||||
url = f"{base_url}/chat/completions"
|
||||
else:
|
||||
url = base_url or "https://open.bigmodel.cn/api/paas/v4/chat/completions"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data['choices'][0]['message']['content']
|
||||
|
||||
def _call_alibaba(self, api_key: str, model: str, base_url: str, prompt: str, max_tokens: int, response_format: str) -> str:
|
||||
"""调用阿里云通义千问 API (OpenAI 兼容接口)"""
|
||||
url = f"{base_url}/chat/completions"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data['choices'][0]['message']['content']
|
||||
|
||||
def _call_openai(self, api_key: str, model: str, base_url: str, prompt: str, max_tokens: int, response_format: str) -> str:
|
||||
"""调用 OpenAI API"""
|
||||
url = base_url or "https://api.openai.com/v1/chat/completions"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": [
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data['choices'][0]['message']['content']
|
||||
|
||||
def quick_generate(self, prompt: str, max_tokens: int = 100) -> str:
|
||||
"""快速生成(用于轻量扫描)"""
|
||||
# 模拟AI响应(实际应调用API)
|
||||
return self._mock_generate(prompt, max_tokens)
|
||||
success, response = self._call_api_with_fallback(prompt, max_tokens, 'text')
|
||||
return response
|
||||
|
||||
def generate_json(self, prompt: str) -> Dict[str, Any]:
|
||||
"""生成JSON格式响应"""
|
||||
response = self.quick_generate(prompt, max_tokens=1000)
|
||||
success, response = self._call_api_with_fallback(prompt, 1000, 'json')
|
||||
try:
|
||||
return json.loads(response)
|
||||
except:
|
||||
|
||||
Reference in New Issue
Block a user