From b50bf6729aa508ef0ec9761764b9fbb9c6f0cfeb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 09:01:57 -0400 Subject: [PATCH 1/3] Record each row's live calls, cache replays and live time Prediction rows now carry live_calls, cache_hit_calls and live_elapsed_seconds, filled by every row writer (sync household and single-output, chunked, batch, row repair) from the same cache_hit flag the spend ledger records. Rows written before the columns existed stay empty, and the supervisor still accepts their scenario files. Co-Authored-By: Claude Opus 5.5 --- policybench/analysis.py | 70 ++++++++++++++++++++++++++++++++++ policybench/batch_eval.py | 9 +++++ policybench/eval_no_tools.py | 74 +++++++++++++++++++++++++++++++++++- policybench/row_repair.py | 29 +++++++++++++- policybench/runstore.py | 27 +++++++++---- policybench/supervisor.py | 16 +++++++- 6 files changed, 212 insertions(+), 13 deletions(-) diff --git a/policybench/analysis.py b/policybench/analysis.py index 0ae3b331..725ace0d 100644 --- a/policybench/analysis.py +++ b/policybench/analysis.py @@ -1465,6 +1465,76 @@ def model_cost_latency( return out +# Where a household's responses came from; see household_call_provenance. +HOUSEHOLD_CALL_STATUSES = ( + "live", + "partial_replay", + "full_replay", + "no_response", + "unknown", +) + + +def household_call_provenance(predictions: pd.DataFrame) -> pd.DataFrame: + """Where each household's responses came from, by (model, scenario_id). + + ``status`` is ``live`` when the provider answered every response behind the + household's rows, ``full_replay`` when the LiteLLM disk cache replayed + every one, ``partial_replay`` for a mix, and ``no_response`` when no + response reached any row. It is ``unknown`` when any row lacks the + call-provenance columns, as every row written before they existed does. + ``live_elapsed_seconds`` sums the rows' live time, as the rows' + ``elapsed_seconds`` sum to the household's latency: zero for a full + replay, NaN when unknown or untimed (batch rows carry no time). + + Nothing published reads this yet. ``model_cost_latency`` still counts a + replay's lookup time as a call (#228). + """ + keys = ["model", "scenario_id"] + if predictions.empty or not set(keys) <= set(predictions.columns): + return pd.DataFrame( + { + "status": pd.Series(dtype=object), + "live_elapsed_seconds": pd.Series(dtype=float), + }, + index=pd.MultiIndex.from_tuples([], names=keys), + ) + frame = predictions[keys].copy() + for column in ("live_calls", "cache_hit_calls", "live_elapsed_seconds"): + frame[column] = ( + pd.to_numeric(predictions[column], errors="coerce") + if column in predictions.columns + else np.nan + ) + frame["known"] = frame["live_calls"].notna() & frame["cache_hit_calls"].notna() + frame["any_live"] = frame["live_calls"].fillna(0) > 0 + frame["any_hit"] = frame["cache_hit_calls"].fillna(0) > 0 + grouped = frame.groupby(keys) + flags = grouped.agg( + known=("known", "all"), + any_live=("any_live", "any"), + any_hit=("any_hit", "any"), + ) + status = np.select( + [ + ~flags["known"], + flags["any_live"] & flags["any_hit"], + flags["any_hit"], + flags["any_live"], + ], + ["unknown", "partial_replay", "full_replay", "live"], + default="no_response", + ) + live_elapsed = grouped["live_elapsed_seconds"].sum(min_count=1) + return pd.DataFrame( + { + "status": pd.Series(status, index=flags.index, dtype=object), + "live_elapsed_seconds": live_elapsed.where(flags["known"]), + }, + index=flags.index, + ) + + def analyze_no_tools( ground_truth: pd.DataFrame, predictions: pd.DataFrame, diff --git a/policybench/batch_eval.py b/policybench/batch_eval.py index 461e1f66..66f52da1 100644 --- a/policybench/batch_eval.py +++ b/policybench/batch_eval.py @@ -13,6 +13,9 @@ is not model latency. Models evaluated in batch mode show no latency on the leaderboard rather than a misleading one. Batch turnaround is recorded in the run's ``batches/`` state files. +- ``live_calls`` is 1 on rows a batch response answered and 0 otherwise, and + ``cache_hit_calls`` is always 0: batch requests bypass the LiteLLM disk + cache. ``live_elapsed_seconds`` is empty, like ``elapsed_seconds``. - Cost columns are reconstructed from token usage at standard synchronous rates (litellm's price map plus ``PRICE_OVERRIDES_PER_1M``), matching the leaderboard's comparison basis. Actual spend is ~half of the reported @@ -687,6 +690,12 @@ def split(value): "provider_resolved_model": ( result.provider_resolved_model if result else None ), + # A batch request goes to the provider's batch endpoint, never + # through the LiteLLM disk cache, so its one response is live. + # It is untimed, like elapsed_seconds. + "live_calls": 1 if result is not None else 0, + "cache_hit_calls": 0, + "live_elapsed_seconds": None, } ) return rows diff --git a/policybench/eval_no_tools.py b/policybench/eval_no_tools.py index d2f64089..402449ba 100644 --- a/policybench/eval_no_tools.py +++ b/policybench/eval_no_tools.py @@ -54,6 +54,21 @@ logger = logging.getLogger(__name__) +# Where the responses behind a row came from. ``live_calls`` counts the +# responses the provider answered and ``cache_hit_calls`` the ones the LiteLLM +# disk cache (policybench/cache.py) replayed, over the request group whose time +# the row's ``elapsed_seconds`` apportions. Every row of one group repeats the +# counts, so they are never summed across rows. ``live_elapsed_seconds`` is the +# live calls' share of ``elapsed_seconds``, apportioned the same way; a replay +# contributes its cache lookup to ``elapsed_seconds`` and nothing here. Failed +# attempts appear only in the spend ledger. Rows written before these columns +# existed leave all three empty (docs/runbook.md, "Call provenance"). +NO_TOOLS_CALL_PROVENANCE_COLUMNS = ( + "live_calls", + "cache_hit_calls", + "live_elapsed_seconds", +) + # Canonical columns emitted by both no-tools result writers. ``run_id`` is # added only by repeated runs. Keep the worker-schema parity test in sync. NO_TOOLS_RESULT_COLUMNS = ( @@ -83,6 +98,7 @@ "provider_response_id", "provider_system_fingerprint", "provider_resolved_model", + *NO_TOOLS_CALL_PROVENANCE_COLUMNS, ) # litellm resolves an unprefixed model's provider (and prices it) through its @@ -515,6 +531,37 @@ def _response_cache_hit(response) -> bool: return isinstance(hidden, dict) and hidden.get("cache_hit") is True +def _request_call_provenance(elapsed_seconds: float, cache_hit: bool) -> dict: + """Provenance of one answered request: a provider call or a cache replay.""" + return { + "live_calls": 0 if cache_hit else 1, + "cache_hit_calls": 1 if cache_hit else 0, + "live_elapsed_seconds": 0.0 if cache_hit else elapsed_seconds, + } + + +def _aggregate_call_provenance(results: list[dict]) -> dict: + """Sum request provenance over the responses a row's time is summed from. + + No responses means zero calls of either kind and no live time, matching + ``elapsed_seconds``, which is then empty too. A result without provenance + (one built before the fields existed) makes the whole aggregate unknown + rather than undercounting it. + """ + if any( + result.get("live_calls") is None or result.get("cache_hit_calls") is None + for result in results + ): + return dict.fromkeys(NO_TOOLS_CALL_PROVENANCE_COLUMNS) + return { + "live_calls": sum(int(result["live_calls"]) for result in results), + "cache_hit_calls": sum(int(result["cache_hit_calls"]) for result in results), + "live_elapsed_seconds": _sum_optional_numbers( + result.get("live_elapsed_seconds") for result in results + ), + } + + def _extract_usage_metadata( response, model_id: str, messages: list[dict], content: str ) -> dict: @@ -1382,6 +1429,7 @@ def _aggregate_request_results(results: list[dict]) -> dict: "provider_response_id": None, "provider_system_fingerprint": None, "provider_resolved_model": None, + **_aggregate_call_provenance([]), } return { @@ -1437,6 +1485,7 @@ def _aggregate_request_results(results: list[dict]) -> dict: "provider_resolved_model": _first_non_null( result.get("provider_resolved_model") for result in results ), + **_aggregate_call_provenance(results), } @@ -2140,6 +2189,7 @@ def _request_explanations_once( if not _numeric_values_match(explanation_value, float(answer)): explanations[variable] = None usage = _extract_usage_metadata(response, model_id, messages, content or "") + cache_hit = _response_cache_hit(response) spend_record = _spend_ledger_record( call_key=call_key, phase="explanation_repair", @@ -2156,7 +2206,7 @@ def _request_explanations_once( request_completed_at=request_completed_at, elapsed_seconds=elapsed_seconds, usage=usage, - cache_hit=_response_cache_hit(response), + cache_hit=cache_hit, completion_budget_tokens=request_completion_budget, escalated_from_budget_tokens=escalated_from_budget_tokens, ) @@ -2170,6 +2220,7 @@ def _request_explanations_once( "finish_reason": finish_reason, "completion_budget": request_completion_budget, **usage, + **_request_call_provenance(elapsed_seconds, cache_hit), "spend_ledger": [spend_record], } except SensitivityKnobError: @@ -2386,6 +2437,7 @@ def _request_predictions_once( messages, raw_response or content, ) + cache_hit = _response_cache_hit(response) has_parse_error = any( predictions.get(variable) is None for variable in variables ) or ( @@ -2408,7 +2460,7 @@ def _request_predictions_once( request_completed_at=request_completed_at, elapsed_seconds=elapsed_seconds, usage=usage, - cache_hit=_response_cache_hit(response), + cache_hit=cache_hit, completion_budget_tokens=request_completion_budget, escalated_from_budget_tokens=( escalated_from_budget_tokens if attempt == 0 else None @@ -2429,6 +2481,7 @@ def _request_predictions_once( escalated_from_budget_tokens is not None ), **usage, + **_request_call_provenance(elapsed_seconds, cache_hit), "spend_ledger": spend_records, } except SensitivityKnobError: @@ -2704,6 +2757,7 @@ def run_single_no_tools( "provider_resolved_model": _first_non_null( result.get("provider_resolved_model") for result in chunk_results ), + **_aggregate_call_provenance(chunk_results), "repair_disagreements": [ disagreement for result in chunk_results @@ -3471,6 +3525,7 @@ def run_no_tools_eval( "total_cost_usd": None, "cost_is_estimated": None, "estimated_cost_usd": None, + **_aggregate_call_provenance([]), "failure_sources": {}, "budget_escalation_count": 0, } @@ -3491,6 +3546,7 @@ def run_no_tools_eval( total_cost_usd = result.get("total_cost_usd") cost_is_estimated = result.get("cost_is_estimated") estimated_cost_usd = result.get("estimated_cost_usd") + live_elapsed_seconds = result.get("live_elapsed_seconds") for variable in scenario_programs: prediction = result["predictions"].get(variable) explanation = result.get("explanations", {}).get(variable) @@ -3576,6 +3632,16 @@ def run_no_tools_eval( "provider_resolved_model": result.get( "provider_resolved_model" ), + # Counts cover the household's whole request group, + # so every row repeats them; live time splits like + # elapsed_seconds. + "live_calls": result.get("live_calls"), + "cache_hit_calls": result.get("cache_hit_calls"), + "live_elapsed_seconds": ( + live_elapsed_seconds / batch_size + if live_elapsed_seconds is not None + else None + ), } ) @@ -3723,6 +3789,7 @@ def run_no_tools_single_output_eval( "provider_response_id": None, "provider_system_fingerprint": None, "provider_resolved_model": None, + **_aggregate_call_provenance([]), "failure_sources": {}, "budget_escalation_count": 0, } @@ -3773,6 +3840,9 @@ def run_no_tools_single_output_eval( "provider_resolved_model": result.get( "provider_resolved_model" ), + "live_calls": result.get("live_calls"), + "cache_hit_calls": result.get("cache_hit_calls"), + "live_elapsed_seconds": result.get("live_elapsed_seconds"), } ) diff --git a/policybench/row_repair.py b/policybench/row_repair.py index 8b2ccb6e..c52b6969 100644 --- a/policybench/row_repair.py +++ b/policybench/row_repair.py @@ -209,6 +209,9 @@ def _empty_repair_row( "provider_response_id": None, "provider_system_fingerprint": None, "provider_resolved_model": None, + "live_calls": 0, + "cache_hit_calls": 0, + "live_elapsed_seconds": None, } @@ -267,6 +270,9 @@ def repair_prediction_row( "provider_system_fingerprint" ), "provider_resolved_model": result.get("provider_resolved_model"), + "live_calls": result.get("live_calls"), + "cache_hit_calls": result.get("cache_hit_calls"), + "live_elapsed_seconds": result.get("live_elapsed_seconds"), } except Exception as exc: row = _empty_repair_row( @@ -288,10 +294,31 @@ def repair_prediction_row( def _append_rows(path: Path, rows: list[dict]) -> None: + """Append rows under the file's header, widening it first if needed. + + A repair round resumed by a newer harness can carry columns the file's + header lacks (the call-provenance columns, for one). Appended under the old + header they would shift every later cell, so the file is first rewritten + with the new columns at the end, each existing cell kept as its exact text. + """ if not rows: return frame = pd.DataFrame(rows) - frame.to_csv(path, mode="a", index=False, header=not path.exists()) + header: list[str] = [] + if path.exists(): + try: + header = list(pd.read_csv(path, nrows=0).columns) + except pd.errors.EmptyDataError: + header = [] + if not header: + frame.to_csv(path, index=False) + return + added = [column for column in frame.columns if column not in header] + if added: + existing = pd.read_csv(path, dtype=str, keep_default_na=False, na_filter=False) + header = [*header, *added] + existing.reindex(columns=header, fill_value="").to_csv(path, index=False) + frame.reindex(columns=header).to_csv(path, mode="a", index=False, header=False) def _completed_repair_keys( diff --git a/policybench/runstore.py b/policybench/runstore.py index d3c2af38..da355b5d 100644 --- a/policybench/runstore.py +++ b/policybench/runstore.py @@ -135,6 +135,24 @@ "provider_response_id", "provider_system_fingerprint", "provider_resolved_model", + # Call provenance (eval_no_tools.NO_TOOLS_CALL_PROVENANCE_COLUMNS): absent + # from CSVs written before it existed, which then export without it. + "live_calls", + "cache_hit_calls", + "live_elapsed_seconds", +) + +# Typed copies kept in ``responses.usage_json`` for querying, alongside tokens. +_USAGE_JSON_COLUMNS = ( + "prompt_tokens", + "completion_tokens", + "total_tokens", + "reasoning_tokens", + "cached_prompt_tokens", + "cache_write_prompt_tokens", + "live_calls", + "cache_hit_calls", + "live_elapsed_seconds", ) # Prediction-level columns (one value per output id). @@ -828,14 +846,7 @@ def _response_record_from_csv( ) -> dict[str, Any]: usage = { field: _store_scalar(record.get(field)) - for field in ( - "prompt_tokens", - "completion_tokens", - "total_tokens", - "reasoning_tokens", - "cached_prompt_tokens", - "cache_write_prompt_tokens", - ) + for field in _USAGE_JSON_COLUMNS if field in record } usage = {k: v for k, v in usage.items() if v is not None} diff --git a/policybench/supervisor.py b/policybench/supervisor.py index 5420d627..551f5716 100644 --- a/policybench/supervisor.py +++ b/policybench/supervisor.py @@ -237,7 +237,10 @@ def _expected_outputs_for_scenario(self, index: int) -> list[str]: ) def _scenario_complete(self, index: int) -> bool: - from policybench.eval_no_tools import NO_TOOLS_RESULT_COLUMNS + from policybench.eval_no_tools import ( + NO_TOOLS_CALL_PROVENANCE_COLUMNS, + NO_TOOLS_RESULT_COLUMNS, + ) path = self.scenario_csv(index) if not path.exists(): @@ -256,7 +259,16 @@ def _scenario_complete(self, index: int) -> bool: except Exception as error: self._raise_stale_scenario_output(path, f"could not read CSV: {error}") - missing_columns = sorted(set(NO_TOOLS_RESULT_COLUMNS) - set(frame.columns)) + # A scenario finished before the call-provenance columns existed is + # still complete: its rows predate them and leave them empty. A file + # with only some of them was not written by any harness version. + provenance = set(NO_TOOLS_CALL_PROVENANCE_COLUMNS) + missing_columns = sorted( + set(NO_TOOLS_RESULT_COLUMNS) - provenance - set(frame.columns) + ) + present_provenance = provenance & set(frame.columns) + if present_provenance and present_provenance != provenance: + missing_columns += sorted(provenance - present_provenance) if missing_columns: self._raise_stale_scenario_output( path, From ce67c355374aed9f2aead63956f78830628717a6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 09:38:23 -0400 Subject: [PATCH 2/3] Test call provenance end to end and keep the classifier out of analysis.py Hypothesis drives run_no_tools_eval and the single-output writer through the real request path (chunks, repair rounds, budget escalation and explanation repair) with a fake clock and per-call cache flags. It checks each row's invariants, the counts against the spend ledger's records of the same calls, the CSV, and a byte-identical run-store round trip. Further tests cover batch rows, row repair and its header-widening append, the retry merge, resumed outputs, run-store CSVs that predate the columns, and supervised resumes of scenario files written before them. household_call_provenance moves to policybench/call_provenance.py, so analysis.py, whose hash the paper prints as the scoring code, is unchanged. The runbook documents the columns. Co-Authored-By: Claude Opus 5.5 --- docs/runbook.md | 45 ++ docs/runstore.md | 1 + policybench/analysis.py | 70 --- policybench/call_provenance.py | 84 +++ tests/test_call_provenance.py | 909 +++++++++++++++++++++++++++++++++ tests/test_supervisor.py | 14 + 6 files changed, 1053 insertions(+), 70 deletions(-) create mode 100644 policybench/call_provenance.py create mode 100644 tests/test_call_provenance.py diff --git a/docs/runbook.md b/docs/runbook.md index 825252a6..79039a10 100644 --- a/docs/runbook.md +++ b/docs/runbook.md @@ -437,3 +437,48 @@ for model_dir in sorted((run / "chunks").glob("*")): ) PY ``` + +## 8. Call provenance + +The harness caches provider responses on disk (`policybench/cache.py`). A +request identical to an earlier one returns the cached response, and the +timer measures the lookup, a few hundredths of a second, as if it were the +call. Each prediction row therefore records where its responses came from: + +| Column | Meaning | +| --- | --- | +| `live_calls` | responses behind the row that the provider answered | +| `cache_hit_calls` | responses behind the row that the disk cache replayed | +| `live_elapsed_seconds` | the live calls' share of `elapsed_seconds` | + +The counts cover the request group whose time the row's `elapsed_seconds` +apportions: the whole household for `run_no_tools_eval` (every chunk, repair +round, budget escalation and explanation repair), the one output for +single-output runs and row repairs, and the one request for batch rows. +Every row of a group repeats them, so never sum them across rows. +`live_elapsed_seconds` is divided among the rows like `elapsed_seconds`, so a +household's live time is the sum over its rows. A replay adds its lookup time +to `elapsed_seconds` and nothing to `live_elapsed_seconds`. Failed attempts +are not counted here; the spend ledger records them. + +Batch rows have `cache_hit_calls` 0 (batch requests never read the disk +cache), `live_calls` 1 when a response answered them, and no live time, like +their empty `elapsed_seconds`. Rows written before these columns existed, +including the frozen board's, leave all three empty, and a supervised run +resumed across the change accepts those scenario files. + +The values hold these invariants, which `tests/test_call_provenance.py` +checks with property-based tests against the real request path and against +the spend ledger's records of the same calls: + +- the counts are non-negative integers; +- `0 <= live_elapsed_seconds <= elapsed_seconds`, and both are empty together; +- a row whose responses were all replays has `live_calls` 0 and + `live_elapsed_seconds` exactly 0; +- a row with no replays has `live_elapsed_seconds` equal to `elapsed_seconds`. + +`call_provenance.household_call_provenance` classifies each household as `live`, +`partial_replay`, `full_replay`, `no_response` or `unknown` (rows that +predate the columns), with its summed live time. No published figure uses it: +`latencySeconds` still takes the median of each household's +`elapsed_seconds`, replays included. diff --git a/docs/runstore.md b/docs/runstore.md index 0748668f..e39e8143 100644 --- a/docs/runstore.md +++ b/docs/runstore.md @@ -94,6 +94,7 @@ The CSV has **22 columns** (one row per `(model, scenario_id, variable)`): | `prompt_tokens`, `completion_tokens`, `total_tokens`, `reasoning_tokens`, `cached_prompt_tokens`, `cache_write_prompt_tokens` | `responses.usage_json` | | `provider_reported_cost_usd`, `reconstructed_cost_usd`, `total_cost_usd`, `cost_is_estimated`, `estimated_cost_usd` | `responses.cost_usd` (+ the exact per-row values for byte-identical export) | | `provider_response_id`, `provider_system_fingerprint`, `provider_resolved_model` | dedicated `responses` columns | +| `live_calls`, `cache_hit_calls`, `live_elapsed_seconds` (call provenance, see the runbook; absent from CSVs written before October 2026) | `responses.usage_json` (+ the exact per-row values) | | any future/unknown column | `predictions.extra_json` | The `run_id` column is optional in the CSV (the eval loop only writes it for diff --git a/policybench/analysis.py b/policybench/analysis.py index 725ace0d..0ae3b331 100644 --- a/policybench/analysis.py +++ b/policybench/analysis.py @@ -1465,76 +1465,6 @@ def model_cost_latency( return out -# Where a household's responses came from; see household_call_provenance. -HOUSEHOLD_CALL_STATUSES = ( - "live", - "partial_replay", - "full_replay", - "no_response", - "unknown", -) - - -def household_call_provenance(predictions: pd.DataFrame) -> pd.DataFrame: - """Where each household's responses came from, by (model, scenario_id). - - ``status`` is ``live`` when the provider answered every response behind the - household's rows, ``full_replay`` when the LiteLLM disk cache replayed - every one, ``partial_replay`` for a mix, and ``no_response`` when no - response reached any row. It is ``unknown`` when any row lacks the - call-provenance columns, as every row written before they existed does. - ``live_elapsed_seconds`` sums the rows' live time, as the rows' - ``elapsed_seconds`` sum to the household's latency: zero for a full - replay, NaN when unknown or untimed (batch rows carry no time). - - Nothing published reads this yet. ``model_cost_latency`` still counts a - replay's lookup time as a call (#228). - """ - keys = ["model", "scenario_id"] - if predictions.empty or not set(keys) <= set(predictions.columns): - return pd.DataFrame( - { - "status": pd.Series(dtype=object), - "live_elapsed_seconds": pd.Series(dtype=float), - }, - index=pd.MultiIndex.from_tuples([], names=keys), - ) - frame = predictions[keys].copy() - for column in ("live_calls", "cache_hit_calls", "live_elapsed_seconds"): - frame[column] = ( - pd.to_numeric(predictions[column], errors="coerce") - if column in predictions.columns - else np.nan - ) - frame["known"] = frame["live_calls"].notna() & frame["cache_hit_calls"].notna() - frame["any_live"] = frame["live_calls"].fillna(0) > 0 - frame["any_hit"] = frame["cache_hit_calls"].fillna(0) > 0 - grouped = frame.groupby(keys) - flags = grouped.agg( - known=("known", "all"), - any_live=("any_live", "any"), - any_hit=("any_hit", "any"), - ) - status = np.select( - [ - ~flags["known"], - flags["any_live"] & flags["any_hit"], - flags["any_hit"], - flags["any_live"], - ], - ["unknown", "partial_replay", "full_replay", "live"], - default="no_response", - ) - live_elapsed = grouped["live_elapsed_seconds"].sum(min_count=1) - return pd.DataFrame( - { - "status": pd.Series(status, index=flags.index, dtype=object), - "live_elapsed_seconds": live_elapsed.where(flags["known"]), - }, - index=flags.index, - ) - - def analyze_no_tools( ground_truth: pd.DataFrame, predictions: pd.DataFrame, diff --git a/policybench/call_provenance.py b/policybench/call_provenance.py new file mode 100644 index 00000000..3b804ce9 --- /dev/null +++ b/policybench/call_provenance.py @@ -0,0 +1,84 @@ +"""Where each household's responses came from: provider calls or cache replays. + +Prediction rows record ``live_calls``, ``cache_hit_calls`` and +``live_elapsed_seconds`` (``eval_no_tools.NO_TOOLS_CALL_PROVENANCE_COLUMNS``; +docs/runbook.md, "Call provenance"). This module reads them back per +household. It sits outside ``analysis.py`` because nothing published reads it: +``model_cost_latency`` still counts a replay's lookup time as a call (#228). +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from policybench.eval_no_tools import NO_TOOLS_CALL_PROVENANCE_COLUMNS + +HOUSEHOLD_CALL_STATUSES = ( + "live", + "partial_replay", + "full_replay", + "no_response", + "unknown", +) + + +def household_call_provenance(predictions: pd.DataFrame) -> pd.DataFrame: + """Where each household's responses came from, by (model, scenario_id). + + ``status`` is ``live`` when the provider answered every response behind the + household's rows, ``full_replay`` when the LiteLLM disk cache replayed + every one, ``partial_replay`` for a mix, and ``no_response`` when no + response reached any row. It is ``unknown`` when any row lacks the + call-provenance columns, as every row written before they existed does. + The test is per row, so it holds whether a household's rows share one + request group (household runs) or each have their own (single-output runs, + row repairs, batch chunks). + + ``live_elapsed_seconds`` sums the rows' live time, as the rows' + ``elapsed_seconds`` sum to the household's latency: zero for a full + replay, NaN when unknown or untimed (batch rows carry no time). + """ + keys = ["model", "scenario_id"] + if predictions.empty or not set(keys) <= set(predictions.columns): + return pd.DataFrame( + { + "status": pd.Series(dtype=object), + "live_elapsed_seconds": pd.Series(dtype=float), + }, + index=pd.MultiIndex.from_tuples([], names=keys), + ) + frame = predictions[keys].copy() + for column in NO_TOOLS_CALL_PROVENANCE_COLUMNS: + frame[column] = ( + pd.to_numeric(predictions[column], errors="coerce") + if column in predictions.columns + else np.nan + ) + frame["known"] = frame["live_calls"].notna() & frame["cache_hit_calls"].notna() + frame["any_live"] = frame["live_calls"].fillna(0) > 0 + frame["any_hit"] = frame["cache_hit_calls"].fillna(0) > 0 + grouped = frame.groupby(keys) + flags = grouped.agg( + known=("known", "all"), + any_live=("any_live", "any"), + any_hit=("any_hit", "any"), + ) + status = np.select( + [ + ~flags["known"], + flags["any_live"] & flags["any_hit"], + flags["any_hit"], + flags["any_live"], + ], + ["unknown", "partial_replay", "full_replay", "live"], + default="no_response", + ) + live_elapsed = grouped["live_elapsed_seconds"].sum(min_count=1) + return pd.DataFrame( + { + "status": pd.Series(status, index=flags.index, dtype=object), + "live_elapsed_seconds": live_elapsed.where(flags["known"]), + }, + index=flags.index, + ) diff --git a/tests/test_call_provenance.py b/tests/test_call_provenance.py new file mode 100644 index 00000000..3e71fff3 --- /dev/null +++ b/tests/test_call_provenance.py @@ -0,0 +1,909 @@ +"""Per-row call provenance: live calls, cache replays and live time. + +Every row writer fills ``live_calls``, ``cache_hit_calls`` and +``live_elapsed_seconds`` (``eval_no_tools.NO_TOOLS_CALL_PROVENANCE_COLUMNS``). +The invariants these tests execute, for every row a writer emits: + +1. ``live_calls`` and ``cache_hit_calls`` are non-negative integers that count + every response behind the row's time, live or replayed. +2. ``0 <= live_elapsed_seconds <= elapsed_seconds``, and the two are empty + together. +3. A row whose responses were all cache hits has zero live calls and exactly + zero live time. +4. A row with no cache hits has live time exactly equal to its elapsed time. +5. The counts and live time agree with the spend ledger's answered records, + an independent record of the same calls (differential). +6. Batch rows never record a cache hit and carry no live time. +7. The columns round-trip through the run store and CSV byte for byte; a CSV + written before they existed round-trips without them and classifies as + unknown. +8. Adding the columns changes no published number. +""" + +import json +import math +import tempfile +import time +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import litellm +import pandas as pd +import pytest +from hypothesis import HealthCheck, event, given, settings +from hypothesis import strategies as st + +import policybench.eval_no_tools as eval_no_tools +import policybench.row_repair as row_repair +from policybench.analysis import model_cost_latency, usage_summary_by_model +from policybench.batch_eval import BatchUnit, NormalizedResult, rows_from_unit +from policybench.call_provenance import ( + HOUSEHOLD_CALL_STATUSES, + household_call_provenance, +) +from policybench.config import COUNTRY_PROGRAMS +from policybench.eval_no_tools import ( + NO_TOOLS_CALL_PROVENANCE_COLUMNS, + NO_TOOLS_RESULT_COLUMNS, + _aggregate_call_provenance, + _aggregate_request_results, + _build_resume_metadata, + _request_call_provenance, + _write_resume_metadata, + run_no_tools_eval, + run_no_tools_single_output_eval, +) +from policybench.retry_eval import merge_retry_predictions +from policybench.row_repair import ( + _append_rows, + merge_row_repair_predictions, + repair_prediction_row, +) +from policybench.runstore import ( + _RESPONSE_COLUMNS, + RunStore, + export_predictions_csv, + import_run_csv, +) +from policybench.scenarios import Person, Scenario +from policybench.spend_ledger import read_spend_ledger, spend_ledger_path + +PROVENANCE = list(NO_TOOLS_CALL_PROVENANCE_COLUMNS) +US_VARIABLES = COUNTRY_PROGRAMS["us"] +FROZEN_US_PREDICTIONS = ( + Path(__file__).resolve().parents[1] + / "paper/snapshot/20260501/runs" + / "us_full_run_20260612_policyengine_4_16_1_populace/predictions.csv.gz" +) + + +def _scenario(scenario_id: str, income: float = 50_000.0) -> Scenario: + return Scenario( + id=scenario_id, + state="CA", + filing_status="single", + adults=[Person(name="adult1", age=35, employment_income=income)], + year=2026, + ) + + +def test_the_provenance_columns_close_the_row_schema_and_are_response_level(): + assert NO_TOOLS_RESULT_COLUMNS[-3:] == NO_TOOLS_CALL_PROVENANCE_COLUMNS + assert set(NO_TOOLS_CALL_PROVENANCE_COLUMNS) <= set(_RESPONSE_COLUMNS) + + +def test_one_request_is_a_live_call_or_a_replay(): + assert _request_call_provenance(12.5, False) == { + "live_calls": 1, + "cache_hit_calls": 0, + "live_elapsed_seconds": 12.5, + } + assert _request_call_provenance(0.02, True) == { + "live_calls": 0, + "cache_hit_calls": 1, + "live_elapsed_seconds": 0.0, + } + + +def test_no_responses_means_zero_calls_and_no_live_time(): + assert _aggregate_call_provenance([]) == { + "live_calls": 0, + "cache_hit_calls": 0, + "live_elapsed_seconds": None, + } + aggregated = _aggregate_request_results([]) + assert aggregated["elapsed_seconds"] is None + assert aggregated["live_elapsed_seconds"] is None + + +def test_a_result_without_provenance_makes_the_aggregate_unknown(): + results = [ + {"elapsed_seconds": 2.0, **_request_call_provenance(2.0, False)}, + {"elapsed_seconds": 1.0}, + ] + + aggregated = _aggregate_request_results(results) + + assert aggregated["elapsed_seconds"] == 3.0 + assert {column: aggregated[column] for column in PROVENANCE} == dict.fromkeys( + PROVENANCE + ) + + +@given( + st.lists( + st.tuples( + st.floats(min_value=0, max_value=1e5, allow_nan=False), + st.booleans(), + ), + max_size=12, + ) +) +def test_aggregated_provenance_counts_every_response_and_bounds_live_time(pairs): + results = [ + {"elapsed_seconds": elapsed, **_request_call_provenance(elapsed, hit)} + for elapsed, hit in pairs + ] + + aggregated = _aggregate_request_results(results) + + hits = sum(hit for _, hit in pairs) + assert aggregated["cache_hit_calls"] == hits + assert aggregated["live_calls"] == len(pairs) - hits + if not pairs: + assert aggregated["live_elapsed_seconds"] is None + return + live = aggregated["live_elapsed_seconds"] + assert 0 <= live <= aggregated["elapsed_seconds"] + if hits == len(pairs): + assert live == 0 + if hits == 0: + assert live == aggregated["elapsed_seconds"] + assert live == pytest.approx(sum(e for e, hit in pairs if not hit)) + + +# -- end to end through the real request path --------------------------------- + +RESPONSE_MODES = ( + # Weighted toward complete answers so most households finish; the other + # modes drive repair rounds, budget escalation and explanation repairs. + "answer", + "answer", + "answer", + "values_only", + "empty", + "truncated", +) + + +def _answer_arguments(with_explanations: bool) -> str: + outputs = {} + for index, variable in enumerate(US_VARIABLES, start=1): + entry = {"value": float(index)} + if with_explanations: + entry["explanation"] = f"Explanation for {variable}. value = {index}" + outputs[variable] = entry + return json.dumps({"outputs": outputs}) + + +def _mock_response(mode: str, cache_hit: bool): + tool_calls = [] + if mode in ("answer", "values_only"): + tool_calls = [ + SimpleNamespace( + function=SimpleNamespace( + name="submit_outputs", + arguments=_answer_arguments(mode == "answer"), + ) + ) + ] + message = SimpleNamespace(content=None, tool_calls=tool_calls, function_call=None) + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=message, + finish_reason="length" if mode == "truncated" else "stop", + ) + ], + usage=litellm.Usage(prompt_tokens=12, completion_tokens=3, total_tokens=15), + # litellm's caching handler sets this on every disk-cache replay + # (litellm/caching/caching_handler.py); a provider response lacks it. + _hidden_params={"cache_hit": cache_hit}, + ) + + +def _assert_row_invariants(rows: pd.DataFrame) -> None: + for column in ("live_calls", "cache_hit_calls"): + assert rows[column].notna().all() + assert (rows[column] >= 0).all() + assert (rows[column] == rows[column].round()).all() + live = rows["live_elapsed_seconds"] + elapsed = rows["elapsed_seconds"] + assert (live.isna() == elapsed.isna()).all() + timed = elapsed.notna() + assert (live[timed] >= 0).all() + assert (live[timed] <= elapsed[timed]).all() + all_hits = (rows["live_calls"] == 0) & (rows["cache_hit_calls"] > 0) + assert (live[all_hits & timed] == 0).all() + no_hits = rows["cache_hit_calls"] == 0 + assert (live[no_hits & timed] == elapsed[no_hits & timed]).all() + + +@settings( + max_examples=60, + deadline=None, + suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], +) +@given( + data=st.data(), + writer=st.sampled_from(["household", "single_output"]), + model_id=st.sampled_from(["claude-sonnet-4-6", "claude-opus-4-6"]), + output_count=st.integers(min_value=1, max_value=3), + household_count=st.integers(min_value=1, max_value=2), +) +def test_rows_record_each_responses_provenance_and_agree_with_the_ledger( + data, writer, model_id, output_count, household_count +): + """Drive run_no_tools_eval / run_no_tools_single_output_eval through the + real request path. Each mocked response's cache flag, duration and shape + (a full answer, values only, nothing, or a length cut) is drawn per call, + so households mix replays with live calls across chunks (Claude Sonnet + 4.6 answers one output per request), repair rounds, budget escalations + and explanation repairs.""" + programs = US_VARIABLES[:output_count] + scenarios = [ + _scenario(f"scenario_{index:03d}", 40_000.0 + index) + for index in range(household_count) + ] + clock = [1_000.0] + current = {"group": None} + calls: list[tuple[tuple, bool, float]] = [] + real_run_single = eval_no_tools.run_single_no_tools + + def tracking_run_single(scenario, variable, *args, **kwargs): + key = (scenario.id,) if writer == "household" else (scenario.id, variable) + current["group"] = key + return real_run_single(scenario, variable, *args, **kwargs) + + def fake_completion(**_kwargs): + mode = data.draw(st.sampled_from(RESPONSE_MODES), label="mode") + cache_hit = data.draw(st.booleans(), label="cache_hit") + seconds = data.draw( + st.floats(min_value=0, max_value=900, allow_nan=False), + label="seconds", + ) + clock[0] += seconds + calls.append((current["group"], cache_hit, seconds)) + return _mock_response(mode, cache_hit) + + fake_time = SimpleNamespace( + perf_counter=lambda: clock[0], + time=time.time, + sleep=lambda _seconds: None, + ) + runner = ( + run_no_tools_eval + if writer == "household" + else (run_no_tools_single_output_eval) + ) + with ( + tempfile.TemporaryDirectory() as tmp, + patch.object(eval_no_tools, "completion", side_effect=fake_completion), + patch.object(eval_no_tools, "time", fake_time), + patch.object(eval_no_tools, "run_single_no_tools", tracking_run_single), + ): + output_path = Path(tmp) / "run" / "predictions.csv" + output_path.parent.mkdir() + frame = runner( + scenarios, + models={"Model": model_id}, + programs=programs, + output_path=str(output_path), + ) + ledger = read_spend_ledger(spend_ledger_path(output_path)) + + assert list(frame.columns) == list(NO_TOOLS_RESULT_COLUMNS) + _assert_row_invariants(frame) + + group_columns = ( + ["scenario_id"] + if writer == "household" + else [ + "scenario_id", + "variable", + ] + ) + for key, rows in frame.groupby(group_columns): + group_calls = [call for call in calls if call[0] == tuple(key)] + hits = sum(hit for _, hit, _ in group_calls) + # One request group per household (or per output): every row + # repeats its counts. + assert rows["live_calls"].nunique() == 1 + assert rows["cache_hit_calls"].nunique() == 1 + assert int(rows["cache_hit_calls"].iloc[0]) == hits + assert int(rows["live_calls"].iloc[0]) == len(group_calls) - hits + live_seconds = sum(seconds for _, hit, seconds in group_calls if not hit) + assert rows["live_elapsed_seconds"].sum() == pytest.approx( + live_seconds, rel=1e-9, abs=1e-6 + ) + + # Differential: the spend ledger records the same calls + # independently, one record per physical request. + records = [ + record + for record in ledger + if record["scenario_id"] == key[0] + and (writer == "household" or record["variables"] == [key[1]]) + and record["status"] in ("ok", "parse_error") + ] + assert len(records) == len(group_calls) + assert sum(record["cache_hit"] for record in records) == hits + assert sum( + record["elapsed_seconds"] + for record in records + if not record["cache_hit"] + ) == pytest.approx(rows["live_elapsed_seconds"].sum(), rel=1e-9, abs=1e-6) + assert sum(record["elapsed_seconds"] for record in records) == ( + pytest.approx(rows["elapsed_seconds"].sum(), rel=1e-9, abs=1e-6) + ) + + # The analysis classifier reads the same facts back from the rows. + status = household_call_provenance(frame)["status"] + for scenario in scenarios: + household_calls = [call for call in calls if call[0][0] == scenario.id] + hits = sum(hit for _, hit, _ in household_calls) + expected = ( + "full_replay" + if hits == len(household_calls) + else "live" + if hits == 0 + else "partial_replay" + ) + assert status[("Model", scenario.id)] == expected + event(f"household {expected}") + event(f"household calls {min(len(household_calls), 6)}+") + + # The CSV carries the same values, and the run store returns it byte + # for byte with the counts queryable from usage_json. + written = pd.read_csv(output_path, float_precision="round_trip") + pd.testing.assert_frame_equal( + written[PROVENANCE].astype(float), + frame[PROVENANCE].astype(float), + check_exact=True, + ) + db_path = Path(tmp) / "run.db" + run_id = import_run_csv(db_path, output_path) + exported = export_predictions_csv(db_path, run_id, Path(tmp) / "out.csv") + assert exported.read_bytes() == output_path.read_bytes() + with RunStore(db_path) as store: + stored = store.connection.execute( + "SELECT scenario_id, usage_json FROM responses WHERE run_id = ?", + (run_id,), + ).fetchall() + for row in stored: + usage = json.loads(row["usage_json"]) + expected_rows = frame[frame["scenario_id"] == row["scenario_id"]] + assert usage["cache_hit_calls"] in set(expected_rows["cache_hit_calls"]) + assert usage["live_calls"] in set(expected_rows["live_calls"]) + + +def test_a_household_that_raised_records_no_responses(monkeypatch, tmp_path): + def raising_run_single(*_args, **_kwargs): + raise ValueError("malformed provider payload") + + monkeypatch.setattr(eval_no_tools, "run_single_no_tools", raising_run_single) + + frame = run_no_tools_eval( + [_scenario("scenario_000")], + models={"Model": "claude-opus-4-6"}, + programs=US_VARIABLES[:2], + output_path=str(tmp_path / "predictions.csv"), + ) + + assert set(frame["live_calls"]) == {0} + assert set(frame["cache_hit_calls"]) == {0} + assert frame["live_elapsed_seconds"].isna().all() + assert frame["elapsed_seconds"].isna().all() + assert set(household_call_provenance(frame)["status"]) == {"no_response"} + + +def test_resumed_output_keeps_older_rows_empty(monkeypatch, tmp_path): + """Rows written before the columns existed stay empty when a run resumes + and later households fill them.""" + scenarios = [_scenario("scenario_000"), _scenario("scenario_001", 60_000.0)] + models = {"Model": "claude-opus-4-6"} + programs = US_VARIABLES[:2] + output_path = tmp_path / "predictions.csv" + _write_resume_metadata( + str(output_path), + _build_resume_metadata( + task="eval_no_tools_batch", + scenarios=scenarios, + models=models, + programs=programs, + run_id=None, + include_explanations=True, + ), + ) + older_columns = [c for c in NO_TOOLS_RESULT_COLUMNS if c not in PROVENANCE] + older = pd.DataFrame( + [ + { + **dict.fromkeys(older_columns), + "call_id": "Model:scenario_000", + "model": "Model", + "scenario_id": "scenario_000", + "variable": variable, + "prediction": 1.0, + "explanation": "Explanation. value = 1", + "raw_response": "{}", + "elapsed_seconds": 0.01, + } + for variable in programs + ], + columns=older_columns, + ) + older.to_csv(output_path, index=False) + + def live_run_single(_scenario, variables, *_args, **_kwargs): + return { + "predictions": dict.fromkeys(variables, 2.0), + "explanations": dict.fromkeys(variables, "Explanation. value = 2"), + "raw_response": "{}", + "error": None, + "elapsed_seconds": 8.0, + **_request_call_provenance(8.0, False), + } + + monkeypatch.setattr(eval_no_tools, "run_single_no_tools", live_run_single) + run_no_tools_eval( + scenarios, + models=models, + programs=programs, + output_path=str(output_path), + ) + + written = pd.read_csv(output_path) + assert list(written.columns) == list(NO_TOOLS_RESULT_COLUMNS) + resumed = written[written["scenario_id"] == "scenario_000"] + fresh = written[written["scenario_id"] == "scenario_001"] + assert resumed[PROVENANCE].isna().all().all() + assert set(fresh["live_calls"]) == {1} + assert set(fresh["cache_hit_calls"]) == {0} + assert fresh["live_elapsed_seconds"].tolist() == [4.0, 4.0] + status = household_call_provenance(written)["status"] + assert status[("Model", "scenario_000")] == "unknown" + assert status[("Model", "scenario_001")] == "live" + + +# -- batch, row repair and retry writers -------------------------------------- + + +@given( + output_count=st.integers(min_value=1, max_value=6), + answered=st.booleans(), + prompt_tokens=st.integers(min_value=0, max_value=100_000), +) +def test_batch_rows_are_live_untimed_and_never_cache_hits( + output_count, answered, prompt_tokens +): + unit = BatchUnit( + scenario_id="scenario_000", + variables=US_VARIABLES[:output_count], + chunk_index=0, + ) + result = ( + NormalizedResult( + custom_id=unit.custom_id, + prompt_tokens=prompt_tokens, + completion_tokens=10, + ) + if answered + else None + ) + + rows = rows_from_unit( + model_name="Model", + model_id="claude-sonnet-4-6", + unit=unit, + predictions={}, + explanations={}, + raw_response=None, + error=None if answered else "batch_missing: no result entry returned", + result=result, + ) + + assert len(rows) == output_count + for row in rows: + assert row["cache_hit_calls"] == 0 + assert row["live_calls"] == int(answered) + assert row["live_elapsed_seconds"] is None + assert row["elapsed_seconds"] is None + frame = pd.DataFrame(rows) + expected = "live" if answered else "no_response" + assert set(household_call_provenance(frame)["status"]) == {expected} + + +def test_row_repair_carries_the_repair_requests_provenance(monkeypatch): + calls = iter( + [ + { + "predictions": {"payroll_tax": None}, + "explanations": {"payroll_tax": None}, + "raw_response": None, + "error": "Missing predictions after repair: payroll_tax", + "elapsed_seconds": 0.02, + **_request_call_provenance(0.02, True), + }, + ValueError("provider returned malformed JSON"), + { + "predictions": {"payroll_tax": 7.0}, + "explanations": {"payroll_tax": "Payroll. value = 7"}, + "raw_response": "{}", + "error": None, + "elapsed_seconds": 5.5, + **_request_call_provenance(5.5, False), + }, + ] + ) + + def scripted_run_single(*_args, **_kwargs): + outcome = next(calls) + if isinstance(outcome, Exception): + raise outcome + return outcome + + monkeypatch.setattr(row_repair, "run_single_no_tools", scripted_run_single) + monkeypatch.setitem(row_repair.MODELS, "Model", "claude-opus-4-6") + + final, attempts = repair_prediction_row( + country="us", + scenario=_scenario("scenario_000"), + model="Model", + variable="payroll_tax", + attempts_per_row=3, + ) + + assert [ + (row["live_calls"], row["cache_hit_calls"], row["live_elapsed_seconds"]) + for row in attempts + ] == [(0, 1, 0.0), (0, 0, None), (1, 0, 5.5)] + assert final["live_elapsed_seconds"] == 5.5 + + +def test_append_rows_widens_an_older_header_and_keeps_its_cells(tmp_path): + path = tmp_path / "row_repair_attempts.csv" + path.write_text( + "model,scenario_id,variable,prediction,explanation\n" + 'm,s1,tax,1.10,"Tax, federal. value = 1.10"\n' + ) + + _append_rows( + path, + [ + { + "variable": "benefit", + "model": "m", + "scenario_id": "s2", + "prediction": 2.0, + "explanation": "Benefit. value = 2", + "live_calls": 1, + "cache_hit_calls": 0, + "live_elapsed_seconds": 3.5, + } + ], + ) + _append_rows( + path, + [{"model": "m", "scenario_id": "s3", "variable": "tax", "prediction": 4.0}], + ) + + lines = path.read_text().splitlines() + assert lines[0] == ( + "model,scenario_id,variable,prediction,explanation," + "live_calls,cache_hit_calls,live_elapsed_seconds" + ) + assert lines[1] == 'm,s1,tax,1.10,"Tax, federal. value = 1.10",,,' + frame = pd.read_csv(path) + assert frame["scenario_id"].tolist() == ["s1", "s2", "s3"] + assert frame.loc[1, "variable"] == "benefit" + assert frame.loc[1, "live_elapsed_seconds"] == 3.5 + assert math.isnan(frame.loc[2, "live_calls"]) + + +def test_append_rows_replaces_a_headerless_empty_file(tmp_path): + path = tmp_path / "row_repair_attempts.csv" + pd.DataFrame().to_csv(path, index=False) + + _append_rows(path, [{"model": "m", "live_calls": 1}]) + + assert pd.read_csv(path).to_dict("records") == [{"model": "m", "live_calls": 1}] + + +def test_row_repair_merge_fills_only_the_repaired_row(tmp_path): + source_path = tmp_path / "source.csv" + repair_path = tmp_path / "repair.csv" + pd.DataFrame( + { + "model": ["model_a", "model_a"], + "scenario_id": ["s1", "s1"], + "variable": ["tax", "benefit"], + "prediction": [1.0, None], + "explanation": ["Tax. value = 1", None], + "raw_response": [None, None], + "error": [None, "Missing prediction after repair: benefit"], + } + ).to_csv(source_path, index=False) + pd.DataFrame( + { + "country": ["us"], + "attempt": [1], + "model": ["model_a"], + "scenario_id": ["s1"], + "variable": ["benefit"], + "prediction": [22.0], + "explanation": ["Repair benefit. value = 22"], + "raw_response": ["raw"], + "error": [None], + "elapsed_seconds": [6.0], + "live_calls": [1], + "cache_hit_calls": [0], + "live_elapsed_seconds": [6.0], + } + ).to_csv(repair_path, index=False) + + outputs = merge_row_repair_predictions( + source_predictions=source_path, + repair_predictions=repair_path, + output_dir=tmp_path / "merged", + ) + + merged = pd.read_csv(outputs["merged_predictions"]).set_index("variable") + assert merged.loc["benefit", "live_calls"] == 1 + assert merged.loc["benefit", "live_elapsed_seconds"] == 6.0 + assert merged.loc["tax", PROVENANCE].isna().all() + assert set(household_call_provenance(merged.reset_index())["status"]) == {"unknown"} + + +def test_retry_merge_carries_the_retry_rows_provenance(tmp_path): + source_path = tmp_path / "source.csv" + retry_path = tmp_path / "retry.csv" + target_path = tmp_path / "target.csv" + pd.DataFrame( + { + "model": ["model_a", "model_a", "model_b"], + "scenario_id": ["s1", "s1", "s1"], + "variable": ["tax", "benefit", "tax"], + "prediction": [1.0, None, 2.0], + "explanation": ["Tax. value = 1", "", "Tax. value = 2"], + "error": [None, "Missing predictions after repair: benefit", None], + } + ).to_csv(source_path, index=False) + pd.DataFrame( + { + "model": ["model_a", "model_a"], + "scenario_id": ["s1", "s1"], + "variable": ["tax", "benefit"], + "prediction": [10.0, 20.0], + "explanation": ["Retry tax. value = 10", "Retry benefit. value = 20"], + "error": [None, None], + "elapsed_seconds": [0.01, 0.01], + "live_calls": [0, 0], + "cache_hit_calls": [1, 1], + "live_elapsed_seconds": [0.0, 0.0], + } + ).to_csv(retry_path, index=False) + pd.DataFrame({"model": ["model_a"], "scenario_id": ["s1"]}).to_csv( + target_path, index=False + ) + + outputs = merge_retry_predictions( + source_predictions=source_path, + retry_predictions=retry_path, + target_units=target_path, + output_dir=tmp_path / "merged", + ) + + merged = pd.read_csv(outputs["merged_predictions"]) + retried = merged[merged["model"] == "model_a"] + kept = merged[merged["model"] == "model_b"] + assert set(retried["cache_hit_calls"]) == {1} + assert set(retried["live_elapsed_seconds"]) == {0.0} + assert kept[PROVENANCE].isna().all().all() + status = household_call_provenance(merged)["status"] + assert status[("model_a", "s1")] == "full_replay" + assert status[("model_b", "s1")] == "unknown" + + +# -- run store ---------------------------------------------------------------- + +optional_count = st.one_of(st.none(), st.integers(min_value=0, max_value=40)) +optional_seconds = st.one_of( + st.none(), st.floats(min_value=0, max_value=1e4, allow_nan=False) +) + + +@settings(max_examples=50, deadline=None) +@given( + st.lists( + st.tuples(optional_count, optional_count, optional_seconds, optional_seconds), + min_size=1, + max_size=6, + ) +) +def test_run_store_round_trips_provenance_byte_for_byte(households): + frame = pd.DataFrame( + [ + { + "call_id": f"Model:scenario_{index:03d}", + "model": "Model", + "scenario_id": f"scenario_{index:03d}", + "variable": variable, + "prediction": 1.0, + "explanation": "Explanation. value = 1", + "raw_response": "{}", + "error": None, + "elapsed_seconds": elapsed, + "live_calls": live_calls, + "cache_hit_calls": cache_hit_calls, + "live_elapsed_seconds": live_elapsed, + } + for index, (live_calls, cache_hit_calls, live_elapsed, elapsed) in ( + enumerate(households) + ) + for variable in ("payroll_tax", "self_employment_tax") + ] + ) + with tempfile.TemporaryDirectory() as tmp: + csv_path = Path(tmp) / "run" / "predictions.csv" + csv_path.parent.mkdir() + frame.to_csv(csv_path, index=False) + db_path = Path(tmp) / "run.db" + + run_id = import_run_csv(db_path, csv_path) + exported = export_predictions_csv(db_path, run_id, Path(tmp) / "out.csv") + + assert exported.read_bytes() == csv_path.read_bytes() + with RunStore(db_path) as store: + usage = { + row["scenario_id"]: json.loads(row["usage_json"]) + for row in store.connection.execute( + "SELECT scenario_id, usage_json FROM responses" + ).fetchall() + } + for index, (live_calls, cache_hit_calls, _, _) in enumerate(households): + stored = usage[f"scenario_{index:03d}"] + assert stored.get("live_calls") == live_calls + assert stored.get("cache_hit_calls") == cache_hit_calls + + +def test_run_store_round_trips_a_csv_that_predates_the_columns(tmp_path): + frame = pd.DataFrame( + { + "call_id": ["Model:scenario_000"] * 2, + "model": ["Model"] * 2, + "scenario_id": ["scenario_000"] * 2, + "variable": ["payroll_tax", "self_employment_tax"], + "prediction": [1.0, 2.0], + "explanation": ["A. value = 1", "B. value = 2"], + "elapsed_seconds": [0.015, 0.015], + } + ) + csv_path = tmp_path / "run" / "predictions.csv" + csv_path.parent.mkdir() + frame.to_csv(csv_path, index=False) + + run_id = import_run_csv(tmp_path / "run.db", csv_path) + exported = export_predictions_csv(tmp_path / "run.db", run_id, tmp_path / "o.csv") + + assert exported.read_bytes() == csv_path.read_bytes() + assert not set(PROVENANCE) & set(pd.read_csv(exported).columns) + + +# -- analysis ------------------------------------------------------------------- + + +def _brute_force_status(rows: list[dict]) -> str: + if any(row["live_calls"] is None or row["cache_hit_calls"] is None for row in rows): + return "unknown" + any_live = any(row["live_calls"] > 0 for row in rows) + any_hit = any(row["cache_hit_calls"] > 0 for row in rows) + if any_live and any_hit: + return "partial_replay" + if any_hit: + return "full_replay" + if any_live: + return "live" + return "no_response" + + +provenance_rows = st.lists( + st.fixed_dictionaries( + { + "variable": st.sampled_from(US_VARIABLES[:4]), + "elapsed_seconds": optional_seconds, + "live_calls": optional_count, + "cache_hit_calls": optional_count, + "live_elapsed_seconds": optional_seconds, + "prompt_tokens": st.integers(min_value=0, max_value=5_000), + "completion_tokens": st.integers(min_value=0, max_value=5_000), + "total_cost_usd": st.floats(min_value=0, max_value=5, allow_nan=False), + } + ), + min_size=1, + max_size=5, +) + + +@settings(max_examples=150, deadline=None) +@given( + st.dictionaries( + st.tuples(st.sampled_from(["m1", "m2"]), st.sampled_from(["s1", "s2", "s3"])), + provenance_rows, + min_size=1, + ) +) +def test_household_classification_and_published_numbers(households): + records = [ + {"model": model, "scenario_id": scenario_id, "prediction": 1.0, **row} + for (model, scenario_id), rows in households.items() + for row in rows + ] + frame = pd.DataFrame(records) + + provenance = household_call_provenance(frame) + + assert set(provenance.index) == set(households) + assert set(provenance["status"]) <= set(HOUSEHOLD_CALL_STATUSES) + for key, rows in households.items(): + assert provenance.loc[key, "status"] == _brute_force_status(rows) + live = provenance.loc[key, "live_elapsed_seconds"] + values = [row["live_elapsed_seconds"] for row in rows] + if provenance.loc[key, "status"] == "unknown" or all( + value is None for value in values + ): + assert math.isnan(live) + else: + assert live == pytest.approx(sum(v for v in values if v is not None)) + + # Adding the provenance columns changes nothing published. + without = frame.drop(columns=PROVENANCE) + assert model_cost_latency(frame) == model_cost_latency(without) + pd.testing.assert_frame_equal( + usage_summary_by_model(frame), usage_summary_by_model(without) + ) + + +def test_frames_without_the_columns_classify_as_unknown(): + frame = pd.DataFrame( + { + "model": ["m", "m"], + "scenario_id": ["s1", "s2"], + "variable": ["tax", "tax"], + "elapsed_seconds": [0.02, 30.0], + } + ) + + provenance = household_call_provenance(frame) + + assert provenance["status"].tolist() == ["unknown", "unknown"] + assert provenance["live_elapsed_seconds"].isna().all() + assert household_call_provenance(pd.DataFrame()).empty + + +@pytest.mark.skipif( + not FROZEN_US_PREDICTIONS.exists(), reason="frozen snapshot not checked out" +) +def test_the_frozen_us_run_predates_the_columns(): + """The frozen board predates the fields, so every household is unknown: + its replays are identified only by the audit's evidence (#228).""" + header = pd.read_csv(FROZEN_US_PREDICTIONS, nrows=0) + assert not set(PROVENANCE) & set(header.columns) + + sample = pd.read_csv( + FROZEN_US_PREDICTIONS, + nrows=500, + usecols=["model", "scenario_id", "variable", "elapsed_seconds"], + ) + assert set(household_call_provenance(sample)["status"]) == {"unknown"} diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py index 1b4240b9..4bef2eaf 100644 --- a/tests/test_supervisor.py +++ b/tests/test_supervisor.py @@ -20,6 +20,7 @@ from policybench.config import MODELS, PROGRAMS, get_programs from policybench.eval_no_tools import ( + NO_TOOLS_CALL_PROVENANCE_COLUMNS, NO_TOOLS_RESULT_COLUMNS, _build_resume_metadata, run_no_tools_eval, @@ -460,6 +461,19 @@ def test_scenario_completion_requires_full_worker_schema( supervisor._scenario_complete(0) +def test_scenario_written_before_call_provenance_is_still_complete(manifest, tmp_path): + """A supervised run resumed across the change keeps its finished + scenarios: their rows predate the call-provenance columns and stay + empty. Only a file with some but not all of them is stale (the + parametrized test above).""" + supervisor = make_supervisor(manifest, tmp_path) + path, _ = write_current_worker_output(supervisor, 0) + frame = pd.read_csv(path).drop(columns=list(NO_TOOLS_CALL_PROVENANCE_COLUMNS)) + frame.to_csv(path, index=False) + + assert supervisor._scenario_complete(0) is True + + @pytest.mark.parametrize( "runner,include_explanations,mismatch", [ From 182decf60137f38c4a28bbc4175e541ed6f35bfe Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 12:40:28 -0400 Subject: [PATCH 3/3] Empty provenance on rows a legacy repair replaces; test failures and both APIs A row-repair merge that applies a repair file written before the provenance columns existed kept the source row's values, which described the replaced request and could exceed the repair's elapsed time. The merge now empties all three on such rows. The end-to-end property now also raises (non-fatal failures on answer and explanation requests), answers the requested tool (submit_explanations for explanation repairs), and runs the Responses API (GPT-5.4, and GPT-5.5 at three outputs per request). Its oracle is the mock's call history: a group with no failed request must count every returned response exactly, and a group with a failure must equal some subset of them, counts and times together. Shrinking is off for it, since each example runs a whole evaluation. Deterministic tests pin a replayed and a live explanation repair and a Responses API replay. Co-Authored-By: Claude Opus 5.5 --- docs/runbook.md | 9 +- policybench/row_repair.py | 23 +- tests/test_call_provenance.py | 438 ++++++++++++++++++++++++++++------ 3 files changed, 393 insertions(+), 77 deletions(-) diff --git a/docs/runbook.md b/docs/runbook.md index 79039a10..d20fbca1 100644 --- a/docs/runbook.md +++ b/docs/runbook.md @@ -465,11 +465,14 @@ Batch rows have `cache_hit_calls` 0 (batch requests never read the disk cache), `live_calls` 1 when a response answered them, and no live time, like their empty `elapsed_seconds`. Rows written before these columns existed, including the frozen board's, leave all three empty, and a supervised run -resumed across the change accepts those scenario files. +resumed across the change accepts those scenario files. A row-repair merge +that applies an older repair, written without the columns, empties them on +the row it replaces rather than keep the replaced request's values. The values hold these invariants, which `tests/test_call_provenance.py` -checks with property-based tests against the real request path and against -the spend ledger's records of the same calls: +checks with property-based tests that drive the real request path, failed +requests included, against the calls actually made and the spend ledger's +records of them: - the counts are non-negative integers; - `0 <= live_elapsed_seconds <= elapsed_seconds`, and both are empty together; diff --git a/policybench/row_repair.py b/policybench/row_repair.py index c52b6969..756afdfb 100644 --- a/policybench/row_repair.py +++ b/policybench/row_repair.py @@ -8,11 +8,16 @@ from dataclasses import dataclass from pathlib import Path +import numpy as np import pandas as pd from policybench.chunked_eval import model_requires_serial_execution from policybench.config import MODELS -from policybench.eval_no_tools import is_infrastructure_error_text, run_single_no_tools +from policybench.eval_no_tools import ( + NO_TOOLS_CALL_PROVENANCE_COLUMNS, + is_infrastructure_error_text, + run_single_no_tools, +) from policybench.reparse_predictions import reparse_predictions_frame from policybench.scenarios import load_scenarios_from_manifest @@ -462,7 +467,21 @@ def merge_row_repair_predictions( valid = pd.DataFrame(columns=repair.columns) merged = source.copy() + # A repair written before the call-provenance columns existed says nothing + # about its request, so the row it replaces loses the source's values: + # they described the replaced request, and kept beside the repair's time + # they could exceed it. + unknown_provenance = [ + column + for column in NO_TOOLS_CALL_PROVENANCE_COLUMNS + if column in merged.columns and column not in valid.columns + ] if not valid.empty: + # Counts read back as int64 unless a cell is empty; a replacement's + # empty cell or a cleared one needs room for NaN. + for column in NO_TOOLS_CALL_PROVENANCE_COLUMNS: + if column in merged.columns: + merged[column] = merged[column].astype("float64") valid_by_key = valid.set_index(KEY_COLUMNS, drop=False) for index, row in merged.iterrows(): key = tuple(str(row[column]) for column in KEY_COLUMNS) @@ -479,6 +498,8 @@ def merge_row_repair_predictions( if isinstance(replacement[column], str): merged[column] = merged[column].astype("object") merged.at[index, column] = replacement[column] + for column in unknown_provenance: + merged.at[index, column] = np.nan accepted_keys = valid[KEY_COLUMNS].drop_duplicates() if not valid.empty else valid target_rows = row_repair_targets( diff --git a/tests/test_call_provenance.py b/tests/test_call_provenance.py index 3e71fff3..d6e8d157 100644 --- a/tests/test_call_provenance.py +++ b/tests/test_call_provenance.py @@ -11,8 +11,11 @@ 3. A row whose responses were all cache hits has zero live calls and exactly zero live time. 4. A row with no cache hits has live time exactly equal to its elapsed time. -5. The counts and live time agree with the spend ledger's answered records, - an independent record of the same calls (differential). +5. The counts and live time agree with the requests actually made (the mock's + call history, which the spend ledger also records): exactly, when none of + a group's requests failed; as a subset of the returned responses, counts + and times together, when one did (a failed chunk or household keeps none of + its time). 6. Batch rows never record a cache hit and carry no live time. 7. The columns round-trip through the run store and CSV byte for byte; a CSV written before they existed round-trips without them and classifies as @@ -24,6 +27,8 @@ import math import tempfile import time +from dataclasses import dataclass +from itertools import combinations from pathlib import Path from types import SimpleNamespace from unittest.mock import patch @@ -31,7 +36,7 @@ import litellm import pandas as pd import pytest -from hypothesis import HealthCheck, event, given, settings +from hypothesis import HealthCheck, Phase, event, given, settings from hypothesis import strategies as st import policybench.eval_no_tools as eval_no_tools @@ -165,16 +170,32 @@ def test_aggregated_provenance_counts_every_response_and_bounds_live_time(pairs) # -- end to end through the real request path --------------------------------- -RESPONSE_MODES = ( +ANSWER_MODES = ( # Weighted toward complete answers so most households finish; the other - # modes drive repair rounds, budget escalation and explanation repairs. + # modes drive repair rounds, budget escalation, explanation repairs, and + # failed requests (retries, partial rounds, failed chunks and households). "answer", "answer", "answer", "values_only", "empty", "truncated", + "raise", ) +EXPLANATION_MODES = ("answer", "answer", "empty", "raise") +# No layer treats this failure as fatal or as provider instability, so the run +# carries on and the household keeps whatever its other requests returned. +REQUEST_FAILURE = "malformed provider payload" +SECONDS = {"rel": 1e-9, "abs": 1e-6} + + +@dataclass(frozen=True) +class MockCall: + group: tuple + cache_hit: bool + seconds: float + raised: bool + explanation_request: bool def _answer_arguments(with_explanations: bool) -> str: @@ -187,15 +208,52 @@ def _answer_arguments(with_explanations: bool) -> str: return json.dumps({"outputs": outputs}) -def _mock_response(mode: str, cache_hit: bool): +def _explanation_arguments() -> str: + return json.dumps( + { + variable: f"Explanation for {variable}. value = {index}" + for index, variable in enumerate(US_VARIABLES, start=1) + } + ) + + +def _mock_response( + mode: str, cache_hit: bool, *, explanation_request: bool, responses_api: bool +): + """A chat or Responses API response answering the requested tool.""" + function_name = arguments = None + if explanation_request and mode == "answer": + function_name, arguments = "submit_explanations", _explanation_arguments() + elif not explanation_request and mode in ("answer", "values_only"): + function_name = "submit_outputs" + arguments = _answer_arguments(mode == "answer") + truncated = mode == "truncated" + # litellm's caching handler sets this on every disk-cache replay + # (litellm/caching/caching_handler.py); a provider response lacks it. + hidden = {"cache_hit": cache_hit} + if responses_api: + output = [] + if function_name: + output = [ + SimpleNamespace( + type="function_call", name=function_name, arguments=arguments + ) + ] + return SimpleNamespace( + output_text="", + output=output, + status="incomplete" if truncated else "completed", + incomplete_details=( + SimpleNamespace(reason="max_output_tokens") if truncated else None + ), + usage=SimpleNamespace(input_tokens=12, output_tokens=3, total_tokens=15), + _hidden_params=hidden, + ) tool_calls = [] - if mode in ("answer", "values_only"): + if function_name: tool_calls = [ SimpleNamespace( - function=SimpleNamespace( - name="submit_outputs", - arguments=_answer_arguments(mode == "answer"), - ) + function=SimpleNamespace(name=function_name, arguments=arguments) ) ] message = SimpleNamespace(content=None, tool_calls=tool_calls, function_call=None) @@ -203,13 +261,11 @@ def _mock_response(mode: str, cache_hit: bool): choices=[ SimpleNamespace( message=message, - finish_reason="length" if mode == "truncated" else "stop", + finish_reason="length" if truncated else "stop", ) ], usage=litellm.Usage(prompt_tokens=12, completion_tokens=3, total_tokens=15), - # litellm's caching handler sets this on every disk-cache replay - # (litellm/caching/caching_handler.py); a provider response lacks it. - _hidden_params={"cache_hit": cache_hit}, + _hidden_params=hidden, ) @@ -221,6 +277,9 @@ def _assert_row_invariants(rows: pd.DataFrame) -> None: live = rows["live_elapsed_seconds"] elapsed = rows["elapsed_seconds"] assert (live.isna() == elapsed.isna()).all() + responded = (rows["live_calls"] + rows["cache_hit_calls"]) > 0 + # Sync rows are timed exactly when some response stands behind them. + assert (responded == elapsed.notna()).all() timed = elapsed.notna() assert (live[timed] >= 0).all() assert (live[timed] <= elapsed[timed]).all() @@ -230,16 +289,63 @@ def _assert_row_invariants(rows: pd.DataFrame) -> None: assert (live[no_hits & timed] == elapsed[no_hits & timed]).all() +def _row_matches_a_subset( + returned: list[MockCall], + hits: int, + live: int, + live_seconds: float, + seconds: float, +) -> bool: + """Whether some subset of the returned responses has the row's counts and + times: ``live`` live responses whose seconds sum to ``live_seconds`` and + ``hits`` replays that bring the total to ``seconds``.""" + hit_calls = [call for call in returned if call.cache_hit] + live_calls = [call for call in returned if not call.cache_hit] + if hits > len(hit_calls) or live > len(live_calls): + return False + if hits + live == 0: + return math.isnan(seconds) and math.isnan(live_seconds) + if math.comb(len(live_calls), live) * math.comb(len(hit_calls), hits) > 200_000: + event("subset check skipped: too many combinations") + return True + for chosen_live in combinations(live_calls, live): + live_sum = sum(call.seconds for call in chosen_live) + if live_sum != pytest.approx(live_seconds, **SECONDS): + continue + for chosen_hits in combinations(hit_calls, hits): + total = live_sum + sum(call.seconds for call in chosen_hits) + if total == pytest.approx(seconds, **SECONDS): + return True + return False + + +def _expected_status(returned: list[MockCall]) -> str: + hits = sum(call.cache_hit for call in returned) + if not returned: + return "no_response" + if hits == len(returned): + return "full_replay" + return "live" if hits == 0 else "partial_replay" + + @settings( - max_examples=60, + max_examples=100, deadline=None, suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], + # Every example runs a whole evaluation, so shrinking a failure would take + # many minutes; the failing example's labelled draws print unshrunk. + phases=[Phase.explicit, Phase.reuse, Phase.generate], ) @given( data=st.data(), writer=st.sampled_from(["household", "single_output"]), - model_id=st.sampled_from(["claude-sonnet-4-6", "claude-opus-4-6"]), - output_count=st.integers(min_value=1, max_value=3), + # Chat completions with one output per request (Claude Sonnet 4.6) and + # unchunked (Claude Opus 4.6); the Responses API unchunked (GPT-5.4) and + # three outputs per request (GPT-5.5). + model_id=st.sampled_from( + ["claude-sonnet-4-6", "claude-opus-4-6", "gpt-5.4", "gpt-5.5"] + ), + output_count=st.integers(min_value=1, max_value=4), household_count=st.integers(min_value=1, max_value=2), ) def test_rows_record_each_responses_provenance_and_agree_with_the_ledger( @@ -247,18 +353,26 @@ def test_rows_record_each_responses_provenance_and_agree_with_the_ledger( ): """Drive run_no_tools_eval / run_no_tools_single_output_eval through the real request path. Each mocked response's cache flag, duration and shape - (a full answer, values only, nothing, or a length cut) is drawn per call, - so households mix replays with live calls across chunks (Claude Sonnet - 4.6 answers one output per request), repair rounds, budget escalations - and explanation repairs.""" + is drawn per call: an answer to the requested tool, values only, nothing, + a length cut, or a failure. Households mix replays with live calls across + chunks, repair rounds, budget escalations, explanation repairs, retries, + and partial rounds, failed chunks and failed households. + + Oracle: the mock's own call history. The spend ledger must record every + physical request it saw. A group none of whose requests failed must count + every response it returned, with their times. A failure can drop returned + responses from the row (a chunk or household that raised keeps none of its + time), so there the row must equal some subset of the returned responses, + counts and times together.""" programs = US_VARIABLES[:output_count] scenarios = [ _scenario(f"scenario_{index:03d}", 40_000.0 + index) for index in range(household_count) ] + responses_api = eval_no_tools._uses_responses_api(model_id) clock = [1_000.0] current = {"group": None} - calls: list[tuple[tuple, bool, float]] = [] + calls: list[MockCall] = [] real_run_single = eval_no_tools.run_single_no_tools def tracking_run_single(scenario, variable, *args, **kwargs): @@ -266,16 +380,37 @@ def tracking_run_single(scenario, variable, *args, **kwargs): current["group"] = key return real_run_single(scenario, variable, *args, **kwargs) - def fake_completion(**_kwargs): - mode = data.draw(st.sampled_from(RESPONSE_MODES), label="mode") - cache_hit = data.draw(st.booleans(), label="cache_hit") + def fake_request(**kwargs): + explanation_request = "submit_explanations" in json.dumps( + kwargs.get("tools"), default=str + ) + modes = EXPLANATION_MODES if explanation_request else ANSWER_MODES + mode = data.draw(st.sampled_from(modes), label="mode") + cache_hit = data.draw(st.booleans(), label="cache_hit") and mode != "raise" seconds = data.draw( st.floats(min_value=0, max_value=900, allow_nan=False), label="seconds", ) clock[0] += seconds - calls.append((current["group"], cache_hit, seconds)) - return _mock_response(mode, cache_hit) + calls.append( + MockCall( + group=current["group"], + cache_hit=cache_hit, + seconds=seconds, + raised=mode == "raise", + explanation_request=explanation_request, + ) + ) + if mode == "raise": + raise ValueError(REQUEST_FAILURE) + if explanation_request and mode == "answer": + event("an explanation repair was answered") + return _mock_response( + mode, + cache_hit, + explanation_request=explanation_request, + responses_api=responses_api, + ) fake_time = SimpleNamespace( perf_counter=lambda: clock[0], @@ -283,13 +418,12 @@ def fake_completion(**_kwargs): sleep=lambda _seconds: None, ) runner = ( - run_no_tools_eval - if writer == "household" - else (run_no_tools_single_output_eval) + run_no_tools_eval if writer == "household" else run_no_tools_single_output_eval ) with ( tempfile.TemporaryDirectory() as tmp, - patch.object(eval_no_tools, "completion", side_effect=fake_completion), + patch.object(eval_no_tools, "completion", side_effect=fake_request), + patch.object(eval_no_tools, "responses", side_effect=fake_request), patch.object(eval_no_tools, "time", fake_time), patch.object(eval_no_tools, "run_single_no_tools", tracking_run_single), ): @@ -306,63 +440,79 @@ def fake_completion(**_kwargs): assert list(frame.columns) == list(NO_TOOLS_RESULT_COLUMNS) _assert_row_invariants(frame) - group_columns = ( - ["scenario_id"] - if writer == "household" - else [ - "scenario_id", - "variable", - ] - ) + group_columns = ["scenario_id"] + if writer == "single_output": + group_columns.append("variable") for key, rows in frame.groupby(group_columns): - group_calls = [call for call in calls if call[0] == tuple(key)] - hits = sum(hit for _, hit, _ in group_calls) - # One request group per household (or per output): every row - # repeats its counts. - assert rows["live_calls"].nunique() == 1 - assert rows["cache_hit_calls"].nunique() == 1 - assert int(rows["cache_hit_calls"].iloc[0]) == hits - assert int(rows["live_calls"].iloc[0]) == len(group_calls) - hits - live_seconds = sum(seconds for _, hit, seconds in group_calls if not hit) - assert rows["live_elapsed_seconds"].sum() == pytest.approx( - live_seconds, rel=1e-9, abs=1e-6 - ) + group_calls = [call for call in calls if call.group == tuple(key)] + returned = [call for call in group_calls if not call.raised] + failed = len(group_calls) - len(returned) + hits = sum(call.cache_hit for call in returned) - # Differential: the spend ledger records the same calls - # independently, one record per physical request. + # The spend ledger records every physical request the mock saw: + # answered ones with their cache flag, failed ones as errors. records = [ record for record in ledger if record["scenario_id"] == key[0] and (writer == "household" or record["variables"] == [key[1]]) - and record["status"] in ("ok", "parse_error") ] - assert len(records) == len(group_calls) - assert sum(record["cache_hit"] for record in records) == hits - assert sum( - record["elapsed_seconds"] + answered = [ + record for record in records - if not record["cache_hit"] - ) == pytest.approx(rows["live_elapsed_seconds"].sum(), rel=1e-9, abs=1e-6) - assert sum(record["elapsed_seconds"] for record in records) == ( - pytest.approx(rows["elapsed_seconds"].sum(), rel=1e-9, abs=1e-6) + if record["status"] in ("ok", "parse_error") + ] + assert len(answered) == len(returned) + assert sum(record["cache_hit"] for record in answered) == hits + assert sum(record["status"] == "provider_error" for record in records) == ( + failed ) - # The analysis classifier reads the same facts back from the rows. + # One request group per household (or per output): every row + # repeats its counts, and its live time sums like elapsed time. + assert rows["live_calls"].nunique() == 1 + assert rows["cache_hit_calls"].nunique() == 1 + row_hits = int(rows["cache_hit_calls"].iloc[0]) + row_live = int(rows["live_calls"].iloc[0]) + row_live_seconds = rows["live_elapsed_seconds"].sum(min_count=1) + row_seconds = rows["elapsed_seconds"].sum(min_count=1) + if not failed: + assert row_hits == hits + assert row_live == len(returned) - hits + live_seconds = sum(c.seconds for c in returned if not c.cache_hit) + assert row_live_seconds == pytest.approx(live_seconds, **SECONDS) + assert row_seconds == pytest.approx( + sum(call.seconds for call in returned), **SECONDS + ) + # Differential: the ledger's own elapsed times for the same + # calls give the same sums. + assert sum( + record["elapsed_seconds"] + for record in answered + if not record["cache_hit"] + ) == pytest.approx(row_live_seconds, **SECONDS) + assert sum(record["elapsed_seconds"] for record in answered) == ( + pytest.approx(row_seconds, **SECONDS) + ) + else: + event("a request in the group failed") + assert _row_matches_a_subset( + returned, row_hits, row_live, row_live_seconds, row_seconds + ) + + # The household classifier reads the same facts back from the rows. status = household_call_provenance(frame)["status"] for scenario in scenarios: - household_calls = [call for call in calls if call[0][0] == scenario.id] - hits = sum(hit for _, hit, _ in household_calls) - expected = ( - "full_replay" - if hits == len(household_calls) - else "live" - if hits == 0 - else "partial_replay" - ) + household_calls = [c for c in calls if c.group[0] == scenario.id] + if any(call.raised for call in household_calls): + rows = frame[frame["scenario_id"] == scenario.id] + expected = _brute_force_status(rows.to_dict("records")) + else: + expected = _expected_status(household_calls) assert status[("Model", scenario.id)] == expected event(f"household {expected}") event(f"household calls {min(len(household_calls), 6)}+") + event("responses api" if responses_api else "chat completions") # The CSV carries the same values, and the run store returns it byte # for byte with the counts queryable from usage_json. @@ -388,6 +538,93 @@ def fake_completion(**_kwargs): assert usage["live_calls"] in set(expected_rows["live_calls"]) +@pytest.mark.parametrize("repair_replayed", [True, False]) +def test_an_explanation_repair_records_its_own_provenance(repair_replayed): + """Three live answers carry values without explanations (the initial + request and both repair rounds), so the household asks for explanations + alone; that repair is a replay or a live call, and counts as one.""" + clock = [0.0] + requests = [] + + def fake_completion(**kwargs): + explanation_request = "submit_explanations" in json.dumps( + kwargs.get("tools"), default=str + ) + requests.append(explanation_request) + cache_hit = explanation_request and repair_replayed + seconds = 0.02 if cache_hit else 7.0 + clock[0] += seconds + return _mock_response( + "answer" if explanation_request else "values_only", + cache_hit, + explanation_request=explanation_request, + responses_api=False, + ) + + fake_time = SimpleNamespace( + perf_counter=lambda: clock[0], time=time.time, sleep=lambda _s: None + ) + with ( + patch.object(eval_no_tools, "completion", side_effect=fake_completion), + patch.object(eval_no_tools, "time", fake_time), + ): + result = eval_no_tools.run_single_no_tools( + _scenario("scenario_000"), + ["payroll_tax"], + "claude-opus-4-6", + include_explanations=True, + ) + + assert requests == [False, False, False, True] + assert result["explanations"]["payroll_tax"] is not None + assert result["live_calls"] == (3 if repair_replayed else 4) + assert result["cache_hit_calls"] == (1 if repair_replayed else 0) + assert result["live_elapsed_seconds"] == pytest.approx( + 21.0 if repair_replayed else 28.0 + ) + assert result["elapsed_seconds"] == pytest.approx( + 21.02 if repair_replayed else 28.0 + ) + assert result["spend_ledger"][-1]["phase"] == "explanation_repair" + assert result["spend_ledger"][-1]["cache_hit"] is repair_replayed + + +def test_responses_api_requests_record_replays_and_live_calls(): + """A Responses API household whose first answer is a replay with nothing + usable, then a live repair: one call of each kind, and only the repair's + time is live.""" + clock = [50.0] + script = iter([(True, 0.03, "empty"), (False, 41.0, "answer")]) + + def fake_responses(**_kwargs): + cache_hit, seconds, mode = next(script) + clock[0] += seconds + return _mock_response( + mode, cache_hit, explanation_request=False, responses_api=True + ) + + fake_time = SimpleNamespace( + perf_counter=lambda: clock[0], time=time.time, sleep=lambda _s: None + ) + with ( + patch.object(eval_no_tools, "responses", side_effect=fake_responses), + patch.object(eval_no_tools, "time", fake_time), + ): + result = eval_no_tools.run_single_no_tools( + _scenario("scenario_000"), + ["payroll_tax"], + "gpt-5.4", + include_explanations=True, + ) + + assert result["predictions"]["payroll_tax"] is not None + assert result["cache_hit_calls"] == 1 + assert result["live_calls"] == 1 + assert result["live_elapsed_seconds"] == pytest.approx(41.0) + assert result["elapsed_seconds"] == pytest.approx(41.03) + assert [record["cache_hit"] for record in result["spend_ledger"]] == [True, False] + + def test_a_household_that_raised_records_no_responses(monkeypatch, tmp_path): def raising_run_single(*_args, **_kwargs): raise ValueError("malformed provider payload") @@ -666,6 +903,61 @@ def test_row_repair_merge_fills_only_the_repaired_row(tmp_path): assert set(household_call_provenance(merged.reset_index())["status"]) == {"unknown"} +def test_a_repair_that_predates_the_columns_leaves_the_row_unknown(tmp_path): + """An older repair file has no provenance. The row it replaces must not + keep the source's values, which described another request: beside the + repair's 2 seconds they would claim 10 live seconds.""" + source_path = tmp_path / "source.csv" + repair_path = tmp_path / "repair.csv" + pd.DataFrame( + { + "model": ["model_a", "model_a"], + "scenario_id": ["s1", "s1"], + "variable": ["tax", "benefit"], + "prediction": [1.0, None], + "explanation": ["Tax. value = 1", None], + "raw_response": [None, None], + "error": [None, "Missing prediction after repair: benefit"], + "elapsed_seconds": [10.0, 10.0], + "live_calls": [1, 1], + "cache_hit_calls": [0, 0], + "live_elapsed_seconds": [10.0, 10.0], + } + ).to_csv(source_path, index=False) + pd.DataFrame( + { + "country": ["us"], + "attempt": [1], + "model": ["model_a"], + "scenario_id": ["s1"], + "variable": ["benefit"], + "prediction": [22.0], + "explanation": ["Repair benefit. value = 22"], + "raw_response": ["raw"], + "error": [None], + "elapsed_seconds": [2.0], + } + ).to_csv(repair_path, index=False) + + outputs = merge_row_repair_predictions( + source_predictions=source_path, + repair_predictions=repair_path, + output_dir=tmp_path / "merged", + ) + + merged = pd.read_csv(outputs["merged_predictions"]).set_index("variable") + assert merged.loc["benefit", "elapsed_seconds"] == 2.0 + assert merged.loc["benefit", PROVENANCE].isna().all() + assert merged.loc["tax", "live_calls"] == 1 + assert merged.loc["tax", "live_elapsed_seconds"] == 10.0 + timed = merged["live_elapsed_seconds"].notna() + assert ( + merged.loc[timed, "live_elapsed_seconds"] + <= merged.loc[timed, "elapsed_seconds"] + ).all() + assert set(household_call_provenance(merged.reset_index())["status"]) == {"unknown"} + + def test_retry_merge_carries_the_retry_rows_provenance(tmp_path): source_path = tmp_path / "source.csv" retry_path = tmp_path / "retry.csv"