101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
"""LLM Key 兜底：主 Key 余额不足（402 Insufficient Balance）时自动切到兜底 Key。
|
||
|
||
- 主 Key：设置页保存的 llm_api_key / .env 的 DEEPSEEK_API_KEY（千问：DASHSCOPE_API_KEY）
|
||
- 兜底 Key：.env 的 DEEPSEEK_API_KEY_FALLBACK / DASHSCOPE_API_KEY_FALLBACK
|
||
- 切换状态记在 settings 表 `llm_primary_exhausted_at`，重启不丢；超过 LLM_PRIMARY_RETRY_HOURS 后下一次调用先试主 Key，
|
||
仍 402 就再记一次（代价是每个重试周期一次失败调用）。
|
||
|
||
三条调用路径都走这里：NLU / 沉淀提炼（nlu._call_llm）、Agent（llm_factory.FallbackChatModel）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from datetime import datetime, timedelta
|
||
|
||
from backend.config import settings
|
||
from backend.db import database as db
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
EXHAUSTED_KEY = "llm_primary_exhausted_at"
|
||
_TIME_FMT = "%Y-%m-%d %H:%M:%S"
|
||
|
||
# 进程内缓存，避免每次调用都查库；跨进程以 settings 表为准
|
||
_exhausted_at: datetime | None = None
|
||
_loaded = False
|
||
|
||
|
||
def _load() -> None:
|
||
global _exhausted_at, _loaded
|
||
if _loaded:
|
||
return
|
||
_loaded = True
|
||
raw = db.get_setting(EXHAUSTED_KEY)
|
||
try:
|
||
_exhausted_at = datetime.strptime(raw, _TIME_FMT) if raw else None
|
||
except ValueError:
|
||
_exhausted_at = None
|
||
|
||
|
||
def fallback_key(provider: str) -> str:
|
||
if provider == "qwen":
|
||
return settings.dashscope_api_key_fallback
|
||
return settings.deepseek_api_key_fallback
|
||
|
||
|
||
def primary_exhausted() -> bool:
|
||
"""主 Key 是否仍处于「已耗尽」状态（超过重试周期视为可以再试）。"""
|
||
_load()
|
||
if not _exhausted_at:
|
||
return False
|
||
if datetime.now() - _exhausted_at > timedelta(hours=max(0.05, settings.llm_primary_retry_hours)):
|
||
return False
|
||
return True
|
||
|
||
|
||
def mark_primary_exhausted(reason: str = "") -> None:
|
||
global _exhausted_at
|
||
_load()
|
||
_exhausted_at = datetime.now()
|
||
db.set_setting(EXHAUSTED_KEY, _exhausted_at.strftime(_TIME_FMT))
|
||
logger.warning("LLM 主 Key 余额不足，切到兜底 Key。%s", reason[:200])
|
||
|
||
|
||
def clear_primary_exhausted() -> None:
|
||
"""主 Key 调用成功后清掉标记（充值后自动恢复）。"""
|
||
global _exhausted_at
|
||
_load()
|
||
if _exhausted_at is None:
|
||
return
|
||
_exhausted_at = None
|
||
db.set_setting(EXHAUSTED_KEY, "")
|
||
logger.info("LLM 主 Key 恢复可用")
|
||
|
||
|
||
def is_balance_error(error: BaseException | str) -> bool:
|
||
"""DeepSeek 402：{'error': {'message': 'Insufficient Balance', …}}；千问余额不足也带 402 / Arrearage。"""
|
||
status = getattr(error, "status_code", None)
|
||
text = str(error)
|
||
return status == 402 or "Error code: 402" in text or "Insufficient Balance" in text or "Arrearage" in text
|
||
|
||
|
||
def choose(provider: str, primary: str) -> tuple[str, bool]:
|
||
"""返回 (本次应使用的 Key, 是否兜底)。"""
|
||
fallback = fallback_key(provider)
|
||
if fallback and fallback != primary and primary_exhausted():
|
||
return fallback, True
|
||
return primary or fallback, bool(primary == "" and fallback)
|
||
|
||
|
||
def status(provider: str, primary: str) -> dict:
|
||
_load()
|
||
fallback = fallback_key(provider)
|
||
_key, using = choose(provider, primary)
|
||
return {
|
||
"fallback_configured": bool(fallback),
|
||
"using_fallback": using,
|
||
"primary_exhausted_at": _exhausted_at.strftime(_TIME_FMT) if _exhausted_at else None,
|
||
"retry_hours": settings.llm_primary_retry_hours,
|
||
}
|