projectAIpopular/research/routerarena/local_runner.py

"""本地 Runner：不依赖 RouterArena 完整仓库，验证接入方法学。

职责：
1. 加载数据集（mock 或真实 sub_10，真实数据需从 HF 拉取）
2. 跑 adapter.get_prediction 拿到 (global_index, prompt, prediction)
3. 写入 RouterArena 协议预测文件
4. 提供 mock 推理 + Arena Score 计算，验证方法学
5. 输出路由分布报告与排行榜基线对比占位

不依赖：API key、RouterArena 仓库
仅依赖：Python 标准库 + 本项目 router_system
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

# 项目根加入 path
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))


# --- 模型价格（USD / 1M tokens）---
# 依据 RouterArena model_cost/model_cost.json 公开快照（2026-07）
# 这里给的是 input+output 平均近似；精确值在官方文件
MODEL_PRICING: Dict[str, Dict[str, float]] = {
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25},
    "gemini-2.0-flash-001": {"input": 0.075, "output": 0.30},
    "deepseek-chat": {"input": 0.14, "output": 0.28},
    "mistral-medium": {"input": 2.7, "output": 8.1},
}

# 排行榜公开基线（2026-07 快照，来源 RouterArena README L40-66）
LEADERBOARD_BASELINES: Dict[str, Dict[str, float]] = {
    "Cross-Router":        {"arena_score": 75.75, "accuracy": 78.14, "cost_per_1k": 0.40},
    "Hybrid Router":       {"arena_score": 72.08, "accuracy": 71.38, "cost_per_1k": 0.04},
    "R2-Router":           {"arena_score": 71.60, "accuracy": 71.23, "cost_per_1k": 0.06},
    "GPT-5":               {"arena_score": 64.32, "accuracy": 73.96, "cost_per_1k": 10.02},
    "MIRT-BERT":           {"arena_score": 66.89, "accuracy": 66.88, "cost_per_1k": 0.15},
    "NotDiamond":          {"arena_score": 57.29, "accuracy": 60.83, "cost_per_1k": 4.10},
    "RouteLLM":            {"arena_score": 48.07, "accuracy": 47.04, "cost_per_1k": 0.27},
    "RouterDC":            {"arena_score": 33.75, "accuracy": 32.01, "cost_per_1k": 0.07},
}


# --- Arena Score 公式（依据 RouterArena llm_evaluation/run.py L65-87）---
def compute_arena_score(
    cost_per_1k: float,
    accuracy: float,
    beta: float = 0.1,
    c_max: float = 200.0,
    c_min: float = 0.0044,
) -> float:
    if cost_per_1k is None or cost_per_1k <= 0:
        raise ValueError("cost_per_1k must be positive")
    if accuracy is None or not (0 <= accuracy <= 1):
        raise ValueError("accuracy must be in [0, 1]")
    cost_clamped = max(c_min, min(cost_per_1k, c_max))
    C = (math.log2(c_max) - math.log2(cost_clamped)) / (math.log2(c_max) - math.log2(c_min))
    return ((1 + beta) * accuracy * C) / (beta * accuracy + C)


# --- Mock 数据集生成 ---
# 模仿 RouterArena sub_10 的 9 领域结构：9 domains × ~90 queries = 810 ≈ 809
MOCK_DOMAIN_QUERIES: Dict[str, List[str]] = {
    "code": [
        "Implement quicksort in Python",
        "Write a function to reverse a linked list",
        "Debug this TypeError: undefined is not a function",
        "Optimize SQL query with index hints",
        "Implement binary search in Java",
        "用 Python 写一个快速排序函数",
        "解释这段 JavaScript 代码的 TypeError",
        "帮我优化 SQL 索引",
        "Implement merge sort",
        "Convert JSON to CSV in Python",
    ],
    "math": [
        "Solve x^2 - 5x + 6 = 0",
        "Prove the Pythagorean theorem",
        "Calculate the integral of x^2 from 0 to 1",
        "求方程 x^2+3x+2=0 的根",
        "证明勾股定理",
        "计算 3x+5=20 中 x 的值",
        "Find eigenvalues of a 2x2 matrix",
        "Differentiate sin(x) * cos(x)",
        "求 ∫ x^2 dx",
        "Compute dot product of two vectors",
    ],
    "legal": [
        "Is a non-compete clause for 2 years enforceable?",
        "How to calculate severance pay",
        "劳动合同到期不续签是否要给补偿金",
        "加班费怎么计算",
        "违约金上限 30% 合法吗",
        "What counts as wrongful termination",
        "Can I sue my employer for unpaid wages",
        "劳动合同里约定竞业限制是否有效",
        "Statute of limitations for breach of contract",
        "How does arbitration work in employment disputes",
    ],
    "medical": [
        "What foods should hypertensive patients avoid",
        "First aid for burns",
        "高血压患者日常饮食",
        "感冒发烧 38.5 度需要吃退烧药吗",
        "Side effects of common blood pressure medications",
        "When to go to ER for chest pain",
        "烫伤后怎么处理",
        "感冒初期如何缓解症状",
        "How to treat a sprained ankle",
        "What are warning signs of diabetes",
    ],
    "finance": [
        "How to calculate ROI on a fund",
        "What to do when credit card is overdue",
        "基金定投收益率怎么计算",
        "信用卡逾期怎么办",
        "房贷利率是 LPR 加多少",
        "Should I refinance my mortgage",
        "Best way to save for retirement",
        "理财产品和基金的区别",
        "How to read a stock balance sheet",
        "What is dollar-cost averaging",
    ],
    "life": [
        "Travel itinerary for Japan in 7 days",
        "Beginner muscle building plan",
        "日本旅行攻略",
        "健身增肌计划",
        "家常菜推荐",
        "How to meal prep for a week",
        "Best hiking trails near San Francisco",
        "减脂餐怎么搭配",
        "How to start running for beginners",
        "室内绿植推荐",
    ],
    "education": [
        "How to prepare for graduate English exam",
        "Effective study techniques",
        "考研英语怎么备考",
        "高效学习方法",
        "面试技巧有哪些",
        "How to write a research paper",
        "GRE quantitative prep strategy",
        "如何准备技术面试",
        "Best resources for learning Python",
        "时间管理方法",
    ],
    "general": [
        "Why is the sky blue",
        "Explain the Transformer architecture",
        "为什么天空是蓝色的",
        "介绍 Transformer 架构",
        "Write a vacation request email",
        "What is quantum entanglement",
        "请写一封请假邮件",
        "Explain CRISPR in simple terms",
        "What is blockchain",
        "简单介绍下黑洞",
    ],
    "creative": [
        "Write a haiku about autumn",
        "Suggest a name for a coffee shop",
        "Plot twist ideas for a mystery novel",
        "Write a short poem about the ocean",
        "Ideas for a 5-year-old's birthday party",
        "Story opening for a sci-fi short",
        "Suggest tagline for eco-friendly brand",
        "Lyrics for an upbeat summer song",
        "Title ideas for a romance novel",
        "Concept art description for a fantasy creature",
    ],
    # 注：RouterArena 有 9 domains，creative 是第 9 类的代表（写作/创意）
}


def build_mock_dataset() -> List[Dict[str, Any]]:
    """生成与 RouterArena sub_10 协议对齐的 mock 数据集。"""
    data: List[Dict[str, Any]] = []
    idx = 0
    for domain, queries in MOCK_DOMAIN_QUERIES.items():
        for i, q in enumerate(queries):
            data.append({
                "global index": f"mock_{idx:04d}",
                "prompt": q,
                "prompt_formatted": q,
                "domain": domain,  # 仅用于本地诊断，不暴露给路由决策
                "difficulty": ["easy", "medium", "hard"][i % 3],
            })
            idx += 1
    return data


# --- 加载器（支持 mock + 真实 sub_10 JSON）---
def load_dataset(source: str, path: Optional[str] = None) -> List[Dict[str, Any]]:
    """source ∈ {'mock', 'sub_10_file'}；path 为 None 时按 source 推断。"""
    if source == "mock":
        return build_mock_dataset()
    if source == "sub_10_file":
        if not path or not os.path.exists(path):
            raise FileNotFoundError(f"sub_10 dataset not found at {path}")
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        # RouterArena 协议字段：global index, prompt_formatted
        normalized = []
        for entry in data:
            normalized.append({
                "global index": entry.get("global index") or entry.get("global_index"),
                "prompt": entry.get("prompt_formatted") or entry.get("prompt"),
                "domain": None,  # 真实数据无 ground truth
                "difficulty": None,
            })
        return normalized
    raise ValueError(f"Unknown source: {source}")


# --- 模拟推理（不调 API）：用 L0 router 自身的 response 当 generated_answer ---
def mock_inference(router, query: str) -> Dict[str, Any]:
    """用本系统 L0 router 自身的 response 模拟目标 LLM 的输出。

    注意：这只是验证"路由层 + 协议格式"正确，不替代真实 LLM 推理。
    真实评测需要 RouterArena 的 llm_inference/run.py + 目标模型 API key。
    """
    import asyncio
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # 极少见兜底
            return {"generated_answer": "[sync-fallback]", "success": True,
                    "token_usage": {"output_tokens": 50}}
        result = loop.run_until_complete(router._router.route(query))
    except RuntimeError:
        result = asyncio.run(router._router.route(query))

    answer = result.response or ""
    # 估算 token 数（中英文 1 token ≈ 1.5 字符）
    output_tokens = max(1, int(len(answer) / 1.5))
    return {
        "generated_answer": answer,
        "success": True,
        "model_used": result.model_used,
        "token_usage": {"output_tokens": output_tokens, "input_tokens": int(len(query) / 1.5)},
    }


# --- 主流程 ---
def estimate_cost(model_name: str, token_usage: Dict[str, int]) -> float:
    """按 MODEL_PRICING 估算单条推理成本（USD）。"""
    p = MODEL_PRICING.get(model_name)
    if not p:
        return 0.0
    in_tok = token_usage.get("input_tokens", 0)
    out_tok = token_usage.get("output_tokens", 0)
    return (in_tok * p["input"] + out_tok * p["output"]) / 1_000_000


def run_local(
    source: str = "mock",
    dataset_path: Optional[str] = None,
    router_name: str = "es-expert",
    config_path: Optional[str] = None,
    do_mock_inference: bool = True,
    output_dir: str = "research/routerarena/output",
) -> Dict[str, Any]:
    """跑本地端到端流程，输出预测文件 + 诊断报告。"""
    from .adapter import ESExpertRouter

    if config_path is None:
        here = Path(__file__).resolve().parent
        config_path = str(here / "config" / f"{router_name}.json")

    router = ESExpertRouter(router_name=router_name, config_path=config_path)
    dataset = load_dataset(source, dataset_path)
    print(f"[local_runner] router models = {router.models}")
    print(f"[local_runner] dataset size  = {len(dataset)}")

    # 1) 路由决策（按 RouterArena 协议）
    predictions: List[Dict[str, Any]] = []
    diagnostics: List[Dict[str, Any]] = []
    t0 = time.perf_counter()
    for entry in dataset:
        gi = entry["global index"]
        prompt = entry["prompt"]
        selected = router.get_prediction(prompt)
        # 同时记一份诊断（科研用，不影响协议）
        diag = router.diagnostics(prompt)
        diag["global_index"] = gi
        diag["ground_truth_domain"] = entry.get("domain")  # 仅 mock 数据有
        diagnostics.append(diag)
        predictions.append({
            "global index": gi,
            "prompt": prompt,
            "prediction": selected,
            "generated_result": None,
            "cost": None,
            "accuracy": None,
            "for_optimality": False,
        })
    routing_latency_ms = (time.perf_counter() - t0) * 1000 / len(dataset)

    # 2) 模拟推理（mock generated_result）
    if do_mock_inference:
        for pred, diag in zip(predictions, diagnostics):
            gen = mock_inference(router, pred["prompt"])
            pred["generated_result"] = gen
            cost = estimate_cost(pred["prediction"], gen["token_usage"])
            pred["cost"] = cost
            # mock 准确率：仅用于方法学验证（真实评测 RouterArena 用 ground truth）
            # 这里如果 router 选对了 ground truth domain → 1.0，否则按简单启发式
            if diag.get("ground_truth_domain") and diag.get("domain"):
                # 把 L0 的 8 域映射到 RouterArena 的 9 域（creative 算 general）
                gt = diag["ground_truth_domain"]
                pred_domain = diag["domain"]
                # RouterArena creative → 本系统 general 域
                gt_mapped = "general" if gt == "creative" else gt
                pred["accuracy"] = 1.0 if pred_domain == gt_mapped else 0.0
            else:
                pred["accuracy"] = None  # 真实数据无 domain 标签，跳过

    # 3) 写预测文件（RouterArena 协议；router_name 仅取 basename 防路径穿越）
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    safe_name = Path(router_name).name
    pred_path = out_dir / f"{safe_name}.json"
    pred_path.write_text(json.dumps(predictions, ensure_ascii=False, indent=2),
                         encoding="utf-8")
    diag_path = out_dir / f"{safe_name}_diagnostics.json"
    diag_path.write_text(json.dumps(diagnostics, ensure_ascii=False, indent=2),
                         encoding="utf-8")

    # 4) 算指标
    n = len(predictions)
    acc_vals = [p["accuracy"] for p in predictions if p["accuracy"] is not None]
    cost_vals = [p["cost"] for p in predictions if p["cost"] is not None and p["cost"] > 0]
    avg_acc = sum(acc_vals) / len(acc_vals) if acc_vals else 0.0
    total_cost = sum(cost_vals) if cost_vals else 0.0
    cost_per_1k = (total_cost / n * 1000) if n > 0 else 0.0
    try:
        arena_score = compute_arena_score(cost_per_1k, avg_acc) if cost_per_1k > 0 else None
    except ValueError:
        arena_score = None

    # 路由分布
    from collections import Counter
    routing_dist = Counter(p["prediction"] for p in predictions)
    domain_dist = Counter(d["domain"] for d in diagnostics)
    confidence_dist = {
        "min": min(d["confidence"] for d in diagnostics),
        "max": max(d["confidence"] for d in diagnostics),
        "mean": sum(d["confidence"] for d in diagnostics) / len(diagnostics),
    }

    summary = {
        "router_name": router_name,
        "n_queries": n,
        "routing_latency_ms_per_query": routing_latency_ms,
        "domain_distribution": dict(domain_dist),
        "routing_distribution": dict(routing_dist),
        "confidence": confidence_dist,
        "mock_accuracy": avg_acc,
        "total_cost_usd": total_cost,
        "cost_per_1k_usd": cost_per_1k,
        "arena_score_mock": arena_score,
        "prediction_file": str(pred_path),
        "diagnostics_file": str(diag_path),
    }
    summary_path = out_dir / f"{safe_name}_summary.json"
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2),
                            encoding="utf-8")
    return summary


def main():
    parser = argparse.ArgumentParser(description="RouterArena 本地接入烟测（无需 API）")
    parser.add_argument("--source", choices=["mock", "sub_10_file"], default="mock")
    parser.add_argument("--dataset-path", default=None)
    parser.add_argument("--router", default="es-expert")
    parser.add_argument("--config", default=None)
    parser.add_argument("--no-mock-inference", action="store_true",
                        help="只跑路由不模拟推理（用于纯路由层验证）")
    parser.add_argument("--output-dir", default="research/routerarena/output")
    args = parser.parse_args()

    summary = run_local(
        source=args.source,
        dataset_path=args.dataset_path,
        router_name=args.router,
        config_path=args.config,
        do_mock_inference=not args.no_mock_inference,
        output_dir=args.output_dir,
    )
    # 打印关键指标（确保终端 GBK 安全：写到文件再读）
    out_path = os.path.join(args.output_dir, f"{args.router}_summary.json")
    with open(out_path, "r", encoding="utf-8") as f:
        s = json.load(f)
    print("\n========== RouterArena Local Run Summary ==========")
    for k, v in s.items():
        print(f"  {k}: {v}")
    print("====================================================")


if __name__ == "__main__":
    main()