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()
|