341 lines
12 KiB
Python
341 lines
12 KiB
Python
"""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: ...
|