feat(#875): LLM Router shadow mode + consolidate litellm master key
This commit is contained in:
QihuAI Design committed
1 parent
037c965c26
commit
180409d317
3 files changed
+390
-5
No files matched your search
@@ -0,0 +1,340 @@
|
||||
"""LLM Router —— 统一模型路由层（#875 · 影子模式先行）
|
||||
|
||||
设计基线（leagent_accel_belt_design.md §C）：
|
||||
- 统一 Web Claw 与 IM Bot Claw 的模型选择，消除两端重复的 failover / 主密钥硬编码。
|
||||
- 决策顺序：显式指定 → 租户配额 → 能力过滤 → 成本分层 → 健康 failover（熔断）。
|
||||
- 灰度策略：**默认影子模式**（`LLM_ROUTER_SHADOW=1`）—— Router 计算结果只记日志 +
|
||||
`claw_type` 标签 metrics，实际仍走旧逻辑；一致率 ≥98% 后把开关置 0 切主路径。
|
||||
- 上下文三元组（tenant_id + claw_type + trace_id）经 `contextvars` 注入，零签名侵入
|
||||
（复用 observability.py 的 trace_id 同款机制）。
|
||||
|
||||
对外 API：
|
||||
set_route_ctx(tenant_id="", claw_type="web") # 在入口注入上下文
|
||||
get_route_ctx() -> dict # {tenant_id, claw_type, trace_id}
|
||||
route(capability, db, *, prefer_model_id=None, legacy_primary=None, shadow=None)
|
||||
-> RouteDecision(models: list[AIModel], primary, reason, matched_legacy, ...)
|
||||
resolve_api_key(api_url, model_api_key, provider_api_key) -> str|None
|
||||
# 收敛 LiteLLM 主密钥硬编码：litellm 端点统一返回 LITELLM_MASTER_KEY
|
||||
report_failure(model_id, claw_type="web") / report_success(model_id)
|
||||
# 健康熔断：Redis 记录连续失败（TTL 60s），无 Redis 时降级内存
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import threading
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 影子模式开关：默认开启 → 只观测，不改流量
|
||||
# ---------------------------------------------------------------------------
|
||||
SHADOW = os.getenv("LLM_ROUTER_SHADOW", "1") not in ("0", "false", "False", "")
|
||||
|
||||
_LITELLM_MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY") or "sk-litellm-master-key-aiflow-2026"
|
||||
_LITELLM_HOST_MARK = "qihu-aiflow-litellm"
|
||||
|
||||
# 媒体能力（含 "text" 但本质是媒体模型，不能当文本 failover 备选）
|
||||
_MEDIA_CAPS = {"image", "video", "audio", "asr", "embedding", "rerank", "ocr"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 上下文三元组：tenant_id + claw_type（trace_id 复用 observability）
|
||||
# ---------------------------------------------------------------------------
|
||||
_route_tenant_var: ContextVar[str] = ContextVar("aiflow_route_tenant", default="")
|
||||
_route_claw_var: ContextVar[str] = ContextVar("aiflow_route_claw", default="web")
|
||||
|
||||
|
||||
def set_route_ctx(*, tenant_id: str = "", claw_type: str = "web") -> None:
|
||||
"""在请求入口注入路由上下文。tenant_id/claw_type 为空时保留默认值。"""
|
||||
if tenant_id:
|
||||
_route_tenant_var.set(tenant_id)
|
||||
_route_claw_var.set(claw_type or "web")
|
||||
|
||||
|
||||
def get_route_ctx() -> dict:
|
||||
try:
|
||||
from app.core.observability import get_trace_id
|
||||
_trace = get_trace_id()
|
||||
except Exception:
|
||||
_trace = "-"
|
||||
return {
|
||||
"tenant_id": _route_tenant_var.get(),
|
||||
"claw_type": _route_claw_var.get(),
|
||||
"trace_id": _trace,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 决策结果
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class RouteDecision:
|
||||
capability: str
|
||||
claw_type: str
|
||||
tenant_id: str
|
||||
models: list = field(default_factory=list) # 有序候选链 list[AIModel]
|
||||
primary: object | None = None # 首选 AIModel
|
||||
reason: str = ""
|
||||
shadow: bool = True
|
||||
matched_legacy: Optional[bool] = None # 与旧逻辑首选是否一致
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 能力过滤：复用 capability_registry + 内置 _is_text_only
|
||||
# ---------------------------------------------------------------------------
|
||||
def _is_text_only(m) -> bool:
|
||||
"""纯文本模型才能作为文本 failover 备选（消除混入图像/语音模型导致空内容的坑）。"""
|
||||
caps = m.capabilities or []
|
||||
return "text" in caps and not any(c in _MEDIA_CAPS for c in caps)
|
||||
|
||||
|
||||
def _fallbacks(db, exclude_ids: set, limit: int = 4) -> list:
|
||||
"""对齐 IM Bot 的备选构造：全部启用、按 is_default 降序、仅纯文本、取前 limit。"""
|
||||
from app.models import AIModel
|
||||
rows = (
|
||||
db.query(AIModel)
|
||||
.filter(AIModel.is_enabled == True) # noqa: E712
|
||||
.order_by(AIModel.is_default.desc())
|
||||
.all()
|
||||
)
|
||||
out = [m for m in rows if m.id not in exclude_ids and _is_text_only(m)]
|
||||
return out[:limit]
|
||||
|
||||
|
||||
def _tenant_authorized(m, tenant_id: str) -> bool:
|
||||
"""模型归属校验：全局模型(无 tenant)任何人可用；否则仅所属租户可用。"""
|
||||
mt = getattr(m, "tenant_id", None)
|
||||
if not mt:
|
||||
return True
|
||||
if not tenant_id:
|
||||
return False
|
||||
return mt == tenant_id
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 健康熔断：Redis 连续失败计数（TTL 60s），无 Redis 降级内存
|
||||
# ---------------------------------------------------------------------------
|
||||
_BREAKER_TTL = 60.0
|
||||
_BREAKER_LOCAL: dict = {}
|
||||
_BREAKER_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def _redis():
|
||||
try:
|
||||
import redis
|
||||
from app.core.config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD
|
||||
return redis.Redis(
|
||||
host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD,
|
||||
socket_connect_timeout=1, socket_timeout=1,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def report_failure(model_id: str, claw_type: str = "web") -> None:
|
||||
if not model_id:
|
||||
return
|
||||
r = _redis()
|
||||
key = f"aiflow:router:fail:{model_id}"
|
||||
if r:
|
||||
try:
|
||||
r.incr(key)
|
||||
r.expire(key, int(_BREAKER_TTL))
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
with _BREAKER_LOCK:
|
||||
e = _BREAKER_LOCAL.setdefault(model_id, {"fails": 0, "until": 0.0})
|
||||
e["fails"] += 1
|
||||
e["until"] = time.time() + _BREAKER_TTL
|
||||
|
||||
|
||||
def report_success(model_id: str) -> None:
|
||||
if not model_id:
|
||||
return
|
||||
r = _redis()
|
||||
key = f"aiflow:router:fail:{model_id}"
|
||||
if r:
|
||||
try:
|
||||
r.delete(key)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
with _BREAKER_LOCK:
|
||||
_BREAKER_LOCAL.pop(model_id, None)
|
||||
|
||||
|
||||
def _fail_count(model_id: str) -> int:
|
||||
if not model_id:
|
||||
return 0
|
||||
r = _redis()
|
||||
key = f"aiflow:router:fail:{model_id}"
|
||||
if r:
|
||||
try:
|
||||
v = r.get(key)
|
||||
return int(v) if v else 0
|
||||
except Exception:
|
||||
pass
|
||||
with _BREAKER_LOCK:
|
||||
e = _BREAKER_LOCAL.get(model_id)
|
||||
if e and e["until"] > time.time():
|
||||
return e["fails"]
|
||||
_BREAKER_LOCAL.pop(model_id, None)
|
||||
return 0
|
||||
|
||||
|
||||
def _apply_breaker(models: list, claw_type: str) -> list:
|
||||
"""失败次数少的模型优先；0 失败保持原序。"""
|
||||
if not models:
|
||||
return models
|
||||
scored = [( _fail_count(getattr(m, "id", "")), i, m) for i, m in enumerate(models)]
|
||||
scored.sort(key=lambda x: (x[0], x[1]))
|
||||
return [m for _, _, m in scored]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 主密钥收敛：litellm 端点统一用 master key
|
||||
# ---------------------------------------------------------------------------
|
||||
def resolve_api_key(api_url: Optional[str], model_api_key: Optional[str],
|
||||
provider_api_key: Optional[str]) -> Optional[str]:
|
||||
if api_url and _LITELLM_HOST_MARK in api_url:
|
||||
return _LITELLM_MASTER_KEY
|
||||
return model_api_key or provider_api_key
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# metrics：复用 otel meter，带 claw_type 标签（失败安全）
|
||||
# ---------------------------------------------------------------------------
|
||||
_counter = None
|
||||
_match_counter = None
|
||||
_hist = None
|
||||
|
||||
|
||||
def _metric_objs():
|
||||
global _counter, _match_counter, _hist
|
||||
if _counter is not None:
|
||||
return _counter, _match_counter, _hist
|
||||
try:
|
||||
from app.core.otel import get_meter
|
||||
m = get_meter()
|
||||
if m is None:
|
||||
return None, None, None
|
||||
_counter = m.create_counter("aiflow_llm_route_total", "LLM router decisions (shadow+live)")
|
||||
_match_counter = m.create_counter(
|
||||
"aiflow_llm_route_legacy_match_total", "Router vs legacy primary consistency"
|
||||
)
|
||||
_hist = m.create_histogram("aiflow_llm_route_ms", "Router decision latency ms")
|
||||
except Exception:
|
||||
return None, None, None
|
||||
return _counter, _match_counter, _hist
|
||||
|
||||
|
||||
def _emit_metrics(capability: str, claw_type: str, primary, shadow: bool, elapsed_ms: float) -> None:
|
||||
c, mc, h = _metric_objs()
|
||||
if c is None:
|
||||
return
|
||||
try:
|
||||
c.add(1, {
|
||||
"claw_type": claw_type,
|
||||
"capability": capability,
|
||||
"decision": primary.name if primary else "none",
|
||||
"shadow": str(shadow),
|
||||
})
|
||||
if h:
|
||||
h.record(elapsed_ms, {"claw_type": claw_type})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _emit_match(claw_type: str, matched: bool) -> None:
|
||||
c, mc, h = _metric_objs()
|
||||
if mc is None:
|
||||
return
|
||||
try:
|
||||
mc.add(1, {"claw_type": claw_type, "matched": "true" if matched else "false"})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# route：同步（与 resolve_auto_model 一致，可直接在 async 引擎内调用）
|
||||
# ---------------------------------------------------------------------------
|
||||
def route(capability: str, db, *, prefer_model_id: Optional[str] = None,
|
||||
legacy_primary: Optional[str] = None, shadow: Optional[bool] = None) -> RouteDecision:
|
||||
_t0 = time.perf_counter()
|
||||
shadow = SHADOW if shadow is None else shadow
|
||||
ctx = get_route_ctx()
|
||||
claw_type = ctx["claw_type"]
|
||||
tenant_id = ctx["tenant_id"]
|
||||
reasons: list = []
|
||||
|
||||
models: list = []
|
||||
try:
|
||||
from app.models import AIModel
|
||||
from app.capability_registry import resolve_auto_model
|
||||
|
||||
# 1) 显式指定
|
||||
if prefer_model_id:
|
||||
m = (
|
||||
db.query(AIModel)
|
||||
.filter(AIModel.id == prefer_model_id, AIModel.is_enabled == True) # noqa: E712
|
||||
.first()
|
||||
)
|
||||
if m and _tenant_authorized(m, tenant_id):
|
||||
models = [m] + _fallbacks(db, exclude_ids={m.id})
|
||||
reasons.append(f"explicit={m.name}")
|
||||
else:
|
||||
reasons.append("explicit rejected (disabled or tenant mismatch)")
|
||||
|
||||
# 2) 能力过滤：平台默认 + 文本备选（配额/成本分档为预留钩子，未强校验）
|
||||
if not models:
|
||||
primary = resolve_auto_model(capability, db)
|
||||
if primary and primary.is_enabled:
|
||||
models = [primary] + _fallbacks(db, exclude_ids={primary.id})
|
||||
reasons.append(f"auto={primary.name}")
|
||||
else:
|
||||
reasons.append("no primary resolved")
|
||||
|
||||
# 5) 健康 failover 重排
|
||||
models = _apply_breaker(models, claw_type)
|
||||
except Exception as _e:
|
||||
logger.warning(f"[llm_router] route error cap={capability} claw={claw_type}: {_e}")
|
||||
reasons.append(f"error:{type(_e).__name__}")
|
||||
|
||||
primary = models[0] if models else None
|
||||
matched = None
|
||||
if legacy_primary is not None:
|
||||
matched = bool(primary and primary.name == legacy_primary)
|
||||
|
||||
elapsed_ms = (time.perf_counter() - _t0) * 1000
|
||||
_emit_metrics(capability, claw_type, primary, shadow, elapsed_ms)
|
||||
if matched is not None:
|
||||
_emit_match(claw_type, matched)
|
||||
|
||||
decision = RouteDecision(
|
||||
capability=capability,
|
||||
claw_type=claw_type,
|
||||
tenant_id=tenant_id,
|
||||
models=models,
|
||||
primary=primary,
|
||||
reason="; ".join(reasons) or "default",
|
||||
shadow=shadow,
|
||||
matched_legacy=matched,
|
||||
)
|
||||
logger.info(
|
||||
f"[llm_router] claw={claw_type} cap={capability} primary="
|
||||
f"{primary.name if primary else None} shadow={shadow} matched={matched} "
|
||||
f"reason={decision.reason}"
|
||||
)
|
||||
return decision
|
||||
|
||||
|
||||
# 预留钩子（#875 后续迭代启用）：租户配额 / 成本分档
|
||||
# def _check_tenant_quota(tenant_id, capability) -> tuple[bool, str]: ...
|
||||
# def _cost_tier_models(tenant_id, capability) -> list: ...
|
||||
Reference in new issue
Block a user