model_configs.py

from __future__ import annotations

"""
Model definitions for multi-model DSPy/OpenRouter benchmarking.

Each model is identified by its **openrouter_id** (e.g. openai/gpt-oss-120b).
Prices are in USD per 1M tokens.
"""

from dataclasses import dataclass
from typing import Optional, Sequence


@dataclass(frozen=True)
class ModelConfig:
    """
    Static metadata for an LLM model used in benchmarking.

    - openrouter_id: OpenRouter model slug (the single identifier for API and CLI).
    - display_name: human-readable label for logs / tables.
    - context_window_tokens: advertised maximum context window (tokens).
    - input_cost_per_million: price for 1M input tokens (USD).
    - output_cost_per_million: price for 1M output tokens (USD).
    - notes: optional description.
    """

    openrouter_id: str
    display_name: str
    context_window_tokens: Optional[int]
    input_cost_per_million: float
    output_cost_per_million: float
    notes: Optional[str] = None


MODELS: list[ModelConfig] = [
    ModelConfig(
        openrouter_id="openai/gpt-oss-120b",
        display_name="OpenAI: gpt-oss-120b",
        context_window_tokens=131_000,
        input_cost_per_million=0.039,
        output_cost_per_million=0.19,
        notes=(
            "117B-parameter MoE; 5.1B activated per forward; MXFP4; "
            "high-reasoning, agentic, and general-purpose; native tools."
        ),
    ),
    ModelConfig(
        openrouter_id="google/gemini-3-flash-preview",
        display_name="Google: Gemini 3 Flash",
        context_window_tokens=1_050_000,
        input_cost_per_million=0.10,
        output_cost_per_million=0.40,
        notes=(
            "Latest Gemini Flash; fast multimodal with strong coding and "
            "function-calling."
        ),
    ),
    ModelConfig(
        openrouter_id="z-ai/glm-5.1",
        display_name="Z.ai: GLM 5.1",
        context_window_tokens=203_000,
        input_cost_per_million=0.95,
        output_cost_per_million=3.15,
        notes=(
            "Latest flagship GLM with enhanced programming capabilities and "
            "more stable multi-step reasoning / agent execution."
        ),
    ),
    ModelConfig(
        openrouter_id="minimax/minimax-m2.7",
        display_name="MiniMax: MiniMax M2.7",
        context_window_tokens=197_000,
        input_cost_per_million=0.30,
        output_cost_per_million=1.20,
        notes=(
            "Released Mar 18, 2026; 2.12T tokens; flagship model optimized for "
            "office workflows (Word, Excel, PPT); scores 80.2% on SWE-Bench "
            "Verified; strong planning and agentic capabilities."
        ),
    ),
    ModelConfig(
        openrouter_id="x-ai/grok-4.1-fast",
        display_name="xAI: Grok 4.1 Fast",
        context_window_tokens=2_000_000,
        input_cost_per_million=0.20,
        output_cost_per_million=0.50,
        notes=(
            "Grok 4.1 Fast; 2M-token context; $0.20/M input, $0.50/M output."
        ),
    ),
    ModelConfig(
        openrouter_id="qwen/qwen3.5-35b-a3b",
        display_name="Qwen: Qwen3.5-35B-A3B",
        context_window_tokens=262_000,
        input_cost_per_million=0.25,
        output_cost_per_million=2.00,
        notes="Vision-language MoE; linear attention; comparable to Qwen3.5-27B.",
    ),
    ModelConfig(
        openrouter_id="qwen/qwen3.5-27b",
        display_name="Qwen: Qwen3.5-27B",
        context_window_tokens=262_000,
        input_cost_per_million=0.30,
        output_cost_per_million=2.40,
        notes="Dense VLM with linear attention; comparable to Qwen3.5-122B-A10B.",
    ),
    ModelConfig(
        openrouter_id="qwen/qwen3.5-122b-a10b",
        display_name="Qwen: Qwen3.5-122B-A10B",
        context_window_tokens=262_000,
        input_cost_per_million=0.40,
        output_cost_per_million=3.20,
        notes="VLM MoE; second to Qwen3.5-397B-A17B; strong text and visual.",
    ),
    ModelConfig(
        openrouter_id="qwen/qwen3.5-9b",
        display_name="Qwen: Qwen3.5-9B",
        context_window_tokens=262_000,
        input_cost_per_million=0.10,
        output_cost_per_million=0.15,
        notes=(
            "Released Mar 10, 2026; 9B-parameter multimodal foundation model "
            "from the Qwen3.5 family; strong reasoning, coding, and visual "
            "understanding via early fusion of multimodal tokens."
        ),
    ),
    ModelConfig(
        openrouter_id="allenai/olmo-3.1-32b-instruct",
        display_name="AllenAI: Olmo 3.1 32B Instruct",
        context_window_tokens=65_536,
        input_cost_per_million=0.20,
        output_cost_per_million=0.60,
        notes="Released Jan 6, 2026; 65K context.",
    ),
    ModelConfig(
        openrouter_id="nvidia/nemotron-3-nano-30b-a3b",
        display_name="NVIDIA: Nemotron 3 Nano 30B A3B",
        context_window_tokens=262_000,
        input_cost_per_million=0.05,
        output_cost_per_million=0.20,
        notes="Small MoE; open-weights; agentic AI; high compute efficiency.",
    ),
    ModelConfig(
        openrouter_id="nvidia/nemotron-3-super-120b-a12b:free",
        display_name="NVIDIA: Nemotron 3 Super 120B A12B (Free Tier)",
        context_window_tokens=262_000,
        input_cost_per_million=0.20,
        output_cost_per_million=0.80,
        notes=(
            "Large Nemotron 3 Super 120B variant on OpenRouter; "
            "configured here with approximate list pricing so it can "
            "participate in intelligence-per-dollar benchmarking."
        ),
    ),
    ModelConfig(
        openrouter_id="mistralai/devstral-2512",
        display_name="Mistral: Devstral 2 2512",
        context_window_tokens=262_144,
        input_cost_per_million=0.40,
        output_cost_per_million=2.00,
        notes="Released Dec 9, 2025; 262K context.",
    ),
    ModelConfig(
        openrouter_id="deepseek/deepseek-v3.2",
        display_name="DeepSeek V3.2",
        context_window_tokens=131_072,
        input_cost_per_million=0.26,
        output_cost_per_million=0.38,
        notes="Released Dec 8, 2025; 131K context.",
    ),
    ModelConfig(
        openrouter_id="inception/mercury-2",
        display_name="Inception: Mercury 2",
        context_window_tokens=128_000,
        input_cost_per_million=0.25,
        output_cost_per_million=0.75,
        notes=(
            "Released Mar 4, 2026; 9.15B tokens; first reasoning diffusion LLM "
            "(dLLM); >1,000 tokens/sec; supports tunable reasoning levels, "
            "native tool use, and schema-aligned JSON output."
        ),
    ),
    ModelConfig(
        openrouter_id="anthropic/claude-sonnet-4.6",
        display_name="Anthropic: Claude Sonnet 4.6",
        context_window_tokens=200_000,
        input_cost_per_million=3.0,
        output_cost_per_million=15.0,
        notes=(
            "Gold-only: use for --gold-model when generating gold standards; "
            "excluded from default multi-model eval (too expensive for repeated runs)."
        ),
    ),
]

# Model IDs that are only used for gold generation, not in default multi-eval sweep.
GOLD_ONLY_MODEL_IDS: frozenset[str] = frozenset({"anthropic/claude-sonnet-4.6"})


MODEL_BY_ID: dict[str, ModelConfig] = {m.openrouter_id: m for m in MODELS}


def get_default_models() -> Sequence[ModelConfig]:
    """
    Return the default ordered list of models for benchmarking.

    Excludes gold-only models (e.g. Claude Sonnet) so typical multi_eval runs
    stay cheap. Use --gold-model anthropic/claude-sonnet-4.6 when generating
    golds; that model is in MODEL_BY_ID but not in this list.
    """
    return [m for m in MODELS if m.openrouter_id not in GOLD_ONLY_MODEL_IDS]