317 lines
11 KiB
Python
317 lines
11 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__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 配置 ───────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class LLMProviderConfig:
|
|||
|
|
name: str
|
|||
|
|
base_url: str
|
|||
|
|
model: str
|
|||
|
|
max_retries: int = 2
|
|||
|
|
rate_limit: int = 0 # 每分钟请求数，0=无限制
|
|||
|
|
|
|||
|
|
|
|||
|
|
PROVIDERS = {
|
|||
|
|
"deepseek": LLMProviderConfig(
|
|||
|
|
name="deepseek",
|
|||
|
|
base_url="https://api.deepseek.com/v1",
|
|||
|
|
model="deepseek-chat",
|
|||
|
|
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,
|
|||
|
|
),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 限流器 ─────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
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:
|
|||
|
|
raise ValueError(f"Unknown provider: {config.name}")
|
|||
|
|
|
|||
|
|
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 in ("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),
|
|||
|
|
}
|
|||
|
|
else:
|
|||
|
|
# 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),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return LLMResponse(
|
|||
|
|
content=content,
|
|||
|
|
provider=config.name,
|
|||
|
|
tokens=tokens,
|
|||
|
|
latency_ms=latency_ms,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ─── 统一客户端 ───────────────────────────────────────────
|
|||
|
|
|
|||
|
|
class LLMClient:
|
|||
|
|
"""多供应商 LLM 客户端，支持自动切换和速率限制。"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._configs: list[tuple[LLMProviderConfig, str]] = []
|
|||
|
|
self._limiters: dict[str, RateLimiter] = {}
|
|||
|
|
self._active_idx = 0
|
|||
|
|
|
|||
|
|
# 从环境变量加载
|
|||
|
|
deepseek_key = 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.warning("[LLMClient] 未配置任何 API Key，使用 mock 模式")
|
|||
|
|
|
|||
|
|
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:
|
|||
|
|
"""多供应商自动切换调用。"""
|
|||
|
|
if not self._configs:
|
|||
|
|
return LLMResponse(content="{}", provider="mock", latency_ms=0)
|
|||
|
|
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# 限流
|
|||
|
|
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=60_000)
|
|||
|
|
self._active_idx = idx
|
|||
|
|
return response
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"[LLMClient] {config.name} 失败: {e}")
|
|||
|
|
if attempt < 3 and i < len(self._configs) - 1:
|
|||
|
|
self._active_idx = (idx + 1) % len(self._configs)
|
|||
|
|
|
|||
|
|
raise RuntimeError("所有 LLM 供应商均失败")
|
|||
|
|
|
|||
|
|
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()
|