Judge tuning (3/5): Preserve native token usage

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)

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")),
         )

judgearena/cache/sqlite.py

@@ -6,14 +6,16 @@
 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
 from urllib.parse import quote, unquote
 
 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)

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

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]
 

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]