WIP: Judge tuning (5/6): Resolve OpenRouter reference prices

judgearena/pricing.py

@@ -1,9 +1,15 @@
 """Shared reference token pricing and request cost calculation."""
 
+import json
 from dataclasses import dataclass
+from pathlib import Path
+from urllib.error import URLError
+from urllib.request import urlopen
 
 from judgearena.usage import RequestUsage
 
+_OPENROUTER_URL = "https://openrouter.ai/api/v1/models"
+
 
 @dataclass(frozen=True)
 class TokenPrice:
@@ -30,21 +36,89 @@ def reference_cost(usage: RequestUsage, price: TokenPrice) -> float:
     ) / 1_000_000
 
 
-def resolve_prices(models, overrides, *, require_cost=True) -> dict[str, TokenPrice]:
-    """Resolve explicit USD-per-million-token rates before searching."""
+def _fetch_openrouter(catalog_cache: Path | None) -> list[dict]:
+    try:
+        with urlopen(_OPENROUTER_URL, timeout=10) as response:
+            records = json.load(response)["data"]
+        if catalog_cache is not None:
+            catalog_cache.parent.mkdir(parents=True, exist_ok=True)
+            catalog_cache.write_text(json.dumps(records), encoding="utf-8")
+        return records
+    except (URLError, TimeoutError, OSError) as error:
+        if catalog_cache is not None and catalog_cache.exists():
+            cached = json.loads(catalog_cache.read_text(encoding="utf-8"))
+            if isinstance(cached, list):
+                return cached
+            raise ValueError(
+                "Cached OpenRouter pricing catalog uses an unsupported format; "
+                "refresh the catalog or set price_per_million_tokens overrides"
+            ) from error
+        raise ValueError(
+            "Could not fetch OpenRouter pricing and no cached catalog is available; "
+            "fetch the catalog on a login node or set "
+            "price_per_million_tokens overrides"
+        ) from error
+
+
+def _find_price(model: str, records: list[dict]) -> TokenPrice | None:
+    requested_id = model.split("/", 1)[-1].lower()
+    suffix = "/".join(requested_id.split("/")[-2:])
+    record = next(
+        (
+            item
+            for item in records
+            if str(item.get("hugging_face_id") or "").lower() == requested_id
+        ),
+        None,
+    )
+    if record is None:
+        record = next(
+            (item for item in records if str(item.get("id") or "").lower() == suffix),
+            None,
+        )
+    if record is None:
+        return None
+    rates = record.get("pricing", {})
+    if "prompt" not in rates or "completion" not in rates:
+        return None
+    return TokenPrice(
+        float(rates["prompt"]) * 1_000_000,
+        float(rates["completion"]) * 1_000_000,
+        "openrouter_reference",
+    )
+
+
+def resolve_prices(
+    models,
+    overrides,
+    *,
+    catalog_cache: Path | None = None,
+    require_cost=True,
+) -> dict[str, TokenPrice]:
+    """Resolve explicit or OpenRouter reference USD-per-million-token rates."""
     if not require_cost:
         return {}
     prices = {}
+    unresolved = []
     for model in sorted(models):
-        if model not in overrides:
-            continue
-        rate = overrides[model]
-        if isinstance(rate, (int, float)):
-            rate = {"input": rate, "output": rate}
-        prices[model] = TokenPrice(rate["input"], rate["output"], "user_override")
+        rate = overrides.get(model)
+        if rate is not None:
+            if isinstance(rate, (int, float)):
+                rate = {"input": rate, "output": rate}
+            prices[model] = TokenPrice(
+                float(rate["input"]), float(rate["output"]), "user_override"
+            )
+        else:
+            unresolved.append(model)
+    records = _fetch_openrouter(catalog_cache) if unresolved else []
+    for model in unresolved:
+        price = _find_price(model, records)
+        if price is not None:
+            prices[model] = price
     missing = sorted(set(models) - prices.keys())
     if missing:
         raise ValueError(
-            f"No token price for {missing}; set tune_judge.price_per_million_tokens"
+            f"No token price for {missing}; set price_per_million_tokens "
+            f"overrides for: {missing}"
         )
     return prices

judgearena/tuning/runner.py

@@ -276,6 +276,8 @@ def run_tune_judge(
             prices = resolve_prices(
                 models,
                 tuning.price_per_million_tokens,
+                catalog_cache=Path(cfg.run.store_root or cfg.run.result_folder)
+                / "openrouter_pricing.json",
                 require_cost=True,
             )
             if not tuning.search_only:

tests/test_tuning_pricing.py

@@ -0,0 +1,66 @@
+import json
+from contextlib import nullcontext
+from io import BytesIO
+
+import pytest
+
+from judgearena import pricing
+from judgearena.pricing import TokenPrice, reference_cost, resolve_prices
+from judgearena.usage import RequestUsage
+
+
+def test_override_and_agreement_only_skip_catalog(monkeypatch):
+    monkeypatch.setattr(pricing, "urlopen", lambda *a, **k: pytest.fail("fetched"))
+    assert resolve_prices(["VLLM/local"], {"VLLM/local": 1.2}) == {
+        "VLLM/local": TokenPrice(1.2, 1.2, "user_override")
+    }
+    assert resolve_prices(["VLLM/local"], {}, require_cost=False) == {}
+
+
+@pytest.mark.parametrize(
+    "provider,hf_match", [("VLLM", True), ("OpenRouter", True), ("VLLM", False)]
+)
+def test_reference_matching_and_catalog_snapshot(
+    provider, hf_match, monkeypatch, tmp_path
+):
+    records = [dict(id="org/model", pricing=dict(prompt="1e-6", completion="2e-6"))]
+    if hf_match:
+        records.append(
+            dict(
+                id="other/model",
+                hugging_face_id="ORG/Model",
+                pricing=dict(prompt="3e-6", completion="4e-6"),
+            )
+        )
+    response = json.dumps({"data": records}).encode()
+    monkeypatch.setattr(
+        pricing, "urlopen", lambda *a, **k: nullcontext(BytesIO(response))
+    )
+    model, cache = f"{provider}/org/model", tmp_path / "prices.json"
+    expected = (
+        TokenPrice(3, 4, "openrouter_reference")
+        if hf_match
+        else TokenPrice(1, 2, "openrouter_reference")
+    )
+    assert resolve_prices([model], {}, catalog_cache=cache)[model] == expected
+    assert json.loads(cache.read_text()) == records
+    with pytest.raises(ValueError, match="overrides"):
+        resolve_prices(["VLLM/unknown"], {}, catalog_cache=cache)
+
+
+@pytest.mark.parametrize(
+    "tokens,expected", [((0, 3), 6), ((2, 1), 4), ((None, 3), None)]
+)
+def test_native_counts_and_no_double_counted_reasoning(tokens, expected):
+    usage = RequestUsage(
+        stage="judging",
+        input_tokens=tokens[0],
+        output_tokens=tokens[1],
+        reasoning_tokens=100,
+    )
+    price = TokenPrice(1_000_000, 2_000_000, "ref")
+    if expected is None:
+        with pytest.raises(ValueError, match="missing native"):
+            reference_cost(usage, price)
+    else:
+        assert reference_cost(usage, price) == expected