大模型应用成本管控：基于 Token Plan 的多模型路由网关设计实践

import hashlib
from typing import Literal
class LLMRouter:
    def __init__(self):
        # 定义模型优先级与预算池
        self.models = {
            "light": {"model_id": "qwen-turbo", "monthly_limit": 10_000_000},  # 轻量任务
            "heavy": {"model_id": "qwen-max", "monthly_limit": 2_000_000},    # 复杂任务
        }
  
    def classify_task(self, prompt: str) -> Literal["light", "heavy"]:
        """
        简易分类器：基于关键词+长度判断
        生产环境可替换为小型 Bert 分类模型
        """
        heavy_keywords = ["代码", "推理", "数学", "分析", "debug"]
        if any(k in prompt for k in heavy_keywords) or len(prompt) > 800:
            return "heavy"
        return "light"
  
    def route(self, prompt: str, user_id: str) -> dict:
        task_type = self.classify_task(prompt)
        target = self.models[task_type]
      
        # 检查该模型本月剩余额度
        quota_key = f"llm:quota:{user_id}:{task_type}"
        remaining = redis.decr(quota_key, 0)  # 伪代码示意
      
        if remaining <= 0:
            # 降级策略：旗舰额度用完自动降级到轻量模型
            target = self.models["light"]
          
        return {
            "model": target["model_id"],
            "strategy": "direct" if task_type == target else "fallback",
            "prompt": prompt
        }