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..918cdfed 100644 --- a/judgearena/cache/inference.py +++ b/judgearena/cache/inference.py @@ -19,6 +19,7 @@ CompletionCache, JudgementCache, cache_folder, + cached_usage_from_json, stable_json_dumps, write_descriptor, ) @@ -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=cached_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=cached_usage_from_json(row.get("usage_json")), ) diff --git a/judgearena/cache/sqlite.py b/judgearena/cache/sqlite.py index 2d745d21..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,6 +14,8 @@ import pandas as pd +from judgearena.usage import RequestUsage, request_usage_from_json + COMPLETION_DB_NAME = "completions.db" JUDGEMENT_DB_NAME = "judgements.db" DESCRIPTOR_FILENAME = "metadata.json" @@ -25,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] @@ -93,6 +110,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 +120,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 +172,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 +225,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 +236,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 +261,13 @@ def save( pushed_by, now, resolved_run_id, + cached_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 +310,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 +326,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 +359,14 @@ def save( pushed_by, now, resolved_run_id, + cached_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..56517a73 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,34 @@ 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 + 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, cost_usd=0.25 + ) + + rows = pd.DataFrame( + [ + { + "input_text": "prompt", + "completion": "answer", + "benchmark": "arena", + "instruction_id": "1", + "model": "Dummy/model", + "usage_json": usage, + } + ] + ) + with CompletionCache(db_path) as cache: + cache.save(rows, pushed_by="alice") + result = cache.query().iloc[0] + + stored_usage = json.loads(result["usage_json"]) + 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 a8f0ef94..b6705373 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,32 @@ 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 == () + + +@pytest.mark.parametrize("cache", [CompletionInferenceCache, JudgementInferenceCache]) +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, + "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 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..2f8b72b8 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, request_usage_from_json def _sample(): @@ -51,10 +53,19 @@ def test_annotation_swaps_and_aggregates_battles(monkeypatch): sample[column] = sample[column].map( lambda turns: np.asarray(turns, dtype=object) ) + usage = RequestUsage(stage="judging", model="judge", input_tokens=7) 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="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)) @@ -67,6 +78,10 @@ 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( @@ -88,7 +103,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"), diff --git a/tests/test_meta_eval_foundation.py b/tests/test_meta_eval_foundation.py index 798cbbac..dbc6bcf4 100644 --- a/tests/test_meta_eval_foundation.py +++ b/tests/test_meta_eval_foundation.py @@ -15,6 +15,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 @@ -114,7 +115,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)