576 lines
21 KiB
Python
576 lines
21 KiB
Python
"""统一 LLM 客户端 — 多供应商自动切换 + 限流保护。
|
||
|
||
支持: DeepSeek → MiniMax → Gemini
|
||
特点:
|
||
- 限流保护：多用户并发时自动排队
|
||
- 失败切换：一个供应商失败自动切换下一个
|
||
- 统一接口：所有 LLM 调用走这里
|
||
"""
|
||
import os
|
||
import re
|
||
import json
|
||
import time
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
DEFAULT_TIMEOUT_MS = 180_000
|
||
MAX_OUTPUT_TOKENS = 8192
|
||
|
||
|
||
# ─── 配置 ───────────────────────────────────────────────
|
||
|
||
@dataclass
|
||
class LLMProviderConfig:
|
||
name: str
|
||
base_url: str
|
||
model: str
|
||
max_retries: int = 2
|
||
rate_limit: int = 0 # 每分钟请求数，0=无限制
|
||
|
||
|
||
PROVIDERS = {
|
||
"openai": LLMProviderConfig(
|
||
name="openai",
|
||
base_url="https://api.openai.com/v1",
|
||
model="gpt-4o",
|
||
max_retries=2,
|
||
rate_limit=0,
|
||
),
|
||
"deepseek": LLMProviderConfig(
|
||
name="deepseek",
|
||
base_url="https://api.deepseek.com/v1",
|
||
model="deepseek-v4-pro",
|
||
max_retries=2,
|
||
rate_limit=0,
|
||
),
|
||
"minimax": LLMProviderConfig(
|
||
name="minimax",
|
||
base_url="https://api.minimax.chat/v1",
|
||
model="MiniMax-2.7-Flash",
|
||
max_retries=2,
|
||
rate_limit=30,
|
||
),
|
||
"gemini": LLMProviderConfig(
|
||
name="gemini",
|
||
base_url="https://generativelanguage.googleapis.com/v1/models",
|
||
model="gemini-2.5-flash",
|
||
max_retries=1,
|
||
rate_limit=60,
|
||
),
|
||
}
|
||
|
||
|
||
def _format_exception(exc: Exception) -> str:
|
||
message = str(exc) or repr(exc)
|
||
return f"{exc.__class__.__name__}: {message}"
|
||
|
||
|
||
def _parse_timeout_ms(value: Optional[str], default: int = DEFAULT_TIMEOUT_MS) -> int:
|
||
try:
|
||
timeout_ms = int(str(value or "").strip())
|
||
except (TypeError, ValueError):
|
||
return default
|
||
return timeout_ms if timeout_ms > 0 else default
|
||
|
||
|
||
# ─── 限流器 ─────────────────────────────────────────────
|
||
|
||
class RateLimiter:
|
||
"""Token Bucket 限流器。"""
|
||
|
||
def __init__(self, max_tokens: int, refill_ms: int = 60_000):
|
||
self.max_tokens = max_tokens
|
||
self.refill_ms = refill_ms / 1000 # 转换为秒
|
||
self.tokens = [0.0] * max_tokens
|
||
self.last_refill = time.monotonic()
|
||
|
||
async def acquire(self, timeout_ms: int = 30_000) -> bool:
|
||
"""获取一个令牌，超时返回 False。"""
|
||
if self.max_tokens == 0:
|
||
return True
|
||
|
||
start = time.monotonic()
|
||
timeout_s = timeout_ms / 1000
|
||
|
||
while True:
|
||
self._refill_if_needed()
|
||
now = time.monotonic()
|
||
|
||
# 找最早可用的令牌槽位
|
||
for i, t in enumerate(self.tokens):
|
||
if now - t >= self.refill_ms:
|
||
self.tokens[i] = now
|
||
return True
|
||
|
||
# 所有槽位都在使用中，等待最老的释放
|
||
oldest = min(self.tokens)
|
||
wait = min(self.refill_ms - (now - oldest), timeout_s)
|
||
if wait <= 0 or (time.monotonic() - start) >= timeout_s:
|
||
return False
|
||
await asyncio.sleep(min(wait, 2.0))
|
||
|
||
def _refill_if_needed(self):
|
||
now = time.monotonic()
|
||
if now - self.last_refill >= self.refill_ms:
|
||
self.tokens = [0.0] * self.max_tokens
|
||
self.last_refill = now
|
||
|
||
|
||
# ─── 响应类型 ────────────────────────────────────────────
|
||
|
||
@dataclass
|
||
class LLMResponse:
|
||
content: str
|
||
provider: str
|
||
tokens: Optional[dict] = None # {"input": int, "output": int}
|
||
latency_ms: float = 0
|
||
|
||
|
||
def _parse_json_content(content: str):
|
||
"""从模型输出中提取 JSON，兼容代码块、前后说明和尾逗号。"""
|
||
text = (content or "").strip().lstrip("\ufeff")
|
||
if not text:
|
||
raise ValueError("模型返回了空内容")
|
||
|
||
candidates = re.findall(
|
||
r"```(?:json)?\s*([\s\S]*?)```",
|
||
text,
|
||
flags=re.IGNORECASE,
|
||
)
|
||
candidates.append(text)
|
||
|
||
decoder = json.JSONDecoder(strict=False)
|
||
last_error = None
|
||
for candidate in candidates:
|
||
cleaned = re.sub(r",\s*([\]}])", r"\1", candidate.strip())
|
||
start_positions = [0]
|
||
start_positions.extend(
|
||
index for index, char in enumerate(cleaned) if char in "[{" and index != 0
|
||
)
|
||
for start in start_positions:
|
||
try:
|
||
data, _end = decoder.raw_decode(cleaned[start:])
|
||
return data
|
||
except json.JSONDecodeError as exc:
|
||
last_error = exc
|
||
|
||
detail = str(last_error) if last_error else "未找到 JSON 对象"
|
||
raise ValueError(f"模型输出不是有效 JSON（{detail}）")
|
||
|
||
|
||
# ─── 单供应商调用 ─────────────────────────────────────────
|
||
|
||
async def _call_provider(
|
||
config: LLMProviderConfig,
|
||
api_key: str,
|
||
messages: list[dict],
|
||
timeout_ms: int = 60_000,
|
||
json_mode: bool = False,
|
||
temperature: float = 0.3,
|
||
) -> LLMResponse:
|
||
"""调用单个 LLM 供应商。"""
|
||
start = time.monotonic()
|
||
timeout_s = timeout_ms / 1000
|
||
|
||
headers = {"Content-Type": "application/json"}
|
||
|
||
if config.name in ("deepseek", "minimax"):
|
||
# OpenAI 兼容格式
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
if config.name == "minimax":
|
||
url = f"{config.base_url}/text/chatcompletion_v2"
|
||
else:
|
||
url = f"{config.base_url}/chat/completions"
|
||
body = {
|
||
"model": config.model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": MAX_OUTPUT_TOKENS,
|
||
}
|
||
if json_mode and config.name == "deepseek":
|
||
body["response_format"] = {"type": "json_object"}
|
||
# DeepSeek V4 默认开启 Thinking Mode。结构化提取若不显式关闭，
|
||
# 可能耗尽输出额度后只返回 reasoning_content，content 为空。
|
||
body["thinking"] = {"type": "disabled"}
|
||
elif config.name == "gemini":
|
||
# Gemini 格式
|
||
model_part = f"{config.model}:generateContent" if ":" not in config.model else config.model
|
||
url = f"{config.base_url}/{model_part}?key={api_key}"
|
||
contents = []
|
||
for m in messages:
|
||
role = "model" if m["role"] == "assistant" else "user"
|
||
contents.append({"role": role, "parts": [{"text": m["content"]}]})
|
||
body = {
|
||
"contents": contents,
|
||
"generationConfig": {"temperature": temperature, "maxOutputTokens": MAX_OUTPUT_TOKENS},
|
||
}
|
||
if json_mode:
|
||
body["generationConfig"]["responseMimeType"] = "application/json"
|
||
else:
|
||
# 自定义供应商：OpenAI 兼容格式
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
base = config.base_url.rstrip("/")
|
||
url = f"{base}/chat/completions"
|
||
body = {
|
||
"model": config.model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": MAX_OUTPUT_TOKENS,
|
||
}
|
||
if json_mode:
|
||
body["response_format"] = {"type": "json_object"}
|
||
|
||
async with httpx.AsyncClient(timeout=timeout_s) as client:
|
||
resp = await client.post(url, json=body, headers=headers)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
latency_ms = (time.monotonic() - start) * 1000
|
||
|
||
# 解析响应
|
||
if config.name == "gemini":
|
||
content = ""
|
||
candidates = data.get("candidates", [])
|
||
if candidates:
|
||
parts = candidates[0].get("content", {}).get("parts", [])
|
||
if parts:
|
||
content = parts[0].get("text", "")
|
||
usage_meta = data.get("usageMetadata")
|
||
tokens = None
|
||
if usage_meta:
|
||
tokens = {
|
||
"input": usage_meta.get("promptTokenCount", 0),
|
||
"output": usage_meta.get("candidatesTokenCount", 0),
|
||
}
|
||
else:
|
||
# OpenAI 兼容格式（deepseek / minimax / 自定义供应商）
|
||
content = ""
|
||
choices = data.get("choices", [])
|
||
if choices:
|
||
content = choices[0].get("message", {}).get("content", "")
|
||
usage = data.get("usage")
|
||
tokens = None
|
||
if usage:
|
||
tokens = {
|
||
"input": usage.get("prompt_tokens", 0),
|
||
"output": usage.get("completion_tokens", 0),
|
||
}
|
||
|
||
return LLMResponse(
|
||
content=content,
|
||
provider=config.name,
|
||
tokens=tokens,
|
||
latency_ms=latency_ms,
|
||
)
|
||
|
||
|
||
async def _call_dify(
|
||
model: str,
|
||
messages: list[dict],
|
||
timeout_ms: int = 60_000,
|
||
) -> LLMResponse:
|
||
"""通过 Dify Chat API 调用模型（复用 Dify 已配置的 API Key）。"""
|
||
import flask
|
||
dify_base_url = flask.current_app.config.get("DIFY_BASE_URL", os.getenv("DIFY_BASE_URL", "http://localhost:5001"))
|
||
dify_api_key = flask.current_app.config.get("BAODAN_CHAT_API_KEY", os.getenv("DIFY_CHAT_APP_API_KEY", ""))
|
||
|
||
if not dify_api_key:
|
||
raise RuntimeError("未配置 DIFY_CHAT_APP_API_KEY，无法使用 Dify 模式")
|
||
|
||
# 合并 messages 为 query（Dify Chat API 不支持多轮 messages 格式）
|
||
system_parts = []
|
||
user_parts = []
|
||
for m in messages:
|
||
if m["role"] == "system":
|
||
system_parts.append(m["content"])
|
||
else:
|
||
user_parts.append(m["content"])
|
||
query = "\n\n".join(user_parts)
|
||
if system_parts:
|
||
query = "\n\n".join(system_parts) + "\n\n" + query
|
||
|
||
# 如果有指定模型，在 query 前加上模型提示
|
||
if model:
|
||
query = f"[请使用模型 {model} 回答]\n\n{query}"
|
||
|
||
start = time.monotonic()
|
||
timeout_s = timeout_ms / 1000
|
||
url = f"{dify_base_url.rstrip('/')}/v1/chat-messages"
|
||
headers = {
|
||
"Authorization": f"Bearer {dify_api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
body = {
|
||
"inputs": {},
|
||
"query": query,
|
||
"response_mode": "blocking",
|
||
"user": "insurance-system",
|
||
}
|
||
|
||
async with httpx.AsyncClient(timeout=timeout_s) as client:
|
||
resp = await client.post(url, json=body, headers=headers)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
latency_ms = (time.monotonic() - start) * 1000
|
||
content = data.get("answer", "")
|
||
tokens = None
|
||
usage = data.get("metadata", {}).get("usage", {})
|
||
if usage:
|
||
tokens = {
|
||
"input": usage.get("prompt_tokens", 0),
|
||
"output": usage.get("completion_tokens", 0),
|
||
}
|
||
|
||
return LLMResponse(content=content, provider="dify", tokens=tokens, latency_ms=latency_ms)
|
||
|
||
|
||
# ─── 统一客户端 ───────────────────────────────────────────
|
||
|
||
class LLMClient:
|
||
"""多供应商 LLM 客户端，支持自动切换和速率限制。
|
||
|
||
config_prefix: 数据库配置键前缀，如 "ppt" 读取 ppt_llm_*，"poster" 读取 poster_llm_*。
|
||
"""
|
||
|
||
def __init__(self, config_prefix: str = "ppt"):
|
||
self._configs: list[tuple[LLMProviderConfig, str]] = []
|
||
self._limiters: dict[str, RateLimiter] = {}
|
||
self._active_idx = 0
|
||
self._config_prefix = config_prefix
|
||
self._timeout_ms = _parse_timeout_ms(os.getenv(f"{config_prefix.upper()}_LLM_TIMEOUT_MS"))
|
||
self._db_config_time: float = 0 # 上次从数据库加载配置的时间戳
|
||
self._db_config_ttl: float = 60 # 配置缓存有效期（秒）
|
||
|
||
self._load_env_config()
|
||
|
||
def _load_env_config(self):
|
||
"""从环境变量加载默认配置。"""
|
||
prefix = self._config_prefix.upper()
|
||
deepseek_key = os.environ.get(f"{prefix}_LLM_API_KEY") or os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("OPENAI_API_KEY", "")
|
||
minimax_key = os.environ.get("MINIMAX_API_KEY", "")
|
||
gemini_key = os.environ.get("GEMINI_API_KEY", "")
|
||
|
||
if deepseek_key:
|
||
self._configs.append((PROVIDERS["deepseek"], deepseek_key))
|
||
self._limiters["deepseek"] = RateLimiter(PROVIDERS["deepseek"].rate_limit)
|
||
if minimax_key:
|
||
self._configs.append((PROVIDERS["minimax"], minimax_key))
|
||
self._limiters["minimax"] = RateLimiter(PROVIDERS["minimax"].rate_limit)
|
||
if gemini_key:
|
||
self._configs.append((PROVIDERS["gemini"], gemini_key))
|
||
self._limiters["gemini"] = RateLimiter(PROVIDERS["gemini"].rate_limit)
|
||
|
||
if not self._configs:
|
||
logger.debug("[LLMClient] 未从环境变量加载 API Key，将在首次调用时从数据库读取配置")
|
||
|
||
def _try_load_db_config(self):
|
||
"""尝试从数据库加载模型配置（带缓存，不阻塞）。"""
|
||
now = time.monotonic()
|
||
if now - self._db_config_time < self._db_config_ttl:
|
||
return
|
||
self._db_config_time = now
|
||
prefix = self._config_prefix
|
||
try:
|
||
from insurance.models.system_setting import SystemSetting
|
||
settings = {s.key: s.value for s in SystemSetting.query.filter(
|
||
SystemSetting.key.in_([
|
||
f"{prefix}_llm_provider", f"{prefix}_llm_model",
|
||
f"{prefix}_llm_api_key", f"{prefix}_llm_base_url",
|
||
f"{prefix}_llm_timeout_ms",
|
||
])
|
||
).all()}
|
||
self._timeout_ms = _parse_timeout_ms(
|
||
settings.get(f"{prefix}_llm_timeout_ms"),
|
||
_parse_timeout_ms(os.getenv(f"{prefix.upper()}_LLM_TIMEOUT_MS")),
|
||
)
|
||
provider = settings.get(f"{prefix}_llm_provider", "").strip()
|
||
if not provider:
|
||
return
|
||
|
||
model = settings.get(f"{prefix}_llm_model", "").strip()
|
||
|
||
# Dify 模式：通过 Dify Chat API 调用，无需独立 API Key
|
||
if provider == "dify":
|
||
cfg = LLMProviderConfig(
|
||
name="dify", base_url="", model=model or "",
|
||
)
|
||
self._configs = [(cfg, "")]
|
||
self._limiters = {"dify": RateLimiter(0)}
|
||
self._active_idx = 0
|
||
logger.info(f"[LLMClient:{prefix}] 使用 Dify 模式: {model}")
|
||
return
|
||
|
||
api_key = settings.get(f"{prefix}_llm_api_key", "").strip()
|
||
if not api_key:
|
||
return
|
||
|
||
base_url = settings.get(f"{prefix}_llm_base_url", "").strip()
|
||
|
||
# 内置供应商：替换对应配置
|
||
if provider in PROVIDERS and not base_url:
|
||
cfg = PROVIDERS[provider]
|
||
if model:
|
||
cfg = LLMProviderConfig(
|
||
name=cfg.name, base_url=cfg.base_url, model=model,
|
||
max_retries=cfg.max_retries, rate_limit=cfg.rate_limit,
|
||
)
|
||
self._configs = [(cfg, api_key)]
|
||
self._limiters = {provider: RateLimiter(cfg.rate_limit)}
|
||
self._active_idx = 0
|
||
logger.info(f"[LLMClient:{prefix}] 使用数据库配置: {provider}/{cfg.model}")
|
||
return
|
||
|
||
# 自定义供应商
|
||
if not base_url:
|
||
return
|
||
|
||
# 安全校验：拒绝私网地址（SEC-P1-02）
|
||
from insurance.utils.security import is_safe_base_url
|
||
is_safe, err_msg = is_safe_base_url(base_url)
|
||
if not is_safe:
|
||
logger.warning(f"[LLMClient:{prefix}] 不安全的 Base URL: {err_msg}")
|
||
return
|
||
|
||
cfg = LLMProviderConfig(
|
||
name=provider, base_url=base_url, model=model or "gpt-4o-mini",
|
||
)
|
||
self._configs = [(cfg, api_key)]
|
||
self._limiters = {provider: RateLimiter(0)}
|
||
self._active_idx = 0
|
||
logger.info(f"[LLMClient:{prefix}] 使用自定义模型: {provider}/{cfg.model}")
|
||
except Exception:
|
||
pass # 无 Flask 上下文或数据库不可用，使用环境变量配置
|
||
|
||
async def chat(self, prompt: str, system_prompt: str = "") -> LLMResponse:
|
||
"""简单聊天。"""
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
messages.append({"role": "user", "content": prompt})
|
||
return await self._call(messages)
|
||
|
||
async def structured_output(
|
||
self,
|
||
prompt: str,
|
||
system_prompt: str = "",
|
||
schema: Optional[dict] = None,
|
||
temperature: float = 0.3,
|
||
) -> tuple[dict, LLMResponse]:
|
||
"""结构化输出（返回 JSON）。返回 (parsed_data, response)。
|
||
|
||
Args:
|
||
temperature: 输出温度。结构化提取建议 0，分析任务可用 0.3。
|
||
"""
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
|
||
full_prompt = prompt
|
||
if schema:
|
||
full_prompt += f"\n\n请以JSON格式输出，格式如下：\n{json.dumps(schema, ensure_ascii=False, indent=2)}"
|
||
full_prompt += "\n重要：只输出JSON，不要任何额外文字。"
|
||
messages.append({"role": "user", "content": full_prompt})
|
||
|
||
response = await self._call(messages, json_mode=True, temperature=temperature)
|
||
try:
|
||
return _parse_json_content(response.content), response
|
||
except ValueError as first_error:
|
||
logger.warning(
|
||
"[LLMClient] 首次结构化输出无效，正在请求模型纠正: %s",
|
||
first_error,
|
||
)
|
||
|
||
repair_messages = messages + [
|
||
{"role": "assistant", "content": (response.content or "")[:12000]},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"上一次回答不是可解析的完整 JSON。请重新输出完整、有效的 JSON；"
|
||
"不要使用 Markdown 代码块，不要解释，不要省略字段。"
|
||
),
|
||
},
|
||
]
|
||
repaired_response = await self._call(repair_messages, json_mode=True, temperature=temperature)
|
||
try:
|
||
return _parse_json_content(repaired_response.content), repaired_response
|
||
except ValueError as repair_error:
|
||
raise ValueError(f"[LLMClient] JSON 解析失败，纠正重试仍无效: {repair_error}")
|
||
|
||
async def _call(
|
||
self,
|
||
messages: list[dict],
|
||
attempt: int = 0,
|
||
json_mode: bool = False,
|
||
temperature: float = 0.3,
|
||
) -> LLMResponse:
|
||
"""多供应商自动切换调用。"""
|
||
self._try_load_db_config()
|
||
if not self._configs:
|
||
raise RuntimeError(
|
||
"未配置任何 LLM API Key，请在管理后台 > 系统配置中设置 "
|
||
"ppt_llm_provider / ppt_llm_api_key，或设置环境变量 DEEPSEEK_API_KEY / OPENAI_API_KEY"
|
||
)
|
||
|
||
start_idx = self._active_idx
|
||
tried = set()
|
||
|
||
for i in range(len(self._configs)):
|
||
idx = (start_idx + i) % len(self._configs)
|
||
config, api_key = self._configs[idx]
|
||
|
||
if config.name in tried:
|
||
continue
|
||
tried.add(config.name)
|
||
|
||
# Dify 模式：通过 Dify Chat API 调用
|
||
if config.name == "dify":
|
||
try:
|
||
return await _call_dify(config.model, messages, timeout_ms=self._timeout_ms)
|
||
except Exception as e:
|
||
logger.warning(f"[LLMClient] Dify 调用失败: {_format_exception(e)}")
|
||
continue
|
||
|
||
# 限流
|
||
limiter = self._limiters.get(config.name)
|
||
if limiter:
|
||
acquired = await limiter.acquire(timeout_ms=30_000)
|
||
if not acquired:
|
||
logger.warning(f"[LLMClient] {config.name} 限流超时")
|
||
continue
|
||
|
||
# 调用
|
||
try:
|
||
response = await _call_provider(
|
||
config,
|
||
api_key,
|
||
messages,
|
||
timeout_ms=self._timeout_ms,
|
||
json_mode=json_mode,
|
||
temperature=temperature,
|
||
)
|
||
self._active_idx = idx
|
||
return response
|
||
except Exception as e:
|
||
logger.warning(f"[LLMClient] {config.name} 失败: {_format_exception(e)}")
|
||
if attempt < 3 and i < len(self._configs) - 1:
|
||
self._active_idx = (idx + 1) % len(self._configs)
|
||
|
||
raise RuntimeError(f"所有 LLM 供应商均失败，请检查 API Key、Base URL、模型名称或超时设置（当前 {self._timeout_ms}ms）")
|
||
|
||
def get_status(self) -> dict:
|
||
"""获取当前供应商信息。"""
|
||
available = [c[0].name for c in self._configs]
|
||
active = self._configs[self._active_idx][0].name if self._configs else "none"
|
||
return {"available": available, "active": active}
|
||
|
||
|
||
# ─── 单例 ────────────────────────────────────────────────
|
||
llm_client = LLMClient(config_prefix="ppt")
|
||
poster_llm_client = LLMClient(config_prefix="poster")
|