146 lines
5.8 KiB
Python
146 lines
5.8 KiB
Python
# 只用标准库
|
||
"""多模型路由：按任务选模型、超时、限流重试、熔断、备用模型降级、成本统计。
|
||
|
||
提供方用假实现模拟（延迟、限流、故障），接口形状对应 OpenAI 兼容的 chat 接口；
|
||
真实项目中 call() 内部换成 openai SDK（base_url 指向各提供方或自部署的 vLLM）。
|
||
运行：python router.py
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import random
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
|
||
|
||
class RateLimited(Exception):
|
||
pass
|
||
|
||
|
||
class ProviderDown(Exception):
|
||
pass
|
||
|
||
|
||
@dataclass
|
||
class Provider:
|
||
name: str
|
||
model: str
|
||
price_in: float # 元 / 千 token（示例数字，按实际价格配置）
|
||
price_out: float
|
||
latency: float # 模拟延迟（秒）
|
||
fail_mode: str = "ok" # ok / rate_limit_once / down / slow
|
||
calls: int = 0
|
||
|
||
async def call(self, messages: list[dict], timeout: float) -> dict:
|
||
self.calls += 1
|
||
if self.fail_mode == "down":
|
||
raise ProviderDown(f"{self.name} 503")
|
||
if self.fail_mode == "rate_limit_once" and self.calls == 1:
|
||
raise RateLimited(f"{self.name} 429")
|
||
delay = self.latency * (20 if self.fail_mode == "slow" else 1)
|
||
await asyncio.wait_for(asyncio.sleep(delay), timeout)
|
||
text = f"[{self.model}] 已处理：{messages[-1]['content'][:10]}"
|
||
return {"text": text, "tokens_in": 200, "tokens_out": 50}
|
||
|
||
|
||
@dataclass
|
||
class Breaker:
|
||
"""连续失败 N 次后打开，冷却时间内直接跳过该提供方；冷却后放行一次试探（半开）。"""
|
||
threshold: int = 3
|
||
cooldown: float = 30.0
|
||
failures: int = 0
|
||
opened_at: float | None = None
|
||
|
||
def allow(self) -> bool:
|
||
if self.opened_at is None:
|
||
return True
|
||
return time.monotonic() - self.opened_at >= self.cooldown
|
||
|
||
def record(self, ok: bool) -> None:
|
||
if ok:
|
||
self.failures, self.opened_at = 0, None
|
||
else:
|
||
self.failures += 1
|
||
if self.failures >= self.threshold:
|
||
self.opened_at = time.monotonic()
|
||
|
||
|
||
@dataclass
|
||
class Router:
|
||
providers: dict[str, Provider]
|
||
routes: dict[str, list[str]] # 任务类型 -> 按优先级排列的提供方
|
||
timeout: float = 1.0
|
||
max_retries: int = 2
|
||
breakers: dict[str, Breaker] = field(default_factory=dict)
|
||
ledger: list[dict] = field(default_factory=list)
|
||
|
||
async def chat(self, task: str, messages: list[dict]) -> dict:
|
||
errors = []
|
||
for name in self.routes[task]:
|
||
p, br = self.providers[name], self.breakers.setdefault(name, Breaker())
|
||
if not br.allow():
|
||
errors.append(f"{name}: 熔断中，跳过")
|
||
continue
|
||
for attempt in range(self.max_retries + 1):
|
||
t0 = time.perf_counter()
|
||
try:
|
||
r = await p.call(messages, self.timeout)
|
||
except RateLimited as e: # 限流：退避后重试同一个提供方
|
||
errors.append(str(e))
|
||
await asyncio.sleep(0.01 * 2 ** attempt + random.random() * 0.01)
|
||
continue
|
||
except (ProviderDown, asyncio.TimeoutError) as e: # 故障 / 超时：换下一个提供方
|
||
br.record(False)
|
||
errors.append(f"{name}: {type(e).__name__} {e}")
|
||
break
|
||
br.record(True)
|
||
cost = r["tokens_in"] / 1000 * p.price_in + r["tokens_out"] / 1000 * p.price_out
|
||
rec = {"task": task, "provider": name, "model": p.model, "ms": round((time.perf_counter() - t0) * 1000),
|
||
"cost": round(cost, 5), "fallback": name != self.routes[task][0], "errors": errors}
|
||
self.ledger.append(rec)
|
||
return {**r, **rec}
|
||
raise RuntimeError("所有模型都不可用：" + "；".join(errors))
|
||
|
||
|
||
def build(fail_main: str = "ok") -> Router:
|
||
providers = {
|
||
"main-large": Provider("main-large", "large-v1", 0.004, 0.016, 0.02, fail_main),
|
||
"main-small": Provider("main-small", "small-v1", 0.0005, 0.002, 0.01),
|
||
"backup-large": Provider("backup-large", "other-large", 0.006, 0.02, 0.03),
|
||
"self-hosted": Provider("self-hosted", "qwen-vllm", 0.0, 0.0, 0.02),
|
||
}
|
||
routes = {
|
||
"classify": ["main-small", "self-hosted"], # 简单任务：小模型
|
||
"answer": ["main-large", "backup-large", "self-hosted"], # 复杂任务：大模型 + 备用
|
||
"sensitive": ["self-hosted"], # 敏感数据只走自部署模型
|
||
}
|
||
return Router(providers, routes, timeout=0.2)
|
||
|
||
|
||
async def main() -> None:
|
||
r = build()
|
||
print(await r.chat("classify", [{"role": "user", "content": "我要退款"}]))
|
||
print(await r.chat("sensitive", [{"role": "user", "content": "员工身份证号查询"}]))
|
||
|
||
r = build(fail_main="rate_limit_once")
|
||
out = await r.chat("answer", [{"role": "user", "content": "解释退款政策"}])
|
||
print("限流后重试：", out["provider"], out["errors"])
|
||
assert out["provider"] == "main-large" and not out["fallback"]
|
||
|
||
r = build(fail_main="slow")
|
||
out = await r.chat("answer", [{"role": "user", "content": "解释退款政策"}])
|
||
print("主模型超时后降级：", out["provider"], out["errors"])
|
||
assert out["provider"] == "backup-large" and out["fallback"]
|
||
|
||
r = build(fail_main="down")
|
||
for _ in range(4):
|
||
await r.chat("answer", [{"role": "user", "content": "q"}])
|
||
print("熔断后主模型实际调用次数：", r.providers["main-large"].calls, "（第 4 次直接跳过）")
|
||
assert r.providers["main-large"].calls == 3
|
||
total = sum(x["cost"] for x in r.ledger)
|
||
print(f"成本台账：{len(r.ledger)} 次调用，共 {total:.4f} 元")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|