wsb1224
0ffd9dffdf
The warning fired at init time when no env var API keys were set, but the actual config was loaded from the database on first use. Changed to debug to avoid confusion in logs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
510 lines
18 KiB
Python
510 lines
18 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
|
||
|
||
|
||
# ─── 配置 ───────────────────────────────────────────────
|
||
|
||
@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
|
||
|
||
|
||
# ─── 单供应商调用 ─────────────────────────────────────────
|
||
|
||
async def _call_provider(
|
||
config: LLMProviderConfig,
|
||
api_key: str,
|
||
messages: list[dict],
|
||
timeout_ms: int = 60_000,
|
||
) -> 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": 0.3,
|
||
"max_tokens": 4096,
|
||
}
|
||
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": 0.3, "maxOutputTokens": 4096},
|
||
}
|
||
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": 0.3,
|
||
"max_tokens": 4096,
|
||
}
|
||
|
||
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,
|
||
) -> tuple[dict, LLMResponse]:
|
||
"""结构化输出（返回 JSON）。返回 (parsed_data, response)。"""
|
||
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
|
||
json_str = response.content.strip()
|
||
# 尝试提取 markdown 代码块或裸 JSON
|
||
match = re.search(r"```(?:json)?\s*([\s\S]*?)```|(\{[\s\S]*\}|\[[\s\S]*\])$", json_str)
|
||
if match:
|
||
json_str = match.group(1) or match.group(2)
|
||
|
||
try:
|
||
data = json.loads(json_str)
|
||
return data, response
|
||
except json.JSONDecodeError:
|
||
# 清理尾逗号
|
||
json_str = re.sub(r",\s*([\]}])", r"\1", json_str)
|
||
try:
|
||
data = json.loads(json_str)
|
||
return data, response
|
||
except json.JSONDecodeError:
|
||
raise ValueError(f"[LLMClient] JSON 解析失败: {json_str[:200]}")
|
||
|
||
async def _call(self, messages: list[dict], attempt: int = 0) -> 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)
|
||
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")
|