288 lines
11 KiB
Python
288 lines
11 KiB
Python
import json
|
||
import logging
|
||
import re
|
||
import time
|
||
|
||
from openai import AsyncOpenAI
|
||
import anthropic
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class LLMClient:
|
||
"""LLM 客户端，支持 OpenAI 和 Anthropic SDK"""
|
||
|
||
def __init__(self):
|
||
self._openai_client = None
|
||
self._anthropic_client = None
|
||
self._api_key = settings.LLM_API_KEY
|
||
self._base_url = settings.LLM_BASE_URL
|
||
self._model = settings.LLM_MODEL
|
||
self._provider = settings.LLM_PROVIDER
|
||
|
||
def _get_openai_client(self) -> AsyncOpenAI:
|
||
if self._openai_client is None:
|
||
self._openai_client = AsyncOpenAI(
|
||
api_key=self._api_key,
|
||
base_url=self._base_url,
|
||
max_retries=0,
|
||
timeout=300.0,
|
||
)
|
||
return self._openai_client
|
||
|
||
def _get_anthropic_client(self) -> anthropic.AsyncAnthropic:
|
||
if self._anthropic_client is None:
|
||
# MiniMax Anthropic 端点
|
||
base_url = self._base_url.replace("/v1", "/anthropic").replace("api.minimax.chat", "api.minimaxi.com")
|
||
if not base_url.endswith("/anthropic"):
|
||
# 如果 base_url 不是标准格式，手动拼接
|
||
base_url = "https://api.minimaxi.com/anthropic"
|
||
self._anthropic_client = anthropic.AsyncAnthropic(
|
||
api_key=self._api_key,
|
||
base_url=base_url,
|
||
timeout=300.0,
|
||
)
|
||
return self._anthropic_client
|
||
|
||
def _use_anthropic(self) -> bool:
|
||
"""判断是否使用 Anthropic SDK"""
|
||
return False # MiniMax thinking 占用太多 token，统一用 OpenAI SDK
|
||
|
||
def update_config(self, base_url: str = None, api_key: str = None, model: str = None):
|
||
"""更新配置（设置页面保存后调用）"""
|
||
if base_url:
|
||
self._base_url = base_url
|
||
if api_key:
|
||
self._api_key = api_key
|
||
if model:
|
||
self._model = model
|
||
# 重置客户端
|
||
self._openai_client = None
|
||
self._anthropic_client = None
|
||
# 更新 provider 判断
|
||
if base_url:
|
||
self._provider = "minimax" if "minimax" in base_url.lower() else "openai"
|
||
|
||
async def chat(self, messages: list[dict], temperature: float = None, max_tokens: int = 16384, timeout: float = 300.0) -> str:
|
||
"""发送对话请求，返回文本内容（不含思考过程）"""
|
||
if temperature is None:
|
||
temperature = settings.TEMPERATURE
|
||
|
||
if self._use_anthropic():
|
||
return await self._chat_anthropic(messages, temperature, max_tokens, timeout)
|
||
else:
|
||
result = await self._chat_openai(messages, temperature, max_tokens, timeout)
|
||
return result["content"]
|
||
|
||
async def chat_detail(self, messages: list[dict], temperature: float = None, max_tokens: int = 16384, timeout: float = 300.0) -> dict:
|
||
"""发送对话请求，返回详细信息（content + reasoning + finish_reason + usage + elapsed）"""
|
||
if temperature is None:
|
||
temperature = settings.TEMPERATURE
|
||
|
||
if self._use_anthropic():
|
||
# Anthropic 暂不支持 detail
|
||
content = await self._chat_anthropic(messages, temperature, max_tokens, timeout)
|
||
return {"content": content, "reasoning_content": "", "finish_reason": "stop", "usage": {}, "elapsed": 0}
|
||
else:
|
||
return await self._chat_openai(messages, temperature, max_tokens, timeout)
|
||
|
||
async def _chat_openai(self, messages: list[dict], temperature: float, max_tokens: int, timeout: float = 300.0) -> dict:
|
||
"""OpenAI SDK 调用，返回详细信息 dict"""
|
||
client = self._get_openai_client()
|
||
t0 = time.time()
|
||
response = await client.chat.completions.create(
|
||
model=self._model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
timeout=timeout,
|
||
)
|
||
elapsed = time.time() - t0
|
||
|
||
content = response.choices[0].message.content or ""
|
||
# 提取 reasoning_content（thinking）
|
||
reasoning = ""
|
||
msg = response.choices[0].message
|
||
if hasattr(msg, 'reasoning_content') and msg.reasoning_content:
|
||
reasoning = msg.reasoning_content
|
||
# 有些 SDK 把 reasoning 放在 model_extra 里
|
||
if not reasoning and hasattr(msg, 'model_extra') and isinstance(msg.model_extra, dict):
|
||
reasoning = msg.model_extra.get("reasoning_content", "")
|
||
|
||
# 检查是否有多个 choice
|
||
if len(response.choices) > 1:
|
||
print(f"[LLM] 警告: 返回了 {len(response.choices)} 个 choices")
|
||
# 检查 finish_reason
|
||
reason = response.choices[0].finish_reason if response.choices else "unknown"
|
||
if reason != "stop":
|
||
print(f"[LLM] 警告: finish_reason={reason}, 可能被截断")
|
||
|
||
# usage 信息
|
||
usage = {}
|
||
if hasattr(response, 'usage') and response.usage:
|
||
usage = {
|
||
"prompt_tokens": response.usage.prompt_tokens,
|
||
"completion_tokens": response.usage.completion_tokens,
|
||
"total_tokens": response.usage.total_tokens,
|
||
}
|
||
if hasattr(response.usage, 'completion_tokens_details') and response.usage.completion_tokens_details:
|
||
details = response.usage.completion_tokens_details
|
||
usage["reasoning_tokens"] = getattr(details, 'reasoning_tokens', 0) or 0
|
||
|
||
# 保存完整响应到文件，供调试
|
||
try:
|
||
resp_dict = response.model_dump() if hasattr(response, 'model_dump') else str(response)
|
||
debug_path = f"/tmp/llm_response_{id(response) % 100000}.json"
|
||
with open(debug_path, "w", encoding="utf-8") as _f:
|
||
json.dump(resp_dict, _f, ensure_ascii=False, indent=2, default=str)
|
||
print(f"[LLM] 完整响应已保存: {debug_path}")
|
||
except Exception:
|
||
pass
|
||
|
||
return {
|
||
"content": content,
|
||
"reasoning_content": reasoning,
|
||
"finish_reason": reason,
|
||
"usage": usage,
|
||
"elapsed": round(elapsed, 2),
|
||
}
|
||
|
||
async def _chat_anthropic(self, messages: list[dict], temperature: float, max_tokens: int, timeout: float = 300.0) -> str:
|
||
"""Anthropic SDK 调用 - thinking 和 text 自动分离"""
|
||
client = self._get_anthropic_client()
|
||
|
||
# 转换消息格式：OpenAI -> Anthropic
|
||
system_msg = ""
|
||
anthropic_messages = []
|
||
for msg in messages:
|
||
role = msg.get("role", "user")
|
||
content = msg.get("content", "")
|
||
if role == "system":
|
||
system_msg = content
|
||
elif role in ("user", "assistant"):
|
||
anthropic_messages.append({"role": role, "content": content})
|
||
|
||
kwargs = {
|
||
"model": self._model,
|
||
"max_tokens": max_tokens,
|
||
"messages": anthropic_messages,
|
||
"thinking": {"type": "disabled"}, # 禁用 thinking，避免占用 token
|
||
}
|
||
if system_msg:
|
||
kwargs["system"] = system_msg
|
||
if temperature is not None:
|
||
kwargs["temperature"] = temperature
|
||
|
||
response = await client.messages.create(**kwargs)
|
||
|
||
# 从 response 中提取文本内容，thinking 自动分离在 block.thinking 中
|
||
text_parts = []
|
||
for block in response.content:
|
||
if hasattr(block, 'text') and block.type == 'text':
|
||
text_parts.append(block.text)
|
||
# thinking 内容自动忽略
|
||
|
||
result = "\n".join(text_parts)
|
||
return result
|
||
|
||
async def chat_json(self, messages: list[dict], temperature: float = None, max_tokens: int = 16384) -> dict:
|
||
"""发送对话请求，期望返回 JSON"""
|
||
raw = await self.chat(messages, temperature, max_tokens=max_tokens)
|
||
return _parse_json(raw)
|
||
|
||
async def test_connection(self) -> dict:
|
||
"""测试 LLM 连接"""
|
||
try:
|
||
response = await self.chat([{"role": "user", "content": "请回复：连接成功"}])
|
||
return {"success": True, "message": f"连接成功，模型回复：{response[:50]}"}
|
||
except Exception as e:
|
||
return {"success": False, "message": f"连接失败：{str(e)}"}
|
||
|
||
|
||
def _parse_json(raw: str) -> dict:
|
||
"""健壮的 JSON 解析，处理 LLM 返回的各种格式问题"""
|
||
print(f"[_parse_json] raw长度={len(raw)}, 前100字={raw[:100]}")
|
||
print(f"[_parse_json] raw后100字={raw[-100:]}")
|
||
raw = raw.strip()
|
||
|
||
# 去掉 markdown 代码块
|
||
if raw.startswith("```"):
|
||
raw = re.sub(r'^```\w*\n?', '', raw)
|
||
raw = re.sub(r'\n?```$', '', raw)
|
||
raw = raw.strip()
|
||
|
||
# 找 JSON 结构的位置
|
||
# 优先找第一个 { }（JSON 对象），用括号匹配确保完整
|
||
first_brace = raw.find('{')
|
||
last_bracket = raw.rfind('[')
|
||
|
||
start = -1
|
||
if first_brace >= 0:
|
||
start = first_brace
|
||
# 从第一个 { 开始，找匹配的 }
|
||
depth = 0
|
||
end = -1
|
||
for i in range(start, len(raw)):
|
||
if raw[i] == '{':
|
||
depth += 1
|
||
elif raw[i] == '}':
|
||
depth -= 1
|
||
if depth == 0:
|
||
end = i + 1
|
||
break
|
||
if end <= start:
|
||
start = -1
|
||
elif last_bracket >= 0:
|
||
start = last_bracket
|
||
end = raw.rfind(']') + 1
|
||
if end <= start:
|
||
start = -1
|
||
|
||
if start < 0:
|
||
start = raw.find('{')
|
||
end = raw.rfind('}') + 1
|
||
|
||
if start < 0 or end <= start:
|
||
raise ValueError(f"无法找到 JSON: {raw[:200]}")
|
||
|
||
json_str = raw[start:end]
|
||
|
||
# 修复常见问题：中文引号替换为英文引号（在JSON值内部的需要转义）
|
||
# 先把 JSON 结构中的英文引号保护起来，再替换中文引号
|
||
json_str = json_str.replace('\u201c', '\u201c').replace('\u201d', '\u201d') # 先保持不变
|
||
json_str = json_str.replace('\u2018', "'").replace('\u2019', "'")
|
||
json_str = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f]', '', json_str)
|
||
json_str = re.sub(r',\s*([}\]])', r'\1', json_str)
|
||
|
||
try:
|
||
return json.loads(json_str)
|
||
except json.JSONDecodeError as e:
|
||
# 如果失败，尝试将中文引号替换为转义的英文引号
|
||
json_str2 = json_str.replace('\u201c', '\\"').replace('\u201d', '\\"')
|
||
try:
|
||
return json.loads(json_str2)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
# 尝试把换行替换为 \n
|
||
json_str2 = json_str.replace('\n', '\\n').replace('\r', '\\r')
|
||
json_str2 = json_str2.replace('\\\\n', '\\n')
|
||
try:
|
||
return json.loads(json_str2)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
# 尝试自动补全缺失的闭合括号（模型输出被截断时常见）
|
||
for extra in ['}', '}}', '}}}', '}]', '}]}']:
|
||
try:
|
||
return json.loads(json_str + extra)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
raise ValueError(f"无法解析 JSON: {json_str[:300]}")
|
||
|
||
|
||
llm_client = LLMClient()
|