From a99202f92e7ff0b05cf2b98256464e564613e2ac Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:54:07 +0200 Subject: [PATCH 1/4] Preserve native request usage through inference caches and judge annotations --- judgearena/benchmarks/meta_eval/annotate.py | 3 + judgearena/cache/inference.py | 9 +- judgearena/cache/sqlite.py | 52 +++++++-- judgearena/evaluate.py | 11 +- judgearena/models.py | 10 +- judgearena/usage.py | 15 ++- tests/test_cache_sqlite.py | 117 ++++++++++++++++++++ tests/test_inference_cache.py | 57 ++++++++-- tests/test_meta_eval_annotations.py | 55 ++++++++- tests/test_meta_eval_foundation.py | 3 +- tests/test_usage.py | 13 +++ 11 files changed, 317 insertions(+), 28 deletions(-) diff --git a/judgearena/benchmarks/meta_eval/annotate.py b/judgearena/benchmarks/meta_eval/annotate.py index ca5a067a..f7bdefed 100644 --- a/judgearena/benchmarks/meta_eval/annotate.py +++ b/judgearena/benchmarks/meta_eval/annotate.py @@ -10,6 +10,7 @@ from judgearena.arenas_utils import extract_turn_text from judgearena.evaluate import JudgeAnnotation, judge_and_parse_prefs +from judgearena.usage import request_usage_to_json if TYPE_CHECKING: from judgearena.config import RunConfig @@ -99,6 +100,8 @@ def _annotation_frame( "judge_top_logprobs_json": _serialize_mapping( annotation.judge_top_logprobs ), + "usage_json": request_usage_to_json(annotation.usage), + "error": annotation.error, "parsed_label": None if parsed is None else parsed.label, "parsed_scores_json": ( None if parsed is None else _serialize_mapping(parsed.scores) diff --git a/judgearena/cache/inference.py b/judgearena/cache/inference.py index 63b79d07..17b497ad 100644 --- a/judgearena/cache/inference.py +++ b/judgearena/cache/inference.py @@ -24,6 +24,7 @@ ) from judgearena.constants import VLLM_DEFAULT_TEMPERATURE, VLLM_DEFAULT_TOP_P from judgearena.inference import InferenceResult +from judgearena.usage import request_usage_from_json _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} VLLM_EXECUTION_ONLY_KWARGS = { @@ -241,10 +242,14 @@ def make_row( "benchmark": self.task, "instruction_id": metadata["instruction_id"], "model": model.model_spec, + "usage_json": output.usage, } def cached_result(self, row: pd.Series) -> InferenceResult: - return InferenceResult(text=str(row["completion"])) + return InferenceResult( + text=str(row["completion"]), + usage=request_usage_from_json(row.get("usage_json")), + ) class JudgementInferenceCache(InferenceCache[JudgementCacheRowMetadata]): @@ -272,6 +277,7 @@ def make_row( "judge": model.model_spec, "top_logprobs": output.first_token_top_logprobs, "orientation": metadata.get("orientation"), + "usage_json": output.usage, } def cached_result(self, row: pd.Series) -> InferenceResult: @@ -281,4 +287,5 @@ def cached_result(self, row: pd.Series) -> InferenceResult: first_token_top_logprobs=( json.loads(top_logprobs) if pd.notna(top_logprobs) else None ), + usage=request_usage_from_json(row.get("usage_json")), ) diff --git a/judgearena/cache/sqlite.py b/judgearena/cache/sqlite.py index 2d745d21..b6c789e7 100644 --- a/judgearena/cache/sqlite.py +++ b/judgearena/cache/sqlite.py @@ -14,6 +14,8 @@ import pandas as pd +from judgearena.usage import request_usage_to_json + COMPLETION_DB_NAME = "completions.db" JUDGEMENT_DB_NAME = "judgements.db" DESCRIPTOR_FILENAME = "metadata.json" @@ -93,6 +95,7 @@ def read_descriptor(folder: Path) -> dict[str, Any]: class _SQLiteCache: table: str schema: str + output_column: str def __init__(self, db_path: Path | str) -> None: self.db_path = Path(db_path) @@ -102,8 +105,17 @@ def _connect(self) -> sqlite3.Connection: if self._conn is None: self.db_path.parent.mkdir(parents=True, exist_ok=True) self._conn = sqlite3.connect(self.db_path) - self._conn.execute(self.schema) - self._conn.commit() + with self._conn: + self._conn.execute("BEGIN IMMEDIATE") + self._conn.execute(self.schema) + columns = { + row[1] + for row in self._conn.execute(f"PRAGMA table_info({self.table})") + } + if "usage_json" not in columns: + self._conn.execute( + f"ALTER TABLE {self.table} ADD COLUMN usage_json TEXT" + ) return self._conn def _query( @@ -145,17 +157,35 @@ def merge_from(self, other_db: Path) -> int: columns = [row[1] for row in conn.execute(f"PRAGMA table_info({self.table})")] column_list = ", ".join(columns) updates = ", ".join( - f"{column} = excluded.{column}" + ( + f"{column} = CASE " + f"WHEN excluded.{column} IS NULL " + f"AND excluded.{self.output_column} = " + f"{self.table}.{self.output_column} " + f"THEN {self.table}.{column} ELSE excluded.{column} END" + if column == "usage_json" + else f"{column} = excluded.{column}" + ) for column in columns if column != "input_hash" ) conn.execute("ATTACH DATABASE ? AS incoming", (str(other_db),)) try: + incoming_columns = { + row[1] + for row in conn.execute(f"PRAGMA incoming.table_info({self.table})") + } + projection = column_list + if "usage_json" not in incoming_columns: + projection = ", ".join( + "NULL AS usage_json" if column == "usage_json" else column + for column in columns + ) with conn: conn.execute( f""" INSERT INTO {self.table} ({column_list}) - SELECT {column_list} FROM incoming.{self.table} WHERE true + SELECT {projection} FROM incoming.{self.table} WHERE true ON CONFLICT(input_hash) DO UPDATE SET {updates} WHERE excluded.pushed_at > {self.table}.pushed_at """ @@ -180,6 +210,7 @@ class CompletionCache(_SQLiteCache): """Completion rows keyed by the exact rendered model input.""" table = "completions" + output_column = "completion" schema = """ CREATE TABLE IF NOT EXISTS completions ( input_hash TEXT PRIMARY KEY, @@ -190,7 +221,8 @@ class CompletionCache(_SQLiteCache): model TEXT NOT NULL, pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, - run_id TEXT NOT NULL + run_id TEXT NOT NULL, + usage_json TEXT ) """ @@ -214,12 +246,13 @@ def save( pushed_by, now, resolved_run_id, + request_usage_to_json(row.get("usage_json")), ) for _, row in rows.iterrows() ] with self._connect() as conn: conn.executemany( - "INSERT OR REPLACE INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "INSERT OR REPLACE INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", values, ) return len(values) @@ -262,6 +295,7 @@ class JudgementCache(_SQLiteCache): """Raw judge completions keyed by the exact rendered judge input.""" table = "judgements" + output_column = "judge_completion" schema = """ CREATE TABLE IF NOT EXISTS judgements ( input_hash TEXT PRIMARY KEY, @@ -277,7 +311,8 @@ class JudgementCache(_SQLiteCache): orientation TEXT, pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, - run_id TEXT NOT NULL + run_id TEXT NOT NULL, + usage_json TEXT ) """ @@ -309,13 +344,14 @@ def save( pushed_by, now, resolved_run_id, + request_usage_to_json(row.get("usage_json")), ) for _, row in rows.iterrows() ] with self._connect() as conn: conn.executemany( "INSERT OR REPLACE INTO judgements " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", values, ) return len(values) diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index 3c3ad2e2..f47fed40 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -8,7 +8,7 @@ from langchain_core.prompts import ChatPromptTemplate from judgearena.log import get_logger -from judgearena.models import InferenceResult, do_inference +from judgearena.models import do_inference from judgearena.prompts.registry import ( DEFAULT_JUDGE_PROMPT_PRESET, ResolvedJudgePrompt, @@ -17,6 +17,7 @@ from judgearena.prompts.registry import ( resolve_run_judge_prompt as _resolve_run_judge_prompt, ) +from judgearena.usage import RequestUsage from judgearena.utils import strip_thinking_tags, truncate if TYPE_CHECKING: @@ -103,6 +104,8 @@ class JudgeAnnotation: # first-token top logprobs, only collected for logprob-weighted presets judge_top_logprobs: dict[str, float] | None = None parsed: ParsedPreference | None = None + usage: RequestUsage | None = None + error: str | None = None def annotate_battles( @@ -188,12 +191,10 @@ def annotate_battles( inputs=inputs, use_tqdm=use_tqdm, return_top_logprobs=collect_top_logprobs, + return_results=True, stage="judging", cache_row_metadata=cache_row_metadata, ) - if not collect_top_logprobs: - judge_results = [InferenceResult(text=text) for text in judge_results] - annotations = [] for judge_input, judge_result, instruction, completion_A, completion_B in zip( inputs, @@ -212,6 +213,8 @@ def annotate_battles( completion_B=completion_B, prompt_preset=resolved_prompt.preset_name, judge_top_logprobs=judge_result.first_token_top_logprobs, + usage=judge_result.usage, + error=judge_result.error, ) ) return annotations diff --git a/judgearena/models.py b/judgearena/models.py index 585db28c..ff0d3a5d 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -714,6 +714,7 @@ def _run_backend_inference( inputs, use_tqdm: bool = False, return_top_logprobs: bool = False, + return_results: bool = False, *, stage: str = "unspecified", ): @@ -813,7 +814,7 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): # LangChain chat models return AIMessage objects, barebones models plain # strings, and ChatVLLM may return structured InferenceResult objects. results = _collect_inference_results(chat_model, res, stage=stage) - if return_top_logprobs: + if return_top_logprobs or return_results: return results return [result.text for result in results] @@ -824,6 +825,7 @@ def do_inference( use_tqdm: bool = False, return_top_logprobs: bool = False, *, + return_results: bool = False, stage: str = "unspecified", cache_row_metadata: list[CacheRowMetadata] | None = None, ): @@ -835,6 +837,7 @@ def do_inference( inputs, use_tqdm, return_top_logprobs, + return_results, stage=stage, ) @@ -845,6 +848,7 @@ def do_inference( inputs, use_tqdm, return_top_logprobs, + return_results, stage=stage, ) if cache_row_metadata is None or len(cache_row_metadata) != len(inputs): @@ -866,6 +870,7 @@ def do_inference( inputs, use_tqdm, return_top_logprobs, + return_results, stage=stage, ) @@ -883,6 +888,7 @@ def do_inference( inputs, use_tqdm, return_top_logprobs, + return_results, stage=stage, ) results: list[InferenceResult | None] = [ @@ -928,7 +934,7 @@ def do_inference( logger.warning("Cache close failed at %s: %s.", store.db_path, exc) resolved_results = [result for result in results if result is not None] - if return_top_logprobs: + if return_top_logprobs or return_results: return resolved_results return [result.text for result in resolved_results] diff --git a/judgearena/usage.py b/judgearena/usage.py index 23b22ae9..0a4721b6 100644 --- a/judgearena/usage.py +++ b/judgearena/usage.py @@ -2,10 +2,11 @@ from __future__ import annotations +import json from collections.abc import Iterator from contextlib import contextmanager from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import asdict, dataclass @dataclass(frozen=True) @@ -35,6 +36,18 @@ def has_token_usage(self) -> bool: ) +def request_usage_to_json(usage: RequestUsage | None) -> str | None: + if usage is None: + return None + return json.dumps(asdict(usage), sort_keys=True, separators=(",", ":")) + + +def request_usage_from_json(value: str | None) -> RequestUsage | None: + if not isinstance(value, str): + return None + return RequestUsage(**json.loads(value)) + + def _sum_optional(requests: tuple[RequestUsage, ...], field: str) -> int | None: values = [getattr(request, field) for request in requests] reported = [int(value) for value in values if value is not None] diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py index e88ae95e..b8d3d47e 100644 --- a/tests/test_cache_sqlite.py +++ b/tests/test_cache_sqlite.py @@ -1,4 +1,5 @@ import json +import sqlite3 import pandas as pd import pytest @@ -12,6 +13,7 @@ input_hash, write_descriptor, ) +from judgearena.usage import RequestUsage DESCRIPTOR = { "model": "Qwen/Qwen3-8B", @@ -181,3 +183,118 @@ def test_merge_from_updates_live_database_in_place(tmp_path): assert cache.merge_from(incoming_path) == 2 assert local_path.stat().st_ino == inode assert cache.query()["completion"].tolist() == ["incoming", "new"] + + +def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): + db_path = tmp_path / COMPLETION_DB_NAME + legacy = sqlite3.connect(db_path) + legacy.execute( + """CREATE TABLE completions ( + input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, + completion TEXT NOT NULL, benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, model TEXT NOT NULL, + pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL + )""" + ) + legacy.commit() + legacy.close() + + rows = pd.DataFrame( + [ + { + "input_text": "prompt", + "completion": "answer", + "benchmark": "arena", + "instruction_id": "1", + "model": "Dummy/model", + "usage_json": RequestUsage( + stage="generation", input_tokens=3, output_tokens=5 + ), + } + ] + ) + with CompletionCache(db_path) as cache: + cache.save(rows, pushed_by="alice") + result = cache.query().iloc[0] + + assert "usage_json" in result.index + assert json.loads(result["usage_json"]) == { + "cached_tokens": None, + "cost_usd": None, + "input_tokens": 3, + "model": None, + "output_tokens": 5, + "reasoning_tokens": None, + "stage": "generation", + "total_tokens": None, + } + + +@pytest.mark.parametrize( + ("incoming_completion", "expect_prior_usage"), + [("old answer", False), ("native answer", True)], +) +def test_merge_legacy_incoming_database_projects_null_usage( + tmp_path, incoming_completion, expect_prior_usage +): + local_path = tmp_path / "local" / COMPLETION_DB_NAME + incoming_path = tmp_path / "legacy" / COMPLETION_DB_NAME + incoming_path.parent.mkdir(parents=True) + legacy = sqlite3.connect(incoming_path) + legacy.execute( + """CREATE TABLE completions ( + input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, + completion TEXT NOT NULL, benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, model TEXT NOT NULL, + pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL + )""" + ) + legacy.execute( + "INSERT INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + input_hash("old prompt"), + "old prompt", + incoming_completion, + "arena", + "1", + "Dummy/model", + "bob", + "2030-01-01T00:00:00+00:00", + "old-run", + ), + ) + legacy.commit() + legacy.close() + + prior_usage = RequestUsage(stage="generation", input_tokens=4) + with CompletionCache(local_path) as cache: + cache.save( + pd.DataFrame( + [ + { + "input_text": "old prompt", + "completion": "native answer", + "benchmark": "arena", + "instruction_id": "1", + "model": "Dummy/model", + "usage_json": prior_usage, + } + ] + ), + pushed_by="alice", + ) + assert cache.merge_from(incoming_path) == 1 + result = cache.query().iloc[0] + + assert result["completion"] == incoming_completion + if expect_prior_usage: + assert json.loads(result["usage_json"])["input_tokens"] == 4 + else: + assert pd.isna(result["usage_json"]) + + incoming = sqlite3.connect(incoming_path) + incoming_columns = { + column[1] for column in incoming.execute("PRAGMA table_info(completions)") + } + incoming.close() + assert "usage_json" not in incoming_columns diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index a8f0ef94..d6d395ff 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -15,7 +15,7 @@ ) from judgearena.inference import InferenceResult from judgearena.models import do_inference, prepare_model -from judgearena.usage import track_usage +from judgearena.usage import RequestUsage, track_usage class EchoModel: @@ -87,6 +87,7 @@ def batch(self, inputs, **_kwargs): InferenceResult( text="m", first_token_top_logprobs={"m": -0.1, "M": -2.0}, + usage=RequestUsage(stage="judging", input_tokens=13), ) for _ in inputs ] @@ -106,16 +107,56 @@ def batch(self, inputs, **_kwargs): return_top_logprobs=True, cache_row_metadata=metadata, ) - second = do_inference( - prepare_model("Dummy/judge", cache=cache), - ["judge prompt"], - return_top_logprobs=True, - cache_row_metadata=metadata, - ) + with track_usage() as tracker: + second = do_inference( + prepare_model("Dummy/judge", cache=cache), + ["judge prompt"], + return_top_logprobs=True, + cache_row_metadata=metadata, + ) assert second[0].text == first[0].text assert second[0].first_token_top_logprobs == first[0].first_token_top_logprobs - assert second[0].usage is None + assert second[0].usage == first[0].usage + assert tracker.snapshot().requests == () + + +def test_cached_usage_is_restored_without_recording_new_spend(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + class UsageModel: + def batch(self, inputs, **_kwargs): + return [ + InferenceResult( + text="answer", + usage=RequestUsage( + stage="generation", input_tokens=7, output_tokens=11 + ), + ) + for _ in inputs + ] + + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: UsageModel()) + metadata = [{"instruction_id": "1"}] + with track_usage() as tracker: + first = do_inference( + prepare_model("Dummy/usage-model", cache=cache), + ["prompt"], + cache_row_metadata=metadata, + return_results=True, + ) + assert tracker.snapshot().requests == (first[0].usage,) + + with track_usage() as tracker: + second = do_inference( + prepare_model("Dummy/usage-model", cache=cache), + ["prompt"], + cache_row_metadata=metadata, + return_results=True, + ) + + assert second[0].usage == first[0].usage + assert tracker.snapshot().requests == () def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): diff --git a/tests/test_meta_eval_annotations.py b/tests/test_meta_eval_annotations.py index b39625be..5611205f 100644 --- a/tests/test_meta_eval_annotations.py +++ b/tests/test_meta_eval_annotations.py @@ -13,7 +13,9 @@ aggregate_battle_preferences, annotate_sample, ) +from judgearena.inference import InferenceResult from judgearena.prompts.registry import resolve_judge_prompt +from judgearena.usage import RequestUsage def _sample(): @@ -53,8 +55,18 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): ) responses = iter( [ - ["score_A: 9\nscore_B: 1", "score_A: 5\nscore_B: 6", "unparseable"], - ["score_A: 1\nscore_B: 7", "unparseable", "unparseable"], + [ + InferenceResult(text=text) + for text in ( + "score_A: 9\nscore_B: 1", + "score_A: 5\nscore_B: 6", + "unparseable", + ) + ], + [ + InferenceResult(text=text) + for text in ("score_A: 1\nscore_B: 7", "unparseable", "unparseable") + ], ] ) monkeypatch.setattr(evaluate, "do_inference", lambda **kwargs: next(responses)) @@ -67,6 +79,7 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): assert rows["orientation"].tolist() == ["direct"] * 3 + ["reversed"] * 3 rendered = rows["judge_input"] + assert pd.isna(rows.loc[2, "error"]) assert rendered[0].index("Alpha answer") < rendered[0].index("Beta answer") assert rendered[3].index("Beta answer") < rendered[3].index("Alpha answer") assert rows.loc[[0, 3], "pref"].tolist() == pytest.approx( @@ -88,7 +101,9 @@ def test_annotation_preserves_parser_label_and_details(monkeypatch): output = ( '{"ordered_models": [{"model": "M", "rank": 1}, {"model": "m", "rank": 2}]}' ) - monkeypatch.setattr(evaluate, "do_inference", lambda **kwargs: [output]) + monkeypatch.setattr( + evaluate, "do_inference", lambda **kwargs: [InferenceResult(text=output)] + ) rows = annotate_sample( _sample(), _config("fixed"), @@ -101,6 +116,40 @@ def test_annotation_preserves_parser_label_and_details(monkeypatch): assert battles["pref"].tolist() == [1.0] +def test_annotation_carries_usage_and_context_skip_error(monkeypatch): + usage = RequestUsage(stage="judging", model="judge", input_tokens=7) + + def fake_do_inference(**kwargs): + assert kwargs["return_results"] is True + assert kwargs["return_top_logprobs"] is False + return [ + InferenceResult(text="", error="context_length"), + InferenceResult(text="", usage=usage), + ] + + monkeypatch.setattr(evaluate, "do_inference", fake_do_inference) + sample = pd.concat([_sample()] * 2, ignore_index=True) + rows = annotate_sample( + sample, + _config("fixed"), + judge_chat_model=object(), + resolved_prompt=resolve_judge_prompt(preset="meta-eval-pair-score"), + ) + assert rows.loc[0, "error"] == "context_length" + assert pd.isna(rows.loc[0, "usage_json"]) + assert rows.loc[1, "error"] is None + assert json.loads(rows.loc[1, "usage_json"]) == { + "stage": "judging", + "model": "judge", + "input_tokens": 7, + "output_tokens": None, + "total_tokens": None, + "reasoning_tokens": None, + "cached_tokens": None, + "cost_usd": None, + } + + def test_aggregate_rejects_incomplete_orientation_sets(): annotations = pd.DataFrame( {"battle_id": ["q1"], "orientation": ["direct"], "pref": [0.2]} diff --git a/tests/test_meta_eval_foundation.py b/tests/test_meta_eval_foundation.py index 4233c900..4b2225af 100644 --- a/tests/test_meta_eval_foundation.py +++ b/tests/test_meta_eval_foundation.py @@ -14,6 +14,7 @@ sample_battles_per_model, select_top_models, ) +from judgearena.inference import InferenceResult from judgearena.prompts.parsing import JUDGE_PARSERS from judgearena.tasks.registry import get_packaged_task from judgearena.tasks.schema import TaskSpec @@ -100,7 +101,7 @@ def test_meta_eval_alpaca_prompt_embeds_json_safe_inputs(monkeypatch): def fake_do_inference(**kwargs): captured.extend(kwargs["inputs"]) - return ["unparsed"] + return [InferenceResult(text="unparsed")] monkeypatch.setattr(evaluate_module, "do_inference", fake_do_inference) instruction = 'Say "hi"\r\nnext \\ path {curly} café' diff --git a/tests/test_usage.py b/tests/test_usage.py index 51e7a0a7..c3726995 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -59,6 +59,19 @@ def test_structured_response_preserves_partial_fallback_usage(capsys): assert "input/output token usage unavailable" in capsys.readouterr().out +def test_return_results_exposes_native_usage_without_logprobs(): + message = AIMessage( + content="answer", + usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + ) + with track_usage() as tracker: + (result,) = do_inference(FakeModel([message]), ["prompt"], return_results=True) + assert result.text == "answer" + assert result.first_token_top_logprobs is None + assert result.usage == tracker.snapshot().requests[0] + assert result.usage.total_tokens == 12 + + def test_run_benchmark_saves_nested_usage_and_cleans_up(tmp_path, monkeypatch, capsys): monkeypatch.setattr(metadata_module, "_get_dependency_versions", lambda **_: {}) monkeypatch.setattr(metadata_module, "_get_git_hash", lambda **_: None) From 1e6c2d58bbfccebf56a4b53f456b525e089ad22c Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:28:24 +0200 Subject: [PATCH 2/4] Keep cached request usage independent of dollar prices --- judgearena/cache/inference.py | 6 ++--- judgearena/cache/sqlite.py | 23 +++++++++++++++---- tests/test_cache_sqlite.py | 18 ++++++++++----- tests/test_inference_cache.py | 43 +++++++++++++++++++++++++++++++++-- 4 files changed, 75 insertions(+), 15 deletions(-) diff --git a/judgearena/cache/inference.py b/judgearena/cache/inference.py index 17b497ad..918cdfed 100644 --- a/judgearena/cache/inference.py +++ b/judgearena/cache/inference.py @@ -19,12 +19,12 @@ CompletionCache, JudgementCache, cache_folder, + cached_usage_from_json, stable_json_dumps, write_descriptor, ) from judgearena.constants import VLLM_DEFAULT_TEMPERATURE, VLLM_DEFAULT_TOP_P from judgearena.inference import InferenceResult -from judgearena.usage import request_usage_from_json _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} VLLM_EXECUTION_ONLY_KWARGS = { @@ -248,7 +248,7 @@ def make_row( def cached_result(self, row: pd.Series) -> InferenceResult: return InferenceResult( text=str(row["completion"]), - usage=request_usage_from_json(row.get("usage_json")), + usage=cached_usage_from_json(row.get("usage_json")), ) @@ -287,5 +287,5 @@ def cached_result(self, row: pd.Series) -> InferenceResult: first_token_top_logprobs=( json.loads(top_logprobs) if pd.notna(top_logprobs) else None ), - usage=request_usage_from_json(row.get("usage_json")), + usage=cached_usage_from_json(row.get("usage_json")), ) diff --git a/judgearena/cache/sqlite.py b/judgearena/cache/sqlite.py index b6c789e7..5a8b4d7e 100644 --- a/judgearena/cache/sqlite.py +++ b/judgearena/cache/sqlite.py @@ -6,7 +6,7 @@ import json import sqlite3 import uuid -from dataclasses import astuple, dataclass, fields +from dataclasses import asdict, astuple, dataclass, fields, replace from datetime import UTC, datetime from pathlib import Path, PurePath, PurePosixPath from typing import Any, Literal @@ -14,7 +14,7 @@ import pandas as pd -from judgearena.usage import request_usage_to_json +from judgearena.usage import RequestUsage, request_usage_from_json COMPLETION_DB_NAME = "completions.db" JUDGEMENT_DB_NAME = "judgements.db" @@ -27,6 +27,21 @@ def stable_json_dumps(value: Any) -> str: return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) +def cached_usage_to_json(usage: RequestUsage | None) -> str | None: + """Persist reusable usage without the original request's dollar charge.""" + if usage is None: + return None + values = asdict(usage) + values.pop("cost_usd") + return stable_json_dumps(values) + + +def cached_usage_from_json(value: str | None) -> RequestUsage | None: + """Restore tokens, discarding charges also present in older cache rows.""" + usage = request_usage_from_json(value) + return replace(usage, cost_usd=None) if usage is not None else None + + def descriptor_hash(descriptor: dict[str, Any]) -> str: return hashlib.sha256(stable_json_dumps(descriptor).encode()).hexdigest()[:16] @@ -246,7 +261,7 @@ def save( pushed_by, now, resolved_run_id, - request_usage_to_json(row.get("usage_json")), + cached_usage_to_json(row.get("usage_json")), ) for _, row in rows.iterrows() ] @@ -344,7 +359,7 @@ def save( pushed_by, now, resolved_run_id, - request_usage_to_json(row.get("usage_json")), + cached_usage_to_json(row.get("usage_json")), ) for _, row in rows.iterrows() ] diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py index b8d3d47e..44ea8f13 100644 --- a/tests/test_cache_sqlite.py +++ b/tests/test_cache_sqlite.py @@ -208,7 +208,14 @@ def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): "instruction_id": "1", "model": "Dummy/model", "usage_json": RequestUsage( - stage="generation", input_tokens=3, output_tokens=5 + stage="generation", + input_tokens=3, + output_tokens=5, + cached_tokens=1, + reasoning_tokens=2, + total_tokens=8, + model="Dummy/model", + cost_usd=0.25, ), } ] @@ -219,14 +226,13 @@ def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): assert "usage_json" in result.index assert json.loads(result["usage_json"]) == { - "cached_tokens": None, - "cost_usd": None, + "cached_tokens": 1, "input_tokens": 3, - "model": None, + "model": "Dummy/model", "output_tokens": 5, - "reasoning_tokens": None, + "reasoning_tokens": 2, "stage": "generation", - "total_tokens": None, + "total_tokens": 8, } diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index d6d395ff..06c9daf2 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,6 +1,8 @@ import json import sqlite3 +from dataclasses import replace +import pandas as pd import pytest from langchain_core.messages import AIMessage from langchain_core.prompts import ChatPromptTemplate @@ -13,6 +15,7 @@ canonicalize_model_input, provider_input_mode, ) +from judgearena.cache.sqlite import cached_usage_to_json from judgearena.inference import InferenceResult from judgearena.models import do_inference, prepare_model from judgearena.usage import RequestUsage, track_usage @@ -130,7 +133,14 @@ def batch(self, inputs, **_kwargs): InferenceResult( text="answer", usage=RequestUsage( - stage="generation", input_tokens=7, output_tokens=11 + stage="generation", + input_tokens=7, + output_tokens=11, + cached_tokens=2, + reasoning_tokens=3, + total_tokens=18, + model="Dummy/usage-model", + cost_usd=0.42, ), ) for _ in inputs @@ -146,6 +156,18 @@ def batch(self, inputs, **_kwargs): return_results=True, ) assert tracker.snapshot().requests == (first[0].usage,) + assert first[0].usage.cost_usd == 0.42 + + stored_usage = json.loads(cached_usage_to_json(first[0].usage)) + assert stored_usage == { + "cached_tokens": 2, + "input_tokens": 7, + "model": "Dummy/usage-model", + "output_tokens": 11, + "reasoning_tokens": 3, + "stage": "generation", + "total_tokens": 18, + } with track_usage() as tracker: second = do_inference( @@ -155,10 +177,27 @@ def batch(self, inputs, **_kwargs): return_results=True, ) - assert second[0].usage == first[0].usage + assert second[0].usage == replace(first[0].usage, cost_usd=None) + assert second[0].usage.cost_usd is None assert tracker.snapshot().requests == () +@pytest.mark.parametrize( + ("cache", "row"), + [ + (CompletionInferenceCache, {"completion": "answer"}), + (JudgementInferenceCache, {"judge_completion": "answer", "top_logprobs": None}), + ], +) +def test_cached_result_discards_legacy_cost(cache, row): + usage = json.dumps({"stage": "generation", "input_tokens": 7, "cost_usd": 0.42}) + result = cache("unused", "arena-hard").cached_result( + pd.Series({**row, "usage_json": usage}) + ) + assert result.usage.input_tokens == 7 + assert result.usage.cost_usd is None + + def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): monkeypatch.setattr(inference.importlib_metadata, "version", lambda _name: "0.10.2") cache = CompletionInferenceCache(tmp_path, "arena-hard") From d9226433992de85f849e44f38499ef27ec7e9706 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:48:00 +0200 Subject: [PATCH 3/4] Compact native usage regression tests around persisted artifacts --- tests/test_cache_sqlite.py | 115 +++++++++++++--------------- tests/test_inference_cache.py | 30 +++----- tests/test_meta_eval_annotations.py | 32 ++------ 3 files changed, 71 insertions(+), 106 deletions(-) diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py index 44ea8f13..0914b930 100644 --- a/tests/test_cache_sqlite.py +++ b/tests/test_cache_sqlite.py @@ -1,5 +1,6 @@ import json import sqlite3 +from dataclasses import asdict import pandas as pd import pytest @@ -21,6 +22,33 @@ "sampling": {"max_tokens": 1024, "temperature": 0.0}, } +LEGACY_COMPLETIONS_SCHEMA = """CREATE TABLE completions ( + input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, + completion TEXT NOT NULL, benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, model TEXT NOT NULL, + pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL +)""" + + +def create_legacy_completion_db(path, *, completion=None): + with sqlite3.connect(path) as db: + db.execute(LEGACY_COMPLETIONS_SCHEMA) + if completion is not None: + db.execute( + "INSERT INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + input_hash("old prompt"), + "old prompt", + completion, + "arena", + "1", + "Dummy/model", + "bob", + "2030-01-01T00:00:00+00:00", + "old-run", + ), + ) + def test_descriptor_paths_separate_completion_and_judgement_caches(tmp_path): model = "VLLM/Qwen/Qwen3-8B" @@ -187,17 +215,17 @@ def test_merge_from_updates_live_database_in_place(tmp_path): def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): db_path = tmp_path / COMPLETION_DB_NAME - legacy = sqlite3.connect(db_path) - legacy.execute( - """CREATE TABLE completions ( - input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, - completion TEXT NOT NULL, benchmark TEXT NOT NULL, - instruction_id TEXT NOT NULL, model TEXT NOT NULL, - pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL - )""" + create_legacy_completion_db(db_path) + usage = RequestUsage( + stage="generation", + input_tokens=3, + output_tokens=5, + cached_tokens=1, + reasoning_tokens=2, + total_tokens=8, + model="Dummy/model", + cost_usd=0.25, ) - legacy.commit() - legacy.close() rows = pd.DataFrame( [ @@ -207,16 +235,7 @@ def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): "benchmark": "arena", "instruction_id": "1", "model": "Dummy/model", - "usage_json": RequestUsage( - stage="generation", - input_tokens=3, - output_tokens=5, - cached_tokens=1, - reasoning_tokens=2, - total_tokens=8, - model="Dummy/model", - cost_usd=0.25, - ), + "usage_json": usage, } ] ) @@ -225,15 +244,10 @@ def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): result = cache.query().iloc[0] assert "usage_json" in result.index - assert json.loads(result["usage_json"]) == { - "cached_tokens": 1, - "input_tokens": 3, - "model": "Dummy/model", - "output_tokens": 5, - "reasoning_tokens": 2, - "stage": "generation", - "total_tokens": 8, - } + stored_usage = json.loads(result["usage_json"]) + expected_usage = asdict(usage) + expected_usage.pop("cost_usd") + assert stored_usage == expected_usage @pytest.mark.parametrize( @@ -246,33 +260,18 @@ def test_merge_legacy_incoming_database_projects_null_usage( local_path = tmp_path / "local" / COMPLETION_DB_NAME incoming_path = tmp_path / "legacy" / COMPLETION_DB_NAME incoming_path.parent.mkdir(parents=True) - legacy = sqlite3.connect(incoming_path) - legacy.execute( - """CREATE TABLE completions ( - input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, - completion TEXT NOT NULL, benchmark TEXT NOT NULL, - instruction_id TEXT NOT NULL, model TEXT NOT NULL, - pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL - )""" - ) - legacy.execute( - "INSERT INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - input_hash("old prompt"), - "old prompt", - incoming_completion, - "arena", - "1", - "Dummy/model", - "bob", - "2030-01-01T00:00:00+00:00", - "old-run", - ), - ) - legacy.commit() - legacy.close() + create_legacy_completion_db(incoming_path, completion=incoming_completion) prior_usage = RequestUsage(stage="generation", input_tokens=4) + + def assert_legacy_schema(): + with sqlite3.connect(incoming_path) as db: + columns = { + column[1] for column in db.execute("PRAGMA table_info(completions)") + } + assert "usage_json" not in columns + + assert_legacy_schema() with CompletionCache(local_path) as cache: cache.save( pd.DataFrame( @@ -297,10 +296,4 @@ def test_merge_legacy_incoming_database_projects_null_usage( assert json.loads(result["usage_json"])["input_tokens"] == 4 else: assert pd.isna(result["usage_json"]) - - incoming = sqlite3.connect(incoming_path) - incoming_columns = { - column[1] for column in incoming.execute("PRAGMA table_info(completions)") - } - incoming.close() - assert "usage_json" not in incoming_columns + assert_legacy_schema() diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 06c9daf2..01111070 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,6 +1,6 @@ import json import sqlite3 -from dataclasses import replace +from dataclasses import asdict, replace import pandas as pd import pytest @@ -15,7 +15,7 @@ canonicalize_model_input, provider_input_mode, ) -from judgearena.cache.sqlite import cached_usage_to_json +from judgearena.cache.sqlite import CompletionCache as SQLiteCompletionCache from judgearena.inference import InferenceResult from judgearena.models import do_inference, prepare_model from judgearena.usage import RequestUsage, track_usage @@ -158,16 +158,11 @@ def batch(self, inputs, **_kwargs): assert tracker.snapshot().requests == (first[0].usage,) assert first[0].usage.cost_usd == 0.42 - stored_usage = json.loads(cached_usage_to_json(first[0].usage)) - assert stored_usage == { - "cached_tokens": 2, - "input_tokens": 7, - "model": "Dummy/usage-model", - "output_tokens": 11, - "reasoning_tokens": 3, - "stage": "generation", - "total_tokens": 18, - } + expected_usage = asdict(first[0].usage) + expected_usage.pop("cost_usd") + with SQLiteCompletionCache(next(tmp_path.glob("**/completions.db"))) as store: + stored_usage = json.loads(store.query().iloc[0]["usage_json"]) + assert stored_usage == expected_usage with track_usage() as tracker: second = do_inference( @@ -182,15 +177,10 @@ def batch(self, inputs, **_kwargs): assert tracker.snapshot().requests == () -@pytest.mark.parametrize( - ("cache", "row"), - [ - (CompletionInferenceCache, {"completion": "answer"}), - (JudgementInferenceCache, {"judge_completion": "answer", "top_logprobs": None}), - ], -) -def test_cached_result_discards_legacy_cost(cache, row): +@pytest.mark.parametrize("cache", [CompletionInferenceCache, JudgementInferenceCache]) +def test_cached_result_discards_legacy_cost(cache): usage = json.dumps({"stage": "generation", "input_tokens": 7, "cost_usd": 0.42}) + row = {"completion": "answer", "judge_completion": "answer", "top_logprobs": None} result = cache("unused", "arena-hard").cached_result( pd.Series({**row, "usage_json": usage}) ) diff --git a/tests/test_meta_eval_annotations.py b/tests/test_meta_eval_annotations.py index 5611205f..9817c070 100644 --- a/tests/test_meta_eval_annotations.py +++ b/tests/test_meta_eval_annotations.py @@ -15,7 +15,7 @@ ) from judgearena.inference import InferenceResult from judgearena.prompts.registry import resolve_judge_prompt -from judgearena.usage import RequestUsage +from judgearena.usage import RequestUsage, request_usage_from_json def _sample(): @@ -53,21 +53,12 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): sample[column] = sample[column].map( lambda turns: np.asarray(turns, dtype=object) ) + batches = [ + ("score_A: 9\nscore_B: 1", "score_A: 5\nscore_B: 6", "unparseable"), + ("score_A: 1\nscore_B: 7", "unparseable", "unparseable"), + ] responses = iter( - [ - [ - InferenceResult(text=text) - for text in ( - "score_A: 9\nscore_B: 1", - "score_A: 5\nscore_B: 6", - "unparseable", - ) - ], - [ - InferenceResult(text=text) - for text in ("score_A: 1\nscore_B: 7", "unparseable", "unparseable") - ], - ] + [[InferenceResult(text=text) for text in batch] for batch in batches] ) monkeypatch.setattr(evaluate, "do_inference", lambda **kwargs: next(responses)) rows = annotate_sample( @@ -138,16 +129,7 @@ def fake_do_inference(**kwargs): assert rows.loc[0, "error"] == "context_length" assert pd.isna(rows.loc[0, "usage_json"]) assert rows.loc[1, "error"] is None - assert json.loads(rows.loc[1, "usage_json"]) == { - "stage": "judging", - "model": "judge", - "input_tokens": 7, - "output_tokens": None, - "total_tokens": None, - "reasoning_tokens": None, - "cached_tokens": None, - "cost_usd": None, - } + assert request_usage_from_json(rows.loc[1, "usage_json"]) == usage def test_aggregate_rejects_incomplete_orientation_sets(): From 7bd561c52382d4d4a7e00c23858feed67ab81e02 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:09:05 +0200 Subject: [PATCH 4/4] Fold native usage coverage into existing cache tests --- tests/test_cache_sqlite.py | 97 +++-------------------------- tests/test_inference_cache.py | 69 +++----------------- tests/test_meta_eval_annotations.py | 46 +++++--------- 3 files changed, 31 insertions(+), 181 deletions(-) diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py index 0914b930..56517a73 100644 --- a/tests/test_cache_sqlite.py +++ b/tests/test_cache_sqlite.py @@ -1,6 +1,5 @@ import json import sqlite3 -from dataclasses import asdict import pandas as pd import pytest @@ -22,33 +21,6 @@ "sampling": {"max_tokens": 1024, "temperature": 0.0}, } -LEGACY_COMPLETIONS_SCHEMA = """CREATE TABLE completions ( - input_hash TEXT PRIMARY KEY, input_text TEXT NOT NULL, - completion TEXT NOT NULL, benchmark TEXT NOT NULL, - instruction_id TEXT NOT NULL, model TEXT NOT NULL, - pushed_by TEXT NOT NULL, pushed_at TEXT NOT NULL, run_id TEXT NOT NULL -)""" - - -def create_legacy_completion_db(path, *, completion=None): - with sqlite3.connect(path) as db: - db.execute(LEGACY_COMPLETIONS_SCHEMA) - if completion is not None: - db.execute( - "INSERT INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - input_hash("old prompt"), - "old prompt", - completion, - "arena", - "1", - "Dummy/model", - "bob", - "2030-01-01T00:00:00+00:00", - "old-run", - ), - ) - def test_descriptor_paths_separate_completion_and_judgement_caches(tmp_path): model = "VLLM/Qwen/Qwen3-8B" @@ -215,16 +187,12 @@ def test_merge_from_updates_live_database_in_place(tmp_path): def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): db_path = tmp_path / COMPLETION_DB_NAME - create_legacy_completion_db(db_path) + with CompletionCache(db_path) as cache: + cache._connect() + with sqlite3.connect(db_path) as db: + db.execute("ALTER TABLE completions DROP COLUMN usage_json") usage = RequestUsage( - stage="generation", - input_tokens=3, - output_tokens=5, - cached_tokens=1, - reasoning_tokens=2, - total_tokens=8, - model="Dummy/model", - cost_usd=0.25, + stage="generation", input_tokens=3, output_tokens=5, cost_usd=0.25 ) rows = pd.DataFrame( @@ -243,57 +211,6 @@ def test_completion_usage_round_trips_and_legacy_database_upgrades(tmp_path): cache.save(rows, pushed_by="alice") result = cache.query().iloc[0] - assert "usage_json" in result.index stored_usage = json.loads(result["usage_json"]) - expected_usage = asdict(usage) - expected_usage.pop("cost_usd") - assert stored_usage == expected_usage - - -@pytest.mark.parametrize( - ("incoming_completion", "expect_prior_usage"), - [("old answer", False), ("native answer", True)], -) -def test_merge_legacy_incoming_database_projects_null_usage( - tmp_path, incoming_completion, expect_prior_usage -): - local_path = tmp_path / "local" / COMPLETION_DB_NAME - incoming_path = tmp_path / "legacy" / COMPLETION_DB_NAME - incoming_path.parent.mkdir(parents=True) - create_legacy_completion_db(incoming_path, completion=incoming_completion) - - prior_usage = RequestUsage(stage="generation", input_tokens=4) - - def assert_legacy_schema(): - with sqlite3.connect(incoming_path) as db: - columns = { - column[1] for column in db.execute("PRAGMA table_info(completions)") - } - assert "usage_json" not in columns - - assert_legacy_schema() - with CompletionCache(local_path) as cache: - cache.save( - pd.DataFrame( - [ - { - "input_text": "old prompt", - "completion": "native answer", - "benchmark": "arena", - "instruction_id": "1", - "model": "Dummy/model", - "usage_json": prior_usage, - } - ] - ), - pushed_by="alice", - ) - assert cache.merge_from(incoming_path) == 1 - result = cache.query().iloc[0] - - assert result["completion"] == incoming_completion - if expect_prior_usage: - assert json.loads(result["usage_json"])["input_tokens"] == 4 - else: - assert pd.isna(result["usage_json"]) - assert_legacy_schema() + assert stored_usage["input_tokens"] == usage.input_tokens + assert "cost_usd" not in stored_usage diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 01111070..b6705373 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,8 +1,6 @@ import json import sqlite3 -from dataclasses import asdict, replace -import pandas as pd import pytest from langchain_core.messages import AIMessage from langchain_core.prompts import ChatPromptTemplate @@ -15,7 +13,6 @@ canonicalize_model_input, provider_input_mode, ) -from judgearena.cache.sqlite import CompletionCache as SQLiteCompletionCache from judgearena.inference import InferenceResult from judgearena.models import do_inference, prepare_model from judgearena.usage import RequestUsage, track_usage @@ -124,66 +121,16 @@ def batch(self, inputs, **_kwargs): assert tracker.snapshot().requests == () -def test_cached_usage_is_restored_without_recording_new_spend(tmp_path, monkeypatch): - cache = CompletionInferenceCache(tmp_path, "arena-hard") - - class UsageModel: - def batch(self, inputs, **_kwargs): - return [ - InferenceResult( - text="answer", - usage=RequestUsage( - stage="generation", - input_tokens=7, - output_tokens=11, - cached_tokens=2, - reasoning_tokens=3, - total_tokens=18, - model="Dummy/usage-model", - cost_usd=0.42, - ), - ) - for _ in inputs - ] - - monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: UsageModel()) - metadata = [{"instruction_id": "1"}] - with track_usage() as tracker: - first = do_inference( - prepare_model("Dummy/usage-model", cache=cache), - ["prompt"], - cache_row_metadata=metadata, - return_results=True, - ) - assert tracker.snapshot().requests == (first[0].usage,) - assert first[0].usage.cost_usd == 0.42 - - expected_usage = asdict(first[0].usage) - expected_usage.pop("cost_usd") - with SQLiteCompletionCache(next(tmp_path.glob("**/completions.db"))) as store: - stored_usage = json.loads(store.query().iloc[0]["usage_json"]) - assert stored_usage == expected_usage - - with track_usage() as tracker: - second = do_inference( - prepare_model("Dummy/usage-model", cache=cache), - ["prompt"], - cache_row_metadata=metadata, - return_results=True, - ) - - assert second[0].usage == replace(first[0].usage, cost_usd=None) - assert second[0].usage.cost_usd is None - assert tracker.snapshot().requests == () - - @pytest.mark.parametrize("cache", [CompletionInferenceCache, JudgementInferenceCache]) -def test_cached_result_discards_legacy_cost(cache): +def test_cached_result_discards_legacy_cost(cache, tmp_path): usage = json.dumps({"stage": "generation", "input_tokens": 7, "cost_usd": 0.42}) - row = {"completion": "answer", "judge_completion": "answer", "top_logprobs": None} - result = cache("unused", "arena-hard").cached_result( - pd.Series({**row, "usage_json": usage}) - ) + row = { + "completion": "answer", + "judge_completion": "answer", + "top_logprobs": None, + "usage_json": usage, + } + result = cache(tmp_path, "arena-hard").cached_result(row) assert result.usage.input_tokens == 7 assert result.usage.cost_usd is None diff --git a/tests/test_meta_eval_annotations.py b/tests/test_meta_eval_annotations.py index 9817c070..2f8b72b8 100644 --- a/tests/test_meta_eval_annotations.py +++ b/tests/test_meta_eval_annotations.py @@ -53,12 +53,20 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): sample[column] = sample[column].map( lambda turns: np.asarray(turns, dtype=object) ) - batches = [ - ("score_A: 9\nscore_B: 1", "score_A: 5\nscore_B: 6", "unparseable"), - ("score_A: 1\nscore_B: 7", "unparseable", "unparseable"), - ] + usage = RequestUsage(stage="judging", model="judge", input_tokens=7) responses = iter( - [[InferenceResult(text=text) for text in batch] for batch in batches] + [ + [ + InferenceResult(text="score_A: 9\nscore_B: 1", usage=usage), + InferenceResult(text="score_A: 5\nscore_B: 6"), + InferenceResult(text="unparseable"), + ], + [ + InferenceResult(text="score_A: 1\nscore_B: 7", usage=usage), + InferenceResult(text="", error="context_length"), + InferenceResult(text="unparseable"), + ], + ] ) monkeypatch.setattr(evaluate, "do_inference", lambda **kwargs: next(responses)) rows = annotate_sample( @@ -71,6 +79,9 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): assert rows["orientation"].tolist() == ["direct"] * 3 + ["reversed"] * 3 rendered = rows["judge_input"] assert pd.isna(rows.loc[2, "error"]) + assert rows.loc[4, "error"] == "context_length" + assert pd.isna(rows.loc[4, "usage_json"]) + assert request_usage_from_json(rows.loc[3, "usage_json"]) == usage assert rendered[0].index("Alpha answer") < rendered[0].index("Beta answer") assert rendered[3].index("Beta answer") < rendered[3].index("Alpha answer") assert rows.loc[[0, 3], "pref"].tolist() == pytest.approx( @@ -107,31 +118,6 @@ def test_annotation_preserves_parser_label_and_details(monkeypatch): assert battles["pref"].tolist() == [1.0] -def test_annotation_carries_usage_and_context_skip_error(monkeypatch): - usage = RequestUsage(stage="judging", model="judge", input_tokens=7) - - def fake_do_inference(**kwargs): - assert kwargs["return_results"] is True - assert kwargs["return_top_logprobs"] is False - return [ - InferenceResult(text="", error="context_length"), - InferenceResult(text="", usage=usage), - ] - - monkeypatch.setattr(evaluate, "do_inference", fake_do_inference) - sample = pd.concat([_sample()] * 2, ignore_index=True) - rows = annotate_sample( - sample, - _config("fixed"), - judge_chat_model=object(), - resolved_prompt=resolve_judge_prompt(preset="meta-eval-pair-score"), - ) - assert rows.loc[0, "error"] == "context_length" - assert pd.isna(rows.loc[0, "usage_json"]) - assert rows.loc[1, "error"] is None - assert request_usage_from_json(rows.loc[1, "usage_json"]) == usage - - def test_aggregate_rejects_incomplete_orientation_sets(): annotations = pd.DataFrame( {"battle_id": ["q1"], "orientation": ["direct"], "pref": [0.2]}