router.py

#!/usr/bin/env python3
"""
SmartRouter - 智能 LLM 路由决策模块（三层架构）

Layer 1: FastRules — 同步规则匹配 (<1ms)
Layer 2: LLM 分类器 — 异步调用便宜模型做任务分类（带缓存）
Layer 3: 模型选择 — 根据 tier 选首选模型
"""

import asyncio
import hashlib
import json
import os
import re
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import yaml


# ===== 配置加载 =====


def _resolve_env_vars(value: str) -> str:
    """替换 ${VAR} 为环境变量值"""

    def replacer(match):
        var_name = match.group(1)
        return os.environ.get(var_name, "")

    return re.sub(r"\$\{(\w+)\}", replacer, str(value))


def _resolve_config(obj: Any) -> Any:
    """递归替换配置中的环境变量"""
    if isinstance(obj, str):
        return _resolve_env_vars(obj)
    elif isinstance(obj, dict):
        return {k: _resolve_config(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_resolve_config(item) for item in obj]
    return obj


def load_config() -> dict:
    """加载 config.yaml，失败则返回最小默认配置"""
    config_path = Path(__file__).parent / "config.yaml"
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            raw = yaml.safe_load(f)
        config = _resolve_config(raw)
        # 确保关键字段存在
        config.setdefault("upstream", {})
        config["upstream"].setdefault(
            "url",
            os.environ.get("NEW_API_URL", "http://127.0.0.1:30080/v1/chat/completions"),
        )
        config["upstream"].setdefault("key", os.environ.get("NEW_API_KEY", ""))
        config.setdefault("classifier", {})
        config["classifier"].setdefault("model", "MiniMax-M2.5")
        config["classifier"].setdefault("cache_size", 500)
        config["classifier"].setdefault("cache_ttl", 3600)
        config["classifier"].setdefault("timeout", 10)
        config.setdefault(
            "tiers", {"general": {"models": ["MiniMax-M2.5"], "description": "通用"}}
        )
        config.setdefault("fast_rules", [])
        config.setdefault(
            "special_models",
            {
                "auto": "smart-auto",
                "smart-auto": "smart-auto",
                "smart-premium": "premium",
            },
        )
        config.setdefault("server", {"port": 30081, "log_routes": True})
        print(f"[SmartRouter] 配置加载成功: {config_path}")
        return config
    except Exception as e:
        print(f"[SmartRouter] 配置加载失败: {e}，使用默认配置")
        return {
            "upstream": {
                "url": os.environ.get(
                    "NEW_API_URL", "http://127.0.0.1:30080/v1/chat/completions"
                ),
                "key": os.environ.get("NEW_API_KEY", ""),
            },
            "classifier": {
                "model": "MiniMax-M2.5",
                "cache_size": 500,
                "cache_ttl": 3600,
                "timeout": 10,
                "prompt": "",
            },
            "tiers": {"general": {"models": ["MiniMax-M2.5"], "description": "通用"}},
            "fast_rules": [],
            "special_models": {
                "auto": "smart-auto",
                "smart-auto": "smart-auto",
                "smart-premium": "premium",
            },
            "server": {"port": 30081, "log_routes": True},
        }


CONFIG = load_config()


def get_all_known_models() -> list:
    """从配置收集所有已知模型名（去重，保序）"""
    seen = set()
    models = []
    for tier_config in CONFIG["tiers"].values():
        for m in tier_config.get("models", []):
            if m not in seen:
                seen.add(m)
                models.append(m)
    return models


_ALL_KNOWN_MODELS = set(get_all_known_models())


# ===== 辅助函数 =====


def _extract_text_from_content(content: Any) -> str:
    """从 string 或 list[{type, text}] 提取文本"""
    if isinstance(content, str):
        return content
    elif isinstance(content, list):
        texts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                texts.append(item.get("text", ""))
        return " ".join(texts)
    return ""


def _extract_messages_content(messages: List[Dict]) -> Tuple[str, str]:
    """提取最后一条 user 消息和 system prompt"""
    user_msg = ""
    system_msg = ""
    if not messages:
        return user_msg, system_msg
    for msg in reversed(messages):
        if msg.get("role") == "user" and not user_msg:
            user_msg = _extract_text_from_content(msg.get("content", ""))
        if msg.get("role") == "system" and not system_msg:
            system_msg = _extract_text_from_content(msg.get("content", ""))
        if user_msg and system_msg:
            break
    return user_msg, system_msg


def _log_route(layer: int, tier: str, model: str, reason: str, preview: str):
    """打印路由日志"""
    if CONFIG.get("server", {}).get("log_routes", False):
        preview_short = preview[:50].replace("\n", " ") if preview else ""
        print(
            f"[SmartRouter] Layer {layer} | {tier} → {model} | {reason} | {preview_short}"
        )


# ===== Layer 1: FastRules =====


def apply_fast_rules(
    messages: List[Dict], request_data: dict
) -> Optional[Tuple[str, str]]:
    """
    快速规则匹配，返回 (tier, reason) 或 None。
    按 config 中 fast_rules 顺序逐条匹配，命中即停。
    """
    user_msg, system_msg = _extract_messages_content(messages)
    rules = CONFIG.get("fast_rules", [])

    for rule in rules:
        condition = rule.get("condition", "")
        tier = rule.get("tier", "general")
        reason = rule.get("reason", rule.get("name", condition))

        if condition == "has_tools":
            tools = request_data.get("tools")
            if tools and isinstance(tools, list) and len(tools) > 0:
                return (tier, reason)

        elif condition == "has_code_block":
            if "```" in user_msg:
                return (tier, reason)

        elif condition == "system_contains":
            keywords = rule.get("keywords", [])
            system_lower = system_msg.lower()
            if any(kw.lower() in system_lower for kw in keywords):
                return (tier, reason)

        elif condition == "short_message":
            max_len = rule.get("max_length", 20)
            if len(user_msg) <= max_len and user_msg:
                # 排除含代码/数学符号的短消息
                if not re.search(r"[`=+\-*/{}()\[\]<>|&^%$#@!\\]", user_msg):
                    return (tier, reason)

    return None


# ===== Layer 2: LLM 分类器（带缓存） =====


class ClassifierCache:
    """简单的 TTL + LRU 缓存"""

    def __init__(self, max_size: int = 500, ttl: int = 3600):
        self._cache: Dict[str, Tuple[str, float]] = {}  # key → (tier, timestamp)
        self._max_size = max_size
        self._ttl = ttl

    def get(self, key: str) -> Optional[str]:
        """获取缓存，过期返回 None"""
        if key not in self._cache:
            return None
        tier, ts = self._cache[key]
        if self._ttl > 0 and (time.time() - ts) > self._ttl:
            del self._cache[key]
            return None
        return tier

    def put(self, key: str, tier: str):
        """写入缓存，满时淘汰最旧条目"""
        if len(self._cache) >= self._max_size and key not in self._cache:
            # 淘汰最旧的
            oldest_key = min(self._cache, key=lambda k: self._cache[k][1])
            del self._cache[oldest_key]
        self._cache[key] = (tier, time.time())

    @property
    def size(self) -> int:
        return len(self._cache)


_classifier_cache = ClassifierCache(
    max_size=CONFIG["classifier"].get("cache_size", 500),
    ttl=CONFIG["classifier"].get("cache_ttl", 3600),
)

# 已知 tier 名称集合（用于校验 LLM 返回值）
_KNOWN_TIERS = set(CONFIG["tiers"].keys())


def _cache_key(text: str) -> str:
    """用消息前 200 字符的 hash 做缓存 key"""
    prefix = text[:200]
    return hashlib.md5(prefix.encode("utf-8")).hexdigest()


async def classify_with_llm(user_message: str, http_session) -> str:
    """
    调用 LLM 做任务分类，返回 tier 名称。
    失败/超时/无法识别 → fallback 到 general。
    """
    if not user_message.strip():
        return "general"

    # 检查缓存
    key = _cache_key(user_message)
    cached = _classifier_cache.get(key)
    if cached is not None:
        return cached

    # 构造分类请求
    classifier_config = CONFIG["classifier"]
    prompt_template = classifier_config.get("prompt", "")
    if not prompt_template:
        return "general"

    # 截取前 500 字符
    truncated = user_message[:500]
    prompt = prompt_template.replace("{user_message}", truncated)

    upstream_url = CONFIG["upstream"]["url"]
    upstream_key = CONFIG["upstream"]["key"]
    classifier_model = classifier_config["model"]
    timeout_sec = classifier_config.get("timeout", 10)

    request_body = {
        "model": classifier_model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 200,
        "temperature": 0,
        "stream": False,
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {upstream_key}",
    }

    try:
        async with asyncio.timeout(timeout_sec):
            async with http_session.post(
                upstream_url, json=request_body, headers=headers
            ) as resp:
                if resp.status != 200:
                    error_text = await resp.text()
                    print(
                        f"[SmartRouter] 分类器调用失败: HTTP {resp.status} | {error_text[:200]}",
                        flush=True,
                    )
                    return "general"
                data = await resp.json()
                content = ""
                choices = data.get("choices", [])
                if choices:
                    msg = choices[0].get("message", {})
                    content = msg.get("content", "").strip().lower()
                    if not content:
                        content = msg.get("reasoning_content", "").strip().lower()
                print(
                    f"[SmartRouter] 分类器原始返回: '{content}' → 消息: '{user_message[:50]}'",
                    flush=True,
                )

    except (asyncio.TimeoutError, Exception) as e:
        print(f"[SmartRouter] 分类器异常: {e}", flush=True)
        return "general"

    # 解析 tier：匹配已知 tier 名称
    tier = _parse_tier(content)

    # 写入缓存
    _classifier_cache.put(key, tier)
    return tier


def _parse_tier(raw: str) -> str:
    """从 LLM 返回的文本中提取 tier 名称"""
    raw_stripped = raw.strip().lower()
    if raw_stripped in _KNOWN_TIERS:
        return raw_stripped

    # 尝试匹配 "任务类型：XXX" 或 "类型: XXX" 模式
    for pattern in [
        r"(?:任务类型|类型|分类|结果|属于|归类)[：:]\s*(\w+)",
        r"(?:应该是|判断为|归为|属于)\s*[「「]?(\w+)[」」]?",
    ]:
        m = re.search(pattern, raw_stripped)
        if m:
            candidate = m.group(1).strip()
            if candidate in _KNOWN_TIERS:
                return candidate

    # 从后往前找最后一个独立出现的 tier 名称（避免 "simple: ... - 不是" 误匹配）
    last_tier = None
    last_pos = -1
    for tier in _KNOWN_TIERS:
        pos = raw_stripped.rfind(tier)
        if pos > last_pos:
            after = raw_stripped[pos + len(tier) : pos + len(tier) + 5]
            if "不是" not in after and "不" not in after[:2]:
                last_tier = tier
                last_pos = pos

    if last_tier:
        return last_tier

    return "general"
    return "general"


# ===== Layer 3: 模型选择 =====


def select_model(tier: str) -> str:
    """根据 tier 返回首选模型名"""
    tier_config = CONFIG["tiers"].get(tier)
    if tier_config and tier_config.get("models"):
        return tier_config["models"][0]
    # fallback 到 general
    general = CONFIG["tiers"].get("general", {})
    if general.get("models"):
        return general["models"][0]
    # 最终 fallback
    return "MiniMax-M2.5"


# ===== 主入口 =====


async def route(
    messages: List[Dict], model: str, request_data: dict, http_session
) -> str:
    """
    主路由函数，返回实际模型名。

    流程:
    1. 特殊模型名处理（smart-premium、手动指定）
    2. Layer 1: FastRules
    3. Layer 2: LLM 分类器
    4. Layer 3: 模型选择
    """
    user_msg, _ = _extract_messages_content(messages)

    # 特殊模型: smart-premium → premium tier
    if model == "smart-premium":
        selected = select_model("premium")
        _log_route(0, "premium", selected, "手动指定 premium", user_msg)
        return selected

    # 手动指定已知模型 → 直接透传
    if model not in ("auto", "smart-auto") and model in _ALL_KNOWN_MODELS:
        _log_route(0, "manual", model, "手动指定已知模型", user_msg)
        return model

    # Layer 1: 快速规则
    fast_result = apply_fast_rules(messages, request_data)
    if fast_result is not None:
        fast_tier, fast_reason = fast_result
        selected = select_model(fast_tier)
        _log_route(1, fast_tier, selected, fast_reason, user_msg)
        return selected

    # Layer 2: LLM 分类器
    tier = await classify_with_llm(user_msg, http_session)
    selected = select_model(tier)
    cache_hit = _classifier_cache.get(_cache_key(user_msg)) is not None
    _log_route(2, tier, selected, f"LLM 分类{'(缓存)' if cache_hit else ''}", user_msg)
    return selected


async def get_route_info(
    messages: List[Dict], model: str, request_data: dict, http_session
) -> dict:
    """返回详细路由信息（调试用）"""
    user_msg, system_msg = _extract_messages_content(messages)

    # 特殊模型: smart-premium
    if model == "smart-premium":
        selected = select_model("premium")
        return {
            "model": selected,
            "requested_model": model,
            "tier": "premium",
            "confidence": 1.0,
            "reasoning": "手动指定 premium 模型",
            "prompt_preview": user_msg[:100] if user_msg else "",
            "layer": 0,
        }

    # 手动指定已知模型
    if model not in ("auto", "smart-auto") and model in _ALL_KNOWN_MODELS:
        return {
            "model": model,
            "requested_model": model,
            "tier": "manual",
            "confidence": 1.0,
            "reasoning": "手动指定已知模型",
            "prompt_preview": user_msg[:100] if user_msg else "",
            "layer": 0,
        }

    # Layer 1: 快速规则
    fast_result = apply_fast_rules(messages, request_data)
    if fast_result is not None:
        fast_tier, fast_reason = fast_result
        selected = select_model(fast_tier)
        return {
            "model": selected,
            "requested_model": model,
            "tier": fast_tier,
            "confidence": 0.9,
            "reasoning": fast_reason,
            "prompt_preview": user_msg[:100] if user_msg else "",
            "layer": 1,
        }

    # Layer 2: LLM 分类器
    tier = await classify_with_llm(user_msg, http_session)
    selected = select_model(tier)
    cache_key = _cache_key(user_msg)
    is_cached = _classifier_cache.get(cache_key) is not None

    return {
        "model": selected,
        "requested_model": model,
        "tier": tier,
        "confidence": 0.85 if not is_cached else 0.9,
        "reasoning": f"LLM 分类器{'(缓存命中)' if is_cached else ''}",
        "prompt_preview": user_msg[:100] if user_msg else "",
        "layer": 2,
        "cache_hit": is_cached,
        "cache_size": _classifier_cache.size,
    }


# ===== 测试 =====

if __name__ == "__main__":
    print("SmartRouter 三层路由测试")
    print("=" * 60)
    print(f"配置 tiers: {list(CONFIG['tiers'].keys())}")
    print(f"分类器模型: {CONFIG['classifier']['model']}")
    print(f"快速规则: {len(CONFIG.get('fast_rules', []))} 条")
    print(f"已知模型: {sorted(_ALL_KNOWN_MODELS)}")
    print()

    # Layer 1 测试
    print("--- Layer 1: FastRules 测试 ---")
    test_cases_l1 = [
        # (messages, request_data, expected_tier_or_none, description)
        (
            [{"role": "user", "content": "你好"}],
            {},
            "simple",
            "短消息问候",
        ),
        (
            [
                {
                    "role": "user",
                    "content": "帮我写一个排序算法\n```python\ndef sort():\n    pass\n```",
                }
            ],
            {},
            "code",
            "包含代码块",
        ),
        (
            [{"role": "user", "content": "查询天气"}],
            {"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
            "code",
            "带 tools 参数",
        ),
        (
            [
                {"role": "system", "content": "你是一个代码助手"},
                {"role": "user", "content": "帮我优化这段逻辑"},
            ],
            {},
            "code",
            "system prompt 含代码关键词",
        ),
        (
            [{"role": "user", "content": "请详细分析一下全球经济形势的变化趋势"}],
            {},
            None,
            "普通长消息 → 需要 LLM 分类",
        ),
    ]

    for messages, req_data, expected, desc in test_cases_l1:
        result = apply_fast_rules(messages, req_data)
        actual_tier = result[0] if result else None
        status = "✓" if actual_tier == expected else "✗"
        print(f"  {status} {desc}: 期望={expected}, 实际={actual_tier}")

    print()
    print("--- Layer 2 + 3 需要 async 环境，请启动服务后用 /route 端点测试 ---")
    print()
    print("测试命令:")
    print(
        '  curl -s http://localhost:30081/route -H "Content-Type: application/json" \\'
    )
    print(
        '    -d \'{"messages":[{"role":"user","content":"你好"}]}\' | python3 -m json.tool'
    )
    print()
    print(
        '  curl -s http://localhost:30081/route -H "Content-Type: application/json" \\'
    )
    print(
        '    -d \'{"messages":[{"role":"user","content":"证明勾股定理"}]}\' | python3 -m json.tool'
    )