430 lines
18 KiB
Python
430 lines
18 KiB
Python
"""统一大模型调用服务 - 支持多模型路由和流式输出"""
|
||
import json
|
||
import httpx
|
||
import asyncio
|
||
from typing import AsyncGenerator, List, Optional, Union
|
||
from openai import AsyncOpenAI
|
||
from anthropic import AsyncAnthropic
|
||
|
||
from config import (
|
||
MODEL_CONFIG,
|
||
OPENAI_API_KEY,
|
||
ANTHROPIC_API_KEY,
|
||
GOOGLE_API_KEY,
|
||
DEEPSEEK_API_KEY,
|
||
ARK_API_KEY,
|
||
ARK_BASE_URL,
|
||
)
|
||
|
||
|
||
def _get_db_model_config(task_type: str, model_config_id: int = None):
|
||
"""从数据库获取指定任务类型的默认模型配置，或指定 ID 的模型"""
|
||
try:
|
||
from database import SessionLocal
|
||
from models.ai_model import AIModelConfig
|
||
db = SessionLocal()
|
||
try:
|
||
# 如果指定了模型 ID，直接用该模型
|
||
if model_config_id:
|
||
model = db.query(AIModelConfig).filter(
|
||
AIModelConfig.id == model_config_id,
|
||
AIModelConfig.is_enabled == True,
|
||
).first()
|
||
if model and model.api_key:
|
||
return {
|
||
"provider": model.provider,
|
||
"model": model.model_id,
|
||
"api_key": model.api_key,
|
||
"base_url": model.base_url,
|
||
"web_search_enabled": model.web_search_enabled,
|
||
"web_search_count": model.web_search_count or 5,
|
||
}
|
||
# 否则找默认模型
|
||
model = db.query(AIModelConfig).filter(
|
||
AIModelConfig.task_type == task_type,
|
||
AIModelConfig.is_default == True,
|
||
AIModelConfig.is_enabled == True,
|
||
).first()
|
||
if model and model.api_key:
|
||
return {
|
||
"provider": model.provider,
|
||
"model": model.model_id,
|
||
"api_key": model.api_key,
|
||
"base_url": model.base_url,
|
||
"web_search_enabled": model.web_search_enabled,
|
||
"web_search_count": model.web_search_count or 5,
|
||
}
|
||
# 没有默认的，找任意一个启用且有Key的
|
||
model = db.query(AIModelConfig).filter(
|
||
AIModelConfig.task_type == task_type,
|
||
AIModelConfig.is_enabled == True,
|
||
AIModelConfig.api_key != "",
|
||
).first()
|
||
if model:
|
||
return {
|
||
"provider": model.provider,
|
||
"model": model.model_id,
|
||
"api_key": model.api_key,
|
||
"base_url": model.base_url,
|
||
"web_search_enabled": model.web_search_enabled,
|
||
"web_search_count": model.web_search_count or 5,
|
||
}
|
||
finally:
|
||
db.close()
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
class AIService:
|
||
"""统一AI服务，根据任务类型路由到不同大模型"""
|
||
|
||
def __init__(self):
|
||
# OpenAI客户端（也用于DeepSeek等兼容API）
|
||
if OPENAI_API_KEY:
|
||
self.openai_client = AsyncOpenAI(api_key=OPENAI_API_KEY)
|
||
else:
|
||
self.openai_client = None
|
||
|
||
# Anthropic客户端
|
||
if ANTHROPIC_API_KEY:
|
||
self.anthropic_client = AsyncAnthropic(api_key=ANTHROPIC_API_KEY)
|
||
else:
|
||
self.anthropic_client = None
|
||
|
||
# DeepSeek客户端（兼容OpenAI接口）
|
||
if DEEPSEEK_API_KEY:
|
||
self.deepseek_client = AsyncOpenAI(
|
||
api_key=DEEPSEEK_API_KEY,
|
||
base_url="https://api.deepseek.com/v1",
|
||
)
|
||
else:
|
||
self.deepseek_client = None
|
||
|
||
def _get_client_for_provider(self, provider: str, api_key: str, base_url: str = ""):
|
||
"""根据provider动态创建客户端"""
|
||
if provider == "anthropic":
|
||
return AsyncAnthropic(api_key=api_key)
|
||
# openai/deepseek/google 都用OpenAI兼容接口
|
||
kwargs = {"api_key": api_key}
|
||
if base_url:
|
||
kwargs["base_url"] = base_url
|
||
return AsyncOpenAI(**kwargs)
|
||
|
||
async def chat(
|
||
self,
|
||
task_type: str,
|
||
messages: List[dict],
|
||
system_prompt: str = "",
|
||
stream: bool = False,
|
||
model_config_id: int = None,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""
|
||
统一对话接口
|
||
|
||
参数:
|
||
task_type: 任务类型 (multimodal/reasoning/lightweight)
|
||
messages: 消息列表 [{"role": "user", "content": "..."}]
|
||
system_prompt: 系统提示词
|
||
stream: 是否流式输出
|
||
model_config_id: 指定模型配置ID（可选，不传则用默认）
|
||
"""
|
||
# 优先从数据库读取模型配置
|
||
db_config = _get_db_model_config(task_type, model_config_id)
|
||
if db_config:
|
||
provider = db_config["provider"]
|
||
model = db_config["model"]
|
||
api_key = db_config["api_key"]
|
||
base_url = db_config["base_url"]
|
||
web_search = db_config.get("web_search_enabled", False)
|
||
web_search_count = db_config.get("web_search_count", 5)
|
||
# 火山方舟/豆包 + 联网搜索（开启后自动使用，失败则降级到普通调用）
|
||
if provider == "ark" and web_search:
|
||
try:
|
||
return await self._chat_ark_web_search(api_key, base_url, model, messages, system_prompt, stream, web_search_count)
|
||
except Exception:
|
||
# 联网搜索调用失败，降级到普通调用
|
||
pass
|
||
# 火山方舟/豆包 不带联网搜索（OpenAI 兼容接口）
|
||
if provider == "ark":
|
||
kwargs = {"api_key": api_key}
|
||
if base_url:
|
||
kwargs["base_url"] = base_url
|
||
else:
|
||
kwargs["base_url"] = ARK_BASE_URL
|
||
client = AsyncOpenAI(**kwargs)
|
||
return await self._chat_openai(client, model, messages, system_prompt, stream)
|
||
if provider == "anthropic":
|
||
client = AsyncAnthropic(api_key=api_key)
|
||
return await self._chat_anthropic_with_client(client, model, messages, system_prompt, stream)
|
||
else:
|
||
kwargs = {"api_key": api_key}
|
||
if base_url:
|
||
kwargs["base_url"] = base_url
|
||
client = AsyncOpenAI(**kwargs)
|
||
# deepseek-reasoner 需要特殊处理
|
||
if model == "deepseek-reasoner":
|
||
return await self._chat_deepseek_reasoner(client, model, messages, system_prompt, stream)
|
||
return await self._chat_openai(client, model, messages, system_prompt, stream)
|
||
|
||
# 回退到 .env 配置
|
||
config = MODEL_CONFIG.get(task_type, MODEL_CONFIG["reasoning"])
|
||
provider = config["provider"]
|
||
model = config["model"]
|
||
|
||
if provider == "anthropic" and self.anthropic_client:
|
||
return await self._chat_anthropic(model, messages, system_prompt, stream)
|
||
elif provider == "openai" and self.openai_client:
|
||
return await self._chat_openai(self.openai_client, model, messages, system_prompt, stream)
|
||
elif provider == "deepseek" and self.deepseek_client:
|
||
return await self._chat_openai(self.deepseek_client, model, messages, system_prompt, stream)
|
||
elif provider == "google":
|
||
if GOOGLE_API_KEY:
|
||
google_client = AsyncOpenAI(
|
||
api_key=GOOGLE_API_KEY,
|
||
base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
|
||
)
|
||
return await self._chat_openai(google_client, model, messages, system_prompt, stream)
|
||
|
||
# 降级
|
||
if self.deepseek_client:
|
||
return await self._chat_openai(self.deepseek_client, "deepseek-chat", messages, system_prompt, stream)
|
||
if self.openai_client:
|
||
return await self._chat_openai(self.openai_client, "gpt-4o-mini", messages, system_prompt, stream)
|
||
if self.anthropic_client:
|
||
return await self._chat_anthropic("claude-sonnet-4-20250514", messages, system_prompt, stream)
|
||
|
||
return "未配置任何AI模型，请到「模型管理」页面配置模型和API Key。"
|
||
|
||
async def _chat_ark_web_search(
|
||
self, api_key: str, base_url: str, model: str,
|
||
messages: List[dict], system_prompt: str, stream: bool,
|
||
search_count: int = 5,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""火山方舟 + 联网搜索（使用 httpx 直接调用，因 web_search 是非标准 tools 类型）"""
|
||
url = f"{base_url or ARK_BASE_URL}/chat/completions"
|
||
full_messages = []
|
||
if system_prompt:
|
||
full_messages.append({"role": "system", "content": system_prompt})
|
||
full_messages.extend(messages)
|
||
|
||
# 限制搜索条数范围 1-50
|
||
search_count = max(1, min(50, search_count or 5))
|
||
payload = {
|
||
"model": model,
|
||
"messages": full_messages,
|
||
"stream": stream,
|
||
"tools": [{"type": "web_search", "web_search": {"enable": True, "search_result_count": search_count}}],
|
||
}
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
if not stream:
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
resp = await client.post(url, json=payload, headers=headers)
|
||
if resp.status_code != 200:
|
||
# 联网搜索调用失败，抛出异常以便降级到普通调用
|
||
raise Exception(f"API调用失败 ({resp.status_code}): {resp.text[:200]}")
|
||
data = resp.json()
|
||
return data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||
else:
|
||
# 流式调用时，先发一个预检测请求确认模型支持 web_search
|
||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||
test_payload = {
|
||
"model": model,
|
||
"messages": [{"role": "user", "content": "test"}],
|
||
"stream": False,
|
||
"max_tokens": 1,
|
||
"tools": [{"type": "web_search", "web_search": {"enable": True, "search_result_count": 1}}],
|
||
}
|
||
try:
|
||
resp = await client.post(url, json=test_payload, headers=headers)
|
||
if resp.status_code == 400:
|
||
# 模型不支持 web_search，抛出异常以便降级
|
||
raise Exception("模型不支持联网搜索")
|
||
except httpx.TimeoutException:
|
||
pass # 超时不影响，继续尝试流式调用
|
||
return self._stream_ark_web_search(url, payload, headers)
|
||
|
||
async def _stream_ark_web_search(
|
||
self, url: str, payload: dict, headers: dict,
|
||
) -> AsyncGenerator[str, None]:
|
||
"""火山方舟联网搜索流式输出"""
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
async with client.stream("POST", url, json=payload, headers=headers) as resp:
|
||
if resp.status_code != 200:
|
||
error_body = await resp.aread()
|
||
yield f"API调用失败 ({resp.status_code}): {error_body.decode()[:200]}"
|
||
return
|
||
buffer = ""
|
||
async for chunk in resp.aiter_text():
|
||
buffer += chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
if not line or not line.startswith("data: "):
|
||
continue
|
||
data_str = line[6:]
|
||
if data_str == "[DONE]":
|
||
continue
|
||
try:
|
||
data = json.loads(data_str)
|
||
choices = data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
content = delta.get("content", "")
|
||
if content:
|
||
yield content
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
async def _chat_openai(
|
||
self, client: AsyncOpenAI, model: str, messages: List[dict],
|
||
system_prompt: str, stream: bool,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""OpenAI兼容接口调用"""
|
||
full_messages = []
|
||
if system_prompt:
|
||
full_messages.append({"role": "system", "content": system_prompt})
|
||
full_messages.extend(messages)
|
||
|
||
if stream:
|
||
return self._stream_openai(client, model, full_messages)
|
||
else:
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=full_messages,
|
||
temperature=0.7,
|
||
max_tokens=4096,
|
||
)
|
||
return response.choices[0].message.content
|
||
|
||
async def _stream_openai(
|
||
self, client: AsyncOpenAI, model: str, messages: List[dict],
|
||
) -> AsyncGenerator[str, None]:
|
||
"""OpenAI流式输出"""
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=0.7,
|
||
max_tokens=4096,
|
||
stream=True,
|
||
)
|
||
async for chunk in response:
|
||
if chunk.choices[0].delta.content:
|
||
yield chunk.choices[0].delta.content
|
||
|
||
async def _chat_deepseek_reasoner(
|
||
self, client: AsyncOpenAI, model: str, messages: List[dict],
|
||
system_prompt: str, stream: bool,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""DeepSeek Reasoner (思考模式) 专用调用
|
||
|
||
注意：deepseek-reasoner 不支持 temperature/top_p/system 等参数
|
||
输出包含 reasoning_content（思考过程）和 content（最终回答）
|
||
"""
|
||
# reasoner 不支持 system role，将 system prompt 合并到第一条用户消息
|
||
full_messages = []
|
||
for msg in messages:
|
||
full_messages.append(msg)
|
||
if system_prompt and full_messages:
|
||
first_user = None
|
||
for m in full_messages:
|
||
if m["role"] == "user":
|
||
first_user = m
|
||
break
|
||
if first_user:
|
||
first_user["content"] = f"[指令] {system_prompt}\n\n[用户输入] {first_user['content']}"
|
||
|
||
if stream:
|
||
return self._stream_deepseek_reasoner(client, model, full_messages)
|
||
else:
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=full_messages,
|
||
max_tokens=8192,
|
||
)
|
||
reasoning = getattr(response.choices[0].message, 'reasoning_content', '') or ''
|
||
content = response.choices[0].message.content or ''
|
||
if reasoning:
|
||
return f"<think>\n{reasoning}\n</think>\n\n{content}"
|
||
return content
|
||
|
||
async def _stream_deepseek_reasoner(
|
||
self, client: AsyncOpenAI, model: str, messages: List[dict],
|
||
) -> AsyncGenerator[str, None]:
|
||
"""DeepSeek Reasoner 流式输出 - 包含思考过程和最终回答"""
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
max_tokens=8192,
|
||
stream=True,
|
||
)
|
||
in_reasoning = False
|
||
reasoning_started = False
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
# 思考过程
|
||
reasoning_content = getattr(delta, 'reasoning_content', None)
|
||
if reasoning_content:
|
||
if not reasoning_started:
|
||
reasoning_started = True
|
||
in_reasoning = True
|
||
yield "<details>\n<summary>💭 思考过程</summary>\n\n"
|
||
yield reasoning_content
|
||
# 最终回答
|
||
if delta.content:
|
||
if in_reasoning:
|
||
in_reasoning = False
|
||
yield "\n</details>\n\n"
|
||
yield delta.content
|
||
|
||
async def _chat_anthropic(
|
||
self, model: str, messages: List[dict],
|
||
system_prompt: str, stream: bool,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""Anthropic接口调用（使用self.anthropic_client）"""
|
||
return await self._chat_anthropic_with_client(self.anthropic_client, model, messages, system_prompt, stream)
|
||
|
||
async def _chat_anthropic_with_client(
|
||
self, client, model: str, messages: List[dict],
|
||
system_prompt: str, stream: bool,
|
||
) -> Union[str, AsyncGenerator[str, None]]:
|
||
"""Anthropic接口调用"""
|
||
if stream:
|
||
return self._stream_anthropic_with_client(client, model, messages, system_prompt)
|
||
else:
|
||
response = await client.messages.create(
|
||
model=model,
|
||
max_tokens=4096,
|
||
system=system_prompt if system_prompt else "You are a helpful assistant.",
|
||
messages=messages,
|
||
)
|
||
return response.content[0].text
|
||
|
||
async def _stream_anthropic(
|
||
self, model: str, messages: List[dict], system_prompt: str,
|
||
) -> AsyncGenerator[str, None]:
|
||
"""Anthropic流式输出（使用self.anthropic_client）"""
|
||
async for text in self._stream_anthropic_with_client(self.anthropic_client, model, messages, system_prompt):
|
||
yield text
|
||
|
||
async def _stream_anthropic_with_client(
|
||
self, client, model: str, messages: List[dict], system_prompt: str,
|
||
) -> AsyncGenerator[str, None]:
|
||
"""Anthropic流式输出"""
|
||
async with client.messages.stream(
|
||
model=model,
|
||
max_tokens=4096,
|
||
system=system_prompt if system_prompt else "You are a helpful assistant.",
|
||
messages=messages,
|
||
) as stream:
|
||
async for text in stream.text_stream:
|
||
yield text
|
||
|
||
|
||
# 单例
|
||
ai_service = AIService()
|