#!/usr/bin/env python3
"""model_selector.py — 模型路由与成本估算决策器（无需 API key，本地可跑）

演示分层路由：按任务类型选模型，估算单次调用成本与限速余量。
价格/限速为「截至 2026-08-30」的官方值，易变，使用前重验：
https://platform.claude.com/docs/en/models/overview
https://platform.claude.com/docs/en/api/rate-limits

运行: python3 model_selector.py
"""

# 基础价（USD / MTok，input, output）与上下文（截至 2026-08-30）
MODELS = {
    "claude-fable-5":    {"in": 10.0, "out": 50.0, "ctx": 1_000_000, "latency": "slow"},
    "claude-opus-5":     {"in": 5.0,  "out": 25.0, "ctx": 1_000_000, "latency": "moderate"},
    "claude-sonnet-5":   {"in": 2.0,  "out": 10.0, "ctx": 1_000_000, "latency": "fast"},
    "claude-haiku-4-5":  {"in": 1.0,  "out": 5.0,  "ctx": 200_000,   "latency": "fastest"},
}

# Start tier 限速（截至 2026-08-30）：RPM / ITPM / OTPM
RATE_LIMITS = {
    "claude-fable-5":    (1_000, 500_000, 100_000),
    "claude-opus-5":     (1_000, 2_000_000, 400_000),
    "claude-sonnet-5":   (1_000, 2_000_000, 400_000),
    "claude-haiku-4-5":  (1_000, 2_000_000, 400_000),
}


def route(task: str, est_input_tokens: int, need_context_over_200k: bool = False) -> str:
    """分层路由规则（示例策略，实际应由评估集校准）。"""
    if need_context_over_200k:
        pool = [m for m, s in MODELS.items() if s["ctx"] > 200_000]
    else:
        pool = list(MODELS)
    if task in ("classify", "extract", "route", "summarize_short"):
        return "claude-haiku-4-5"
    if task in ("chat", "rag_qa", "code_review"):
        return "claude-sonnet-5"
    if task in ("complex_analysis", "agentic_planning", "enterprise_critical"):
        return "claude-opus-5"
    if task == "long_running_agent":
        return "claude-fable-5"
    return "claude-sonnet-5"  # 默认


def cost(model: str, in_tok: int, out_tok: int, *, cache_hit_ratio: float = 0.0,
         batch: bool = False) -> float:
    """估算单次成本（USD）。cache hit 按 0.1x 计费；batch 半价。"""
    s = MODELS[model]
    uncached_in = in_tok * (1 - cache_hit_ratio)
    cached_in = in_tok * cache_hit_ratio
    c = uncached_in / 1e6 * s["in"] + cached_in / 1e6 * s["in"] * 0.1
    c += out_tok / 1e6 * s["out"]
    if batch:
        c *= 0.5
    return c


def throughput_check(model: str, rpm_needed: int, itpm_needed: int) -> bool:
    rpm, itpm, _ = RATE_LIMITS[model]
    return rpm_needed <= rpm and itpm_needed <= itpm


if __name__ == "__main__":
    workloads = [
        ("classify", 2_000, 100, False),
        ("rag_qa", 30_000, 1_500, False),
        ("complex_analysis", 120_000, 8_000, True),   # 需要 >200K 上下文
        ("long_running_agent", 400_000, 60_000, True),
    ]
    total = 0.0
    for task, tin, tout, big_ctx in workloads:
        m = route(task, tin, big_ctx)
        c_nocache = cost(m, tin, tout)
        c_cache = cost(m, tin, tout, cache_hit_ratio=0.8)   # system+历史前缀 80% 命中
        total += c_cache
        print(f"{task:20s} -> {m:18s} 无缓存 ${c_nocache:.4f} | 80%缓存 ${c_cache:.4f} | 批量 ${c_cache/2:.4f}")

    print(f"\n四类工作负载合计（80% 缓存）: ${total:.4f}/轮")
    all_fable = sum(cost("claude-fable-5", t, o) for _, t, o, _ in workloads)
    print(f"全部走 Fable 5 无缓存: ${all_fable:.4f}/轮  =>  分层路由节省 {100*(1-total/all_fable):.0f}%")
    print("（若把 long_running_agent 也降到 Sonnet 5，还能再省 — 但要先用评估集验证质量可接受）")

    assert throughput_check("claude-sonnet-5", 600, 1_500_000)
    assert not throughput_check("claude-fable-5", 600, 800_000)  # Fable ITPM 只有 500K
    print("\n限速校验: Sonnet5 600RPM/1.5M ITPM 通过; Fable5 800K ITPM 超限被拒（500K 上限）")