diff --git a/judgearena/pricing.py b/judgearena/pricing.py index bd01eec4..3b465289 100644 --- a/judgearena/pricing.py +++ b/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 diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 34ff0f87..d3c4b7c1 100644 --- a/judgearena/tuning/runner.py +++ b/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: diff --git a/tests/test_tuning_pricing.py b/tests/test_tuning_pricing.py new file mode 100644 index 00000000..8600f02c --- /dev/null +++ b/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