projectAIpopular/tests/test_routerarena_adapter.py

"""RouterArena 适配器最小测试。

覆盖：
- BaseRouter 接口契约：返回模型名必须在 config.models
- 路由决策：code→gpt-4o-mini、math→claude-3-haiku、低置信度→mistral-medium
- 预测文件格式：包含 RouterArena 必填字段
- Arena Score 公式与官方一致
"""
from __future__ import annotations

import json
import math
import os
import sys
import tempfile
from pathlib import Path

# 让 tests/ 目录能找到项目根
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from research.routerarena.adapter import (  # noqa: E402
    DOMAIN_TO_MODEL_SLOT,
    ESCALATION_MODEL_SLOT,
    ESExpertRouter,
    LOW_CONFIDENCE_THRESHOLD,
)
from research.routerarena.base_router import BaseRouter  # noqa: E402
from research.routerarena.local_runner import (  # noqa: E402
    MODEL_PRICING,
    build_mock_dataset,
    compute_arena_score,
)


CONFIG_PATH = str(
    Path(__file__).resolve().parent.parent
    / "research"
    / "routerarena"
    / "config"
    / "es-expert.json"
)


def test_base_router_loads_config():
    """BaseRouter 应能加载 config 并提取 models 列表。"""
    r = ESExpertRouter(router_name="es-expert", config_path=CONFIG_PATH)
    assert "gpt-4o-mini" in r.models
    assert "claude-3-haiku-20240307" in r.models
    assert "mistral-medium" in r.models
    assert len(r.models) == 5


def test_prediction_in_models():
    """所有 get_prediction 返回值必须在 config.models 中。"""
    r = ESExpertRouter(router_name="es-expert", config_path=CONFIG_PATH)
    samples = [
        "用 Python 写一个快速排序函数",
        "Solve x^2 = 4",
        "Is a non-compete clause enforceable?",
        "高血压患者日常饮食",
        "How to calculate ROI on a fund",
        "Why is the sky blue",
    ]
    for q in samples:
        m = r.get_prediction(q)
        assert m in r.models, f"{q} -> {m} not in {r.models}"


def test_domain_to_slot_mapping_evidence():
    """映射表与 00_integration_plan.md §3.2 严格一致。"""
    assert DOMAIN_TO_MODEL_SLOT["code"] == "gpt-4o-mini"
    assert DOMAIN_TO_MODEL_SLOT["math"] == "claude-3-haiku-20240307"
    assert DOMAIN_TO_MODEL_SLOT["legal"] == "claude-3-haiku-20240307"
    assert DOMAIN_TO_MODEL_SLOT["medical"] == "claude-3-haiku-20240307"
    assert DOMAIN_TO_MODEL_SLOT["finance"] == "claude-3-haiku-20240307"
    assert DOMAIN_TO_MODEL_SLOT["life"] == "gemini-2.0-flash-001"
    assert DOMAIN_TO_MODEL_SLOT["education"] == "deepseek-chat"
    assert DOMAIN_TO_MODEL_SLOT["general"] == "gemini-2.0-flash-001"


def test_escalation_logic():
    """_decide 在低置信度或低质量时升级到 ESCALATION_MODEL_SLOT。"""
    r = ESExpertRouter(router_name="es-expert", config_path=CONFIG_PATH)
    # 高置信度 + 高质量 → 主映射
    assert r._decide("code", 0.95, 0.97) == "gpt-4o-mini"
    # 低置信度 → 升级
    assert r._decide("code", 0.50, 0.97) == ESCALATION_MODEL_SLOT
    # 低质量 → 升级
    assert r._decide("code", 0.95, 0.50) == ESCALATION_MODEL_SLOT
    # 未知 domain → general
    assert r._decide("unknown_xyz", 0.95, 0.97) == "gemini-2.0-flash-001"


def test_compute_arena_score_matches_formula():
    """与 RouterArena run.py L65-87 公式逐项核对（0-1 标度）。"""
    # 已知：Hybrid Router $0.04/1K, 71.38% acc, leaderboard arena=72.08
    # 官方公式 raw 范围是 [0, 1]，leaderboard 显示 ×100 标度
    s = compute_arena_score(0.04, 0.7138, beta=0.1, c_max=200, c_min=0.0044)
    # raw 期望：~0.7208（×100 = leaderboard 的 72.08）
    assert 0.715 < s < 0.725, f"expected ~0.7208, got {s:.4f}"
    # 同时验证 leaderboard 标度对齐
    assert 71.5 < s * 100 < 72.5, f"×100 期望 ~72.08, got {s*100:.4f}"


def test_compute_arena_score_clamping():
    """成本超 c_max/c_min 应被 clamp。"""
    s1 = compute_arena_score(0.005, 0.7)
    s2 = compute_arena_score(0.01, 0.7)
    # 越便宜分越高（在范围内）
    assert s1 > s2 > 0
    # cost = c_max 时 C=0, S=0（公式定义）
    s_max = compute_arena_score(200, 0.7)
    assert s_max == 0.0
    # cost 被 clamp 到 c_min
    s_under = compute_arena_score(0.001, 0.7)
    s_at_min = compute_arena_score(0.0044, 0.7)
    assert abs(s_under - s_at_min) < 1e-9


def test_mock_dataset_protocol():
    """mock 数据集符合 RouterArena 协议字段。"""
    data = build_mock_dataset()
    assert len(data) >= 80  # 接近 sub_10 的 809
    for entry in data:
        assert "global index" in entry
        assert "prompt" in entry
        assert "prompt_formatted" in entry
        assert "domain" in entry  # 仅本地诊断字段
        # prompt 非空
        assert len(entry["prompt"]) > 0


def test_prediction_file_schema():
    """预测文件应包含 RouterArena 必填字段。"""
    from research.routerarena.local_runner import run_local

    with tempfile.TemporaryDirectory() as tmp:
        summary = run_local(
            source="mock",
            router_name="es-expert",
            config_path=CONFIG_PATH,
            do_mock_inference=True,
            output_dir=tmp,
        )
        # summary 含关键字段
        assert summary["n_queries"] >= 80
        assert "routing_distribution" in summary
        assert "mock_accuracy" in summary
        assert "arena_score_mock" in summary
        # 预测文件存在 + 格式合规
        with open(summary["prediction_file"], "r", encoding="utf-8") as f:
            preds = json.load(f)
        for p in preds:
            assert "global index" in p
            assert "prompt" in p
            assert "prediction" in p
            assert "for_optimality" in p
            assert p["prediction"] in MODEL_PRICING or p["prediction"] is None