diff --git a/policybench/annotation_taxonomy.py b/policybench/annotation_taxonomy.py index 5315e0c..58cd172 100644 --- a/policybench/annotation_taxonomy.py +++ b/policybench/annotation_taxonomy.py @@ -10,6 +10,7 @@ "reference_model_issue_fixed", "reference_data_issue_fixed", "parse_contract_failure", + "budget_exhausted_at_ceiling", "needs_review", ) diff --git a/policybench/annotation_validation.py b/policybench/annotation_validation.py index 5608fe8..f17b786 100644 --- a/policybench/annotation_validation.py +++ b/policybench/annotation_validation.py @@ -24,7 +24,11 @@ # labels these without a classifier). Everything else — reference-suspect # classes, prompt ambiguity, needs_review — is unresolved for a frozen # snapshot. -FINAL_FAILURE_SOURCES = {"llm_error", "parse_contract_failure"} +FINAL_FAILURE_SOURCES = { + "llm_error", + "parse_contract_failure", + "budget_exhausted_at_ceiling", +} def _expected_prediction_rows( diff --git a/policybench/audit.py b/policybench/audit.py index 56a768f..b53e111 100644 --- a/policybench/audit.py +++ b/policybench/audit.py @@ -49,6 +49,7 @@ class WrongModel: model: str prediction: str explanation: str + failure_source: str = "" @dataclass(frozen=True) @@ -68,6 +69,11 @@ class AuditCase: def to_manifest_row(self) -> dict: row = asdict(self) + row["recorded_failure_sources"] = { + model.model: model.failure_source + for model in self.wrong_models + if model.failure_source == "budget_exhausted_at_ceiling" + } row["wrong_models"] = [m.model for m in self.wrong_models] missing = [m.model for m in self.wrong_models if m.prediction == MISSING_VALUE] row["missing_models"] = missing @@ -127,6 +133,7 @@ def build_audit_cases( r["prediction"], country=country, variable=variable ), explanation=_clean(r.get("explanation")), + failure_source=_clean(r.get("failure_source")), ) for _, r in deduped.iterrows() ) @@ -315,6 +322,8 @@ def _scenario_questions( reference_suspect=true and a concrete contradiction. - parse_contract_failure: the model's answer was missing or unparseable, not a \ substantive error. +- budget_exhausted_at_ceiling: the provider length-terminated every retry through \ +its maximum allowed completion budget. - needs_review: genuinely cannot tell; name precisely what information is \ missing. @@ -486,7 +495,14 @@ def parse_verdict(path: Path) -> dict | None: def _row_failure_source(meta: dict, model: str, requested_source: str) -> str: """Keep parse-contract labels scoped to genuinely missing predictions.""" + recorded_source = meta.get("recorded_failure_sources", {}).get(model) + if recorded_source == "budget_exhausted_at_ceiling": + return recorded_source source = validate_failure_source(requested_source) + if source == "budget_exhausted_at_ceiling": + if model in set(meta.get("missing_models", [])): + return "parse_contract_failure" + return "llm_error" if source == "parse_contract_failure" and model not in set( meta.get("missing_models", []) ): @@ -520,14 +536,20 @@ def collect_audit(country_dir: Path, audit_dir: Path) -> dict[str, pd.DataFrame] # Every wrong model simply returned no value — a parse failure, not # a substantive error. Classified deterministically, no classifier. note = "All wrong responses were missing or unparseable predictions." + row_sources = [] for model in meta["wrong_models"]: + failure_source = meta.get("recorded_failure_sources", {}).get( + model, + "parse_contract_failure", + ) + row_sources.append(failure_source) row_records.append( { "country": country, "scenario_id": meta["scenario_id"], "variable": meta["variable"], "model": model, - "failure_source": "parse_contract_failure", + "failure_source": failure_source, "failure_subtype": "missing_output", "reference_suspect": False, "annotation": note, @@ -539,7 +561,11 @@ def collect_audit(country_dir: Path, audit_dir: Path) -> dict[str, pd.DataFrame] "scenario_id": meta["scenario_id"], "variable": meta["variable"], "wrong_model_count": len(meta["wrong_models"]), - "case_failure_source": "parse_contract_failure", + "case_failure_source": ( + row_sources[0] + if len(set(row_sources)) == 1 + else "parse_contract_failure" + ), "case_failure_subtype": "missing_output", "reference_suspect": False, "reference_bug_hypothesis": "", @@ -552,6 +578,18 @@ def collect_audit(country_dir: Path, audit_dir: Path) -> dict[str, pd.DataFrame] missing.append(case_id) continue case_source = validate_failure_source(verdict["case_failure_source"]) + recorded_sources = set(meta.get("recorded_failure_sources", {}).values()) + if ( + case_source == "budget_exhausted_at_ceiling" + and "budget_exhausted_at_ceiling" not in recorded_sources + ): + wrong_models = set(meta.get("wrong_models", [])) + missing_models = set(meta.get("missing_models", [])) + case_source = ( + "parse_contract_failure" + if wrong_models and wrong_models <= missing_models + else "llm_error" + ) case_subtype = validate_failure_subtype(verdict["case_failure_subtype"]) reference_suspect = bool(verdict.get("reference_suspect")) rationale = str(verdict.get("rationale", "")).strip() diff --git a/policybench/case_annotations.py b/policybench/case_annotations.py index a6ab13e..47cf030 100644 --- a/policybench/case_annotations.py +++ b/policybench/case_annotations.py @@ -86,10 +86,15 @@ def _expected_prediction_rows( prediction_columns = ["model", "scenario_id", "variable", "prediction"] optional_columns = [ - column for column in ["explanation", "error"] if column in predictions.columns + column + for column in ["explanation", "error", "failure_source"] + if column in predictions.columns ] prediction_details = predictions[prediction_columns + optional_columns].rename( - columns={"error": "prediction_error"} + columns={ + "error": "prediction_error", + "failure_source": "recorded_failure_source", + } ) merged = expected.merge( prediction_details, @@ -139,6 +144,27 @@ def wrong_prediction_rows(country_dir: Path) -> pd.DataFrame: for column in ["failure_source", "failure_subtype"]: if column not in wrong.columns: wrong[column] = pd.NA + recorded_ceiling_exhaustion = pd.Series(False, index=wrong.index) + if "recorded_failure_source" in wrong.columns: + recorded_ceiling_exhaustion = ( + wrong["recorded_failure_source"].astype("string") + == "budget_exhausted_at_ceiling" + ).fillna(False) + wrong.loc[ + recorded_ceiling_exhaustion, + "failure_source", + ] = "budget_exhausted_at_ceiling" + invented_ceiling_exhaustion = ( + wrong["failure_source"].astype("string") == "budget_exhausted_at_ceiling" + ).fillna(False) & ~recorded_ceiling_exhaustion + wrong.loc[ + invented_ceiling_exhaustion & wrong["prediction"].isna(), + "failure_source", + ] = "parse_contract_failure" + wrong.loc[ + invented_ceiling_exhaustion & wrong["prediction"].notna(), + "failure_source", + ] = "llm_error" missing_categories = ( wrong["annotation"].astype("string").fillna("").str.strip() != "" ) & ( @@ -149,12 +175,30 @@ def wrong_prediction_rows(country_dir: Path) -> pd.DataFrame: inferred = wrong.loc[missing_categories, "annotation"].map( infer_failure_category ) - wrong.loc[missing_categories, "failure_source"] = [ - category.failure_source for category in inferred - ] - wrong.loc[missing_categories, "failure_subtype"] = [ - category.failure_subtype for category in inferred - ] + inferred_sources = pd.Series( + [category.failure_source for category in inferred], + index=inferred.index, + ) + inferred_subtypes = pd.Series( + [category.failure_subtype for category in inferred], + index=inferred.index, + ) + missing_source = ( + wrong["failure_source"].astype("string").fillna("").str.strip() == "" + ) + missing_subtype = ( + wrong["failure_subtype"].astype("string").fillna("").str.strip() == "" + ) + source_rows = missing_categories & missing_source + subtype_rows = missing_categories & missing_subtype + source_index = wrong.index[source_rows] + subtype_index = wrong.index[subtype_rows] + wrong.loc[source_index, "failure_source"] = inferred_sources.reindex( + source_index + ) + wrong.loc[subtype_index, "failure_subtype"] = inferred_subtypes.reindex( + subtype_index + ) return wrong diff --git a/policybench/completion_budget.py b/policybench/completion_budget.py new file mode 100644 index 0000000..f85dd2a --- /dev/null +++ b/policybench/completion_budget.py @@ -0,0 +1,39 @@ +"""Helpers for provider-specific completion-budget request fields.""" + +from __future__ import annotations + +# Length-terminated requests may double beyond their normal serving budget. +# A model card can lower this default when its provider has a documented cap. +MAX_ESCALATED_COMPLETION_TOKENS = 128_000 + +COMPLETION_BUDGET_KEYS = ( + "max_completion_tokens", + "max_output_tokens", + "max_tokens", +) + + +def completion_budget_from_kwargs(request_kwargs: dict) -> int | None: + """Return the active completion budget from a provider request.""" + for key in COMPLETION_BUDGET_KEYS: + value = request_kwargs.get(key) + if value is not None: + return int(value) + return None + + +def with_completion_budget(request_kwargs: dict, budget: int) -> dict: + """Copy a request and replace its provider-specific completion budget.""" + updated = dict(request_kwargs) + for key in COMPLETION_BUDGET_KEYS: + if key in updated: + updated[key] = int(budget) + return updated + raise ValueError("Request has no recognized completion-budget field") + + +def next_completion_budget(current: int, ceiling: int) -> int | None: + """Return the next doubled budget, capped at ``ceiling``.""" + if current >= ceiling: + return None + return min(current * 2, ceiling) diff --git a/policybench/eval_no_tools.py b/policybench/eval_no_tools.py index a4ad657..f665288 100644 --- a/policybench/eval_no_tools.py +++ b/policybench/eval_no_tools.py @@ -17,6 +17,12 @@ import pandas as pd from litellm import completion, responses +from policybench.completion_budget import ( + MAX_ESCALATED_COMPLETION_TOKENS, + completion_budget_from_kwargs, + next_completion_budget, + with_completion_budget, +) from policybench.config import ( GPT_56_MODELS, MODELS, @@ -27,6 +33,7 @@ PROMPT_CONTRACT_VERSION, answer_contract_for, card_for, + completion_budget_ceiling_for, explanation_chunk_size_for, ) from policybench.policyengine_runtime import policyengine_bundles_for_countries @@ -38,7 +45,12 @@ ) from policybench.scenarios import Scenario, scenario_to_dict from policybench.spec import expand_programs_for_scenario, metric_type_for_output -from policybench.spend_ledger import spend_ledger_path, upsert_spend_ledger +from policybench.spend_ledger import ( + count_budget_escalations, + read_spend_ledger, + spend_ledger_path, + upsert_spend_ledger, +) logger = logging.getLogger(__name__) @@ -149,7 +161,7 @@ def _env_int(name: str, default: int) -> int: REQUEST_WALL_TIMEOUT_MULTIPLIER = 1.5 CHECKPOINT_EVERY_ROWS = 25 MAX_REPAIR_ROUNDS = _env_int("POLICYBENCH_MAX_REPAIR_ROUNDS", 2) -RESUME_METADATA_VERSION = 2 +RESUME_METADATA_VERSION = 3 DEFAULT_MAX_COMPLETION_TOKENS = 64 EXTENDED_MAX_COMPLETION_TOKENS = 256 EXPLANATION_MAX_COMPLETION_TOKENS = 4096 @@ -518,6 +530,8 @@ def _spend_ledger_record( usage: dict | None = None, error: str | None = None, cache_hit: bool = False, + completion_budget_tokens: int | None = None, + escalated_from_budget_tokens: int | None = None, ) -> dict: """Build one audit row for a physical synchronous provider request.""" usage = usage or {} @@ -548,6 +562,8 @@ def _spend_ledger_record( "error": error, "cache_hit": cache_hit, "cached_response_cost_usd": cached_response_cost_usd, + "completion_budget_tokens": completion_budget_tokens, + "escalated_from_budget_tokens": escalated_from_budget_tokens, **usage_fields, } @@ -629,7 +645,7 @@ def _completion_token_budget( return min(max(base_tokens, dynamic_tokens), cap) -def _completion_controls( +def _uncapped_completion_controls( model_id: str, include_explanations: bool = True, variables: list[str] | None = None, @@ -725,6 +741,29 @@ def _completion_controls( } +def _completion_budget_ceiling(model_id: str) -> int: + """Return the largest completion budget PolicyBench may request.""" + return completion_budget_ceiling_for(model_id) + + +def _completion_controls( + model_id: str, + include_explanations: bool = True, + variables: list[str] | None = None, +) -> dict: + """Return initial controls capped by a documented provider maximum.""" + controls = _uncapped_completion_controls( + model_id, + include_explanations=include_explanations, + variables=variables, + ) + budget = completion_budget_from_kwargs(controls) + ceiling = _completion_budget_ceiling(model_id) + if budget is not None and budget > ceiling: + return with_completion_budget(controls, ceiling) + return controls + + def _request_timeout_seconds(model_id: str) -> int: card = card_for(model_id) if card is not None and card.request_timeout_seconds is not None: @@ -1191,13 +1230,16 @@ def _empty_failed_result( variables: list[str], error: Exception, ) -> dict: + spend_records = _spend_records(error) return { "predictions": {variable: None for variable in variables}, "explanations": {variable: None for variable in variables}, "prediction": None, "error": _format_error(error), "repair_disagreements": [], - "spend_ledger": _spend_records(error), + "failure_sources": {}, + "budget_escalation_count": count_budget_escalations(spend_records), + "spend_ledger": spend_records, **_aggregate_request_results([]), } @@ -1612,6 +1654,61 @@ def _responses_content_and_tool_calls(response) -> tuple[str | None, list[dict]] return content, tool_calls +def _response_finish_reason( + response, + *, + responses_api: bool, + completion_budget_tokens: int | None = None, +) -> str | None: + """Normalize provider completion-limit signals to ``length``.""" + if not responses_api: + choices = getattr(response, "choices", None) or [] + if not choices: + return None + reason = _get_response_item_attr(choices[0], "finish_reason") + return str(reason) if isinstance(reason, str) else None + + reason = _get_response_item_attr(response, "finish_reason") + if reason == "length": + return "length" + status = _get_response_item_attr(response, "status") + details = _get_response_item_attr(response, "incomplete_details") + incomplete_reason = _get_response_item_attr(details, "reason") + if status == "incomplete" and incomplete_reason in { + "max_output_tokens", + "max_tokens", + "length", + }: + return "length" + if ( + status == "incomplete" + and incomplete_reason is None + and completion_budget_tokens is not None + ): + usage = _get_response_item_attr(response, "usage") + output_tokens = _get_usage_value(usage, "completion_tokens") + if output_tokens is None: + output_tokens = _get_usage_value(usage, "output_tokens") + if ( + isinstance(output_tokens, (int, float)) + and not isinstance(output_tokens, bool) + and output_tokens >= completion_budget_tokens + ): + # LiteLLM's chat-to-Responses adapter currently drops the + # underlying chat ``finish_reason``. A bare incomplete response + # that consumed the requested output budget is the remaining + # observable completion-cap signal; lower-token incomplete + # responses (for example filters/refusals) are not escalated. + return "length" + if status == "length": + return "length" + if status == "completed": + return "stop" + if isinstance(reason, str): + return reason + return str(status) if isinstance(status, str) else None + + def extract_predictions( content: str | None, variables: list[str], @@ -1712,6 +1809,9 @@ def _request_explanations_once( answers: dict[str, float], model_id: str, spend_callback: Callable[[list[dict]], None] | None = None, + *, + completion_budget_tokens: int | None = None, + escalated_from_budget_tokens: int | None = None, ) -> dict: answer_contract = _answer_contract_for_model(model_id) prompt = make_explanation_repair_prompt( @@ -1779,6 +1879,12 @@ def _request_explanations_once( request_kwargs["response_format"] = {"type": "json_object"} request_fn = completion + if completion_budget_tokens is not None: + request_kwargs = with_completion_budget( + request_kwargs, + completion_budget_tokens, + ) + request_completion_budget = completion_budget_from_kwargs(request_kwargs) call_key = f"sync:{uuid4().hex}" request_started_at = time.time() started_at = time.perf_counter() @@ -1795,6 +1901,11 @@ def _request_explanations_once( content = getattr(message, "content", None) tool_calls = getattr(message, "tool_calls", None) function_call = getattr(message, "function_call", None) + finish_reason = _response_finish_reason( + response, + responses_api=_uses_responses_api(model_id), + completion_budget_tokens=request_completion_budget, + ) raw_response = _serialize_response_payload( content=content, tool_calls=tool_calls, @@ -1836,6 +1947,8 @@ def _request_explanations_once( elapsed_seconds=elapsed_seconds, usage=usage, cache_hit=_response_cache_hit(response), + completion_budget_tokens=request_completion_budget, + escalated_from_budget_tokens=escalated_from_budget_tokens, ) _emit_spend_records(spend_callback, [spend_record]) return { @@ -1844,6 +1957,8 @@ def _request_explanations_once( "elapsed_seconds": elapsed_seconds, "request_started_at": request_started_at, "request_completed_at": request_completed_at, + "finish_reason": finish_reason, + "completion_budget": request_completion_budget, **usage, "spend_ledger": [spend_record], } @@ -1864,12 +1979,111 @@ def _request_explanations_once( elapsed_seconds=time.perf_counter() - started_at, usage=usage, error=_format_error(error), + completion_budget_tokens=request_completion_budget, + escalated_from_budget_tokens=escalated_from_budget_tokens, ) _emit_spend_records(spend_callback, [spend_record]) _attach_spend_records(error, [spend_record]) raise +def _request_explanations_with_budget_escalation( + scenario: Scenario, + variables: list[str], + answers: dict[str, float], + model_id: str, + *, + base_explanations: dict[str, str | None], + repair_disagreements: list[dict], + spend_callback: Callable[[list[dict]], None] | None, +) -> dict: + """Repair explanations, escalating length-truncated missing cells.""" + explanations = dict(base_explanations) + request_results: list[dict] = [] + exhausted_variables: set[str] = set() + escalation_count = 0 + request_variables = list(variables) + budget_override: int | None = None + escalated_from: int | None = None + + while True: + try: + result = _request_explanations_once( + scenario, + request_variables, + {variable: answers[variable] for variable in request_variables}, + model_id, + spend_callback=spend_callback, + completion_budget_tokens=budget_override, + escalated_from_budget_tokens=escalated_from, + ) + except Exception as error: + prior_spend = [ + record + for prior_result in request_results + for record in _spend_records(prior_result) + ] + if prior_spend: + _attach_spend_records(error, prior_spend) + error._policybench_partial_explanation_round = { + "explanations": explanations, + "request_results": request_results, + "exhausted_variables": exhausted_variables, + "budget_escalation_count": escalation_count, + } + raise + + request_results.append(result) + explanations = _merge_repair_cells( + explanations, + result.get("explanations", {}), + field="explanation", + disagreements=repair_disagreements, + ) + remaining = [ + variable + for variable in request_variables + if not str(explanations.get(variable) or "").strip() + ] + if result.get("finish_reason") != "length" or not remaining: + break + + current_budget = result.get("completion_budget") + if current_budget is None: + break + next_budget = next_completion_budget( + int(current_budget), + _completion_budget_ceiling(model_id), + ) + if next_budget is None: + exhausted_variables.update(remaining) + break + + logger.warning( + "Escalating completion budget: model=%s scenario=%s " + "from_budget=%s to_budget=%s", + model_id, + scenario.id, + current_budget, + next_budget, + ) + escalation_count += 1 + if any( + str(result["explanations"].get(variable) or "").strip() + for variable in request_variables + ): + request_variables = remaining + budget_override = next_budget + escalated_from = int(current_budget) + + return { + "explanations": explanations, + "request_results": request_results, + "exhausted_variables": exhausted_variables, + "budget_escalation_count": escalation_count, + } + + def _request_predictions_once( scenario: Scenario, variables: list[str], @@ -1878,6 +2092,8 @@ def _request_predictions_once( repair: bool = False, include_explanations: bool = True, spend_callback: Callable[[list[dict]], None] | None = None, + completion_budget_tokens: int | None = None, + escalated_from_budget_tokens: int | None = None, ) -> dict: if _uses_responses_api(model_id): messages, request_kwargs = _responses_request_kwargs( @@ -1898,6 +2114,12 @@ def _request_predictions_once( ) request_fn = completion + if completion_budget_tokens is not None: + request_kwargs = with_completion_budget( + request_kwargs, + completion_budget_tokens, + ) + request_completion_budget = completion_budget_from_kwargs(request_kwargs) spend_records = [] phase = "repair" if repair else "initial" for attempt in range(MAX_ATTEMPTS): @@ -1917,6 +2139,11 @@ def _request_predictions_once( content = getattr(message, "content", None) tool_calls = getattr(message, "tool_calls", None) function_call = getattr(message, "function_call", None) + finish_reason = _response_finish_reason( + response, + responses_api=_uses_responses_api(model_id), + completion_budget_tokens=request_completion_budget, + ) raw_response = _serialize_response_payload( content=content, tool_calls=tool_calls, @@ -1968,6 +2195,10 @@ def _request_predictions_once( elapsed_seconds=elapsed_seconds, usage=usage, cache_hit=_response_cache_hit(response), + completion_budget_tokens=request_completion_budget, + escalated_from_budget_tokens=( + escalated_from_budget_tokens if attempt == 0 else None + ), ) ) _emit_spend_records(spend_callback, [spend_records[-1]]) @@ -1978,6 +2209,11 @@ def _request_predictions_once( "elapsed_seconds": elapsed_seconds, "request_started_at": request_started_at, "request_completed_at": request_completed_at, + "finish_reason": finish_reason, + "completion_budget": request_completion_budget, + "budget_escalation_count": int( + escalated_from_budget_tokens is not None + ), **usage, "spend_ledger": spend_records, } @@ -2000,6 +2236,10 @@ def _request_predictions_once( elapsed_seconds=elapsed_seconds, usage=usage, error=_format_error(e), + completion_budget_tokens=request_completion_budget, + escalated_from_budget_tokens=( + escalated_from_budget_tokens if attempt == 0 else None + ), ) ) _emit_spend_records(spend_callback, [spend_records[-1]]) @@ -2013,6 +2253,127 @@ def _request_predictions_once( raise RuntimeError("Request loop exited unexpectedly") +def _request_predictions_with_budget_escalation( + scenario: Scenario, + variables: list[str], + model_id: str, + *, + repair: bool, + include_explanations: bool, + base_predictions: dict[str, float | None], + base_explanations: dict[str, str | None], + repair_disagreements: list[dict], + spend_callback: Callable[[list[dict]], None] | None, +) -> dict: + """Run one logical request, escalating only length-truncated misses.""" + predictions = dict(base_predictions) + explanations = dict(base_explanations) + request_results: list[dict] = [] + exhausted_variables: set[str] = set() + escalation_count = 0 + request_variables = list(variables) + request_is_repair = repair + budget_override: int | None = None + escalated_from: int | None = None + + while True: + try: + result = _request_predictions_once( + scenario, + request_variables, + model_id, + repair=request_is_repair, + include_explanations=include_explanations, + spend_callback=spend_callback, + completion_budget_tokens=budget_override, + escalated_from_budget_tokens=escalated_from, + ) + except Exception as error: + prior_spend = [ + record + for prior_result in request_results + for record in _spend_records(prior_result) + ] + if prior_spend: + _attach_spend_records(error, prior_spend) + error._policybench_partial_budget_round = { + "predictions": predictions, + "explanations": explanations, + "request_results": request_results, + "exhausted_variables": exhausted_variables, + "budget_escalation_count": escalation_count, + } + raise + + request_results.append(result) + predictions = _merge_repair_cells( + predictions, + result["predictions"], + field="prediction", + disagreements=repair_disagreements, + ) + explanations = _merge_repair_cells( + explanations, + result.get("explanations", {}), + field="explanation", + disagreements=repair_disagreements, + ) + + remaining = [ + variable + for variable in request_variables + if predictions.get(variable) is None + or ( + include_explanations + and not str(explanations.get(variable) or "").strip() + ) + ] + if result.get("finish_reason") != "length" or not remaining: + break + + current_budget = result.get("completion_budget") + if current_budget is None: + break + next_budget = next_completion_budget( + int(current_budget), + _completion_budget_ceiling(model_id), + ) + if next_budget is None: + exhausted_variables.update(remaining) + break + + logger.warning( + "Escalating completion budget: model=%s scenario=%s " + "from_budget=%s to_budget=%s", + model_id, + scenario.id, + current_budget, + next_budget, + ) + escalation_count += 1 + yielded_valid_cell = any( + result["predictions"].get(variable) is not None + or ( + include_explanations + and bool(str(result["explanations"].get(variable) or "").strip()) + ) + for variable in request_variables + ) + if yielded_valid_cell: + request_variables = remaining + request_is_repair = True + budget_override = next_budget + escalated_from = int(current_budget) + + return { + "predictions": predictions, + "explanations": explanations, + "request_results": request_results, + "exhausted_variables": exhausted_variables, + "budget_escalation_count": escalation_count, + } + + def run_single_no_tools( scenario: Scenario, variable: str | Iterable[str], @@ -2128,62 +2489,119 @@ def run_single_no_tools( for result in chunk_results for disagreement in result.get("repair_disagreements", []) ], + "failure_sources": { + variable: source + for result in chunk_results + for variable, source in result.get("failure_sources", {}).items() + }, + "budget_escalation_count": sum( + int(result.get("budget_escalation_count", 0)) + for result in chunk_results + ), "spend_ledger": [ record for result in chunk_results for record in _spend_records(result) ], } - request_results = [] - initial_result = _request_predictions_once( - scenario, - variables, - model_id, - repair=False, - include_explanations=include_explanations, - spend_callback=_spend_callback, - ) - request_results.append(initial_result) - spend_records = _spend_records(initial_result) - predictions = dict(initial_result["predictions"]) - explanations = dict(initial_result.get("explanations", {})) + request_results: list[dict] = [] + spend_records: list[dict] = [] + predictions = {variable: None for variable in variables} + explanations = {variable: None for variable in variables} + repair_errors = [] + repair_disagreements: list[dict] = [] + exhausted_variables: set[str] = set() + budget_escalation_count = 0 + request_failed = False + + try: + initial_round = _request_predictions_with_budget_escalation( + scenario, + variables, + model_id, + repair=False, + include_explanations=include_explanations, + base_predictions=predictions, + base_explanations=explanations, + repair_disagreements=repair_disagreements, + spend_callback=_spend_callback, + ) + except Exception as error: + partial_round = getattr(error, "_policybench_partial_budget_round", None) + has_valid_partial = partial_round and any( + value is not None for value in partial_round["predictions"].values() + ) + if not has_valid_partial: + raise + initial_round = partial_round + spend_records.extend(_spend_records(error)) + repair_errors.append(_format_error(error)) + request_failed = True + request_results.extend(initial_round["request_results"]) + if not request_failed: + spend_records.extend( + record + for result in initial_round["request_results"] + for record in _spend_records(result) + ) + predictions = initial_round["predictions"] + explanations = initial_round["explanations"] + exhausted_variables.update(initial_round["exhausted_variables"]) + budget_escalation_count += initial_round["budget_escalation_count"] missing = _missing_variables(predictions) missing_explanations = ( _missing_explanations(explanations, variables) if include_explanations else [] ) - repair_errors = [] - repair_disagreements = [] - for _ in range(MAX_REPAIR_ROUNDS): - repair_targets = sorted(set(missing) | set(missing_explanations)) + for _ in range(0 if request_failed else MAX_REPAIR_ROUNDS): + repair_targets = sorted( + (set(missing) | set(missing_explanations)) - exhausted_variables + ) if not repair_targets: break try: - repair_result = _request_predictions_once( + repair_round = _request_predictions_with_budget_escalation( scenario, repair_targets, model_id, repair=True, include_explanations=include_explanations, + base_predictions=predictions, + base_explanations=explanations, + repair_disagreements=repair_disagreements, spend_callback=_spend_callback, ) except Exception as error: spend_records.extend(_spend_records(error)) + request_failed = True + partial_round = getattr( + error, + "_policybench_partial_budget_round", + None, + ) + if partial_round: + request_results.extend(partial_round["request_results"]) + predictions = partial_round["predictions"] + explanations = partial_round["explanations"] + exhausted_variables.update(partial_round["exhausted_variables"]) + budget_escalation_count += partial_round["budget_escalation_count"] + missing = _missing_variables(predictions) + missing_explanations = ( + _missing_explanations(explanations, variables) + if include_explanations + else [] + ) repair_errors.append(_format_error(error)) break - request_results.append(repair_result) - spend_records.extend(_spend_records(repair_result)) - predictions = _merge_repair_cells( - predictions, - repair_result["predictions"], - field="prediction", - disagreements=repair_disagreements, - ) - explanations = _merge_repair_cells( - explanations, - repair_result.get("explanations", {}), - field="explanation", - disagreements=repair_disagreements, + request_results.extend(repair_round["request_results"]) + spend_records.extend( + record + for result in repair_round["request_results"] + for record in _spend_records(result) ) + predictions = repair_round["predictions"] + explanations = repair_round["explanations"] + exhausted_variables.update(repair_round["exhausted_variables"]) + budget_escalation_count += repair_round["budget_escalation_count"] missing = _missing_variables(predictions) missing_explanations = ( _missing_explanations(explanations, variables) @@ -2191,45 +2609,82 @@ def run_single_no_tools( else [] ) - if include_explanations and not missing and missing_explanations: + if ( + include_explanations + and not request_failed + and not missing + and missing_explanations + ): explanation_answers = { variable: predictions[variable] for variable in missing_explanations if predictions.get(variable) is not None + and variable not in exhausted_variables } if explanation_answers: try: - explanation_result = _request_explanations_once( + explanation_round = _request_explanations_with_budget_escalation( scenario, list(explanation_answers), explanation_answers, model_id, + base_explanations=explanations, + repair_disagreements=repair_disagreements, spend_callback=_spend_callback, ) - request_results.append(explanation_result) - spend_records.extend(_spend_records(explanation_result)) - explanations.update( - { - variable: value - for variable, value in explanation_result.get( - "explanations", {} - ).items() - if value is not None - } + request_results.extend(explanation_round["request_results"]) + spend_records.extend( + record + for result in explanation_round["request_results"] + for record in _spend_records(result) ) + explanations = explanation_round["explanations"] + exhausted_variables.update(explanation_round["exhausted_variables"]) + budget_escalation_count += explanation_round["budget_escalation_count"] missing_explanations = _missing_explanations(explanations, variables) except Exception as error: spend_records.extend(_spend_records(error)) + partial_round = getattr( + error, + "_policybench_partial_explanation_round", + None, + ) + if partial_round: + request_results.extend(partial_round["request_results"]) + explanations = partial_round["explanations"] + exhausted_variables.update(partial_round["exhausted_variables"]) + budget_escalation_count += partial_round["budget_escalation_count"] + missing_explanations = _missing_explanations( + explanations, + variables, + ) repair_errors.append(_format_error(error)) - if missing: + exhausted_missing = sorted(set(missing) & exhausted_variables) + ordinary_missing = sorted(set(missing) - exhausted_variables) + if exhausted_missing: repair_errors.append( - "Missing predictions after repair: " + ", ".join(sorted(missing)) + "budget_exhausted_at_ceiling: " + ", ".join(exhausted_missing) ) - if include_explanations and missing_explanations: + if ordinary_missing: + repair_errors.append( + "Missing predictions after repair: " + ", ".join(ordinary_missing) + ) + ordinary_missing_explanations = sorted( + set(missing_explanations) - exhausted_variables + ) + exhausted_missing_explanations = sorted( + (set(missing_explanations) & exhausted_variables) - set(missing) + ) + if include_explanations and exhausted_missing_explanations: + repair_errors.append( + "budget_exhausted_at_ceiling (explanation): " + + ", ".join(exhausted_missing_explanations) + ) + if include_explanations and ordinary_missing_explanations: repair_errors.append( "Missing explanations after repair: " - + ", ".join(sorted(missing_explanations)) + + ", ".join(ordinary_missing_explanations) ) aggregated = _aggregate_request_results(request_results) @@ -2239,6 +2694,10 @@ def run_single_no_tools( "prediction": predictions[variables[0]] if len(variables) == 1 else None, "error": "; ".join(repair_errors) if repair_errors else None, "repair_disagreements": repair_disagreements, + "failure_sources": { + variable: "budget_exhausted_at_ceiling" for variable in exhausted_missing + }, + "budget_escalation_count": budget_escalation_count, "spend_ledger": spend_records, **aggregated, } @@ -2365,6 +2824,14 @@ def _build_resume_metadata( "models": {name: models[name] for name in sorted(models)}, "policyengine_bundles": policyengine_bundles_for_countries(countries), "response_contract": _response_contract_metadata(), + "completion_budget_escalation": { + "strategy": "double_on_length_with_missing_payload", + "default_ceiling": MAX_ESCALATED_COMPLETION_TOKENS, + "model_ceilings": { + name: _completion_budget_ceiling(models[name]) + for name in sorted(models) + }, + }, } @@ -2372,9 +2839,13 @@ def _write_resume_metadata(output_path: str | None, metadata: dict) -> None: metadata_path = _output_metadata_path(output_path) if metadata_path is None: return + runtime_metadata = dict(metadata) + runtime_metadata["budget_escalation_count"] = count_budget_escalations( + read_spend_ledger(spend_ledger_path(output_path)) + ) metadata_path.parent.mkdir(parents=True, exist_ok=True) metadata_path.write_text( - json.dumps(metadata, indent=2, sort_keys=True), + json.dumps(runtime_metadata, indent=2, sort_keys=True), encoding="utf-8", ) @@ -2413,6 +2884,7 @@ def _validate_resume_metadata(output_path: str | None, expected: dict) -> None: "models", "policyengine_bundles", "response_contract", + "completion_budget_escalation", ): if existing.get(key) != expected.get(key): mismatches.append(key) @@ -2535,10 +3007,10 @@ def run_no_tools_eval( ``CHECKPOINT_EVERY_ROWS`` rows. Returns DataFrame with columns: - model, scenario_id, variable, prediction, explanation, raw_response, - error, elapsed_seconds, prompt_tokens, completion_tokens, total_tokens, - reasoning_tokens, cached_prompt_tokens, cache_write_prompt_tokens, - estimated_cost_usd + model, scenario_id, variable, prediction, explanation, failure_source, + raw_response, error, elapsed_seconds, prompt_tokens, completion_tokens, + total_tokens, reasoning_tokens, cached_prompt_tokens, + cache_write_prompt_tokens, estimated_cost_usd """ if models is None: models = MODELS @@ -2648,6 +3120,8 @@ def run_no_tools_eval( "total_cost_usd": None, "cost_is_estimated": None, "estimated_cost_usd": None, + "failure_sources": {}, + "budget_escalation_count": 0, } batch_size = len(scenario_programs) @@ -2678,6 +3152,9 @@ def run_no_tools_eval( "variable": variable, "prediction": prediction, "explanation": explanation, + "failure_source": result.get("failure_sources", {}).get( + variable + ), "raw_response": result["raw_response"], "error": error, "elapsed_seconds": ( @@ -2890,6 +3367,8 @@ def run_no_tools_single_output_eval( "provider_response_id": None, "provider_system_fingerprint": None, "provider_resolved_model": None, + "failure_sources": {}, + "budget_escalation_count": 0, } call_id = ":".join( @@ -2908,6 +3387,9 @@ def run_no_tools_single_output_eval( "variable": variable, "prediction": result["predictions"].get(variable), "explanation": result.get("explanations", {}).get(variable), + "failure_source": result.get("failure_sources", {}).get( + variable + ), "raw_response": result["raw_response"], "error": error, "elapsed_seconds": result.get("elapsed_seconds"), diff --git a/policybench/full_run_export.py b/policybench/full_run_export.py index d496bfe..e0b460b 100644 --- a/policybench/full_run_export.py +++ b/policybench/full_run_export.py @@ -217,13 +217,47 @@ def merge_annotations( for column in ["annotation", "failure_source", "failure_subtype"] if column in predictions.columns ] - if existing_columns: - predictions = predictions.drop(columns=existing_columns) - return predictions.merge( + recorded_columns = {column: f"_recorded_{column}" for column in existing_columns} + if recorded_columns: + predictions = predictions.rename(columns=recorded_columns) + merged = predictions.merge( annotations, on=["model", "scenario_id", "variable"], how="left", ) + recorded_ceiling_exhaustion = pd.Series(False, index=merged.index) + for column, recorded_column in recorded_columns.items(): + if column not in merged.columns: + merged[column] = merged[recorded_column] + else: + merged[column] = merged[column].combine_first(merged[recorded_column]) + if column == "failure_source": + recorded_ceiling_exhaustion = ( + merged[recorded_column].astype("string") + == "budget_exhausted_at_ceiling" + ).fillna(False) + merged.loc[recorded_ceiling_exhaustion, column] = ( + "budget_exhausted_at_ceiling" + ) + merged = merged.drop(columns=recorded_column) + invented_ceiling_exhaustion = ( + merged["failure_source"].astype("string") == "budget_exhausted_at_ceiling" + ).fillna(False) & ~recorded_ceiling_exhaustion + if invented_ceiling_exhaustion.any(): + prediction_missing = ( + merged["prediction"].isna() + if "prediction" in merged.columns + else pd.Series(True, index=merged.index) + ) + merged.loc[ + invented_ceiling_exhaustion & prediction_missing, + "failure_source", + ] = "parse_contract_failure" + merged.loc[ + invented_ceiling_exhaustion & ~prediction_missing, + "failure_source", + ] = "llm_error" + return merged def merge_case_annotations( diff --git a/policybench/model_cards.py b/policybench/model_cards.py index 3859d34..c3b31a9 100644 --- a/policybench/model_cards.py +++ b/policybench/model_cards.py @@ -27,6 +27,8 @@ from dataclasses import dataclass +from policybench.completion_budget import MAX_ESCALATED_COMPLETION_TOKENS + PROMPT_CONTRACT_VERSION = "2026-08-09-v2-scoring-contract" CLAUDE_EXPLANATION_CHUNK_SIZE = 1 @@ -57,10 +59,13 @@ class ModelCard: # True → 16,384-token completion budget on both explanation arms # (reasoning bills against the same budget as the answer). thinking_budget: bool | None = None - # Overrides the thinking-budget completion ceiling for models whose - # reasoning tail overflows 16,384. Headroom is free — only tokens - # actually generated bill. + # Overrides the thinking-class starting budget for models whose reasoning + # tail overflows 16,384. Headroom is free — only generated tokens bill. completion_token_cap: int | None = None + # Hard provider output limit, when it is lower than PolicyBench's 128k + # escalation ceiling. This is distinct from ``completion_token_cap``, + # which selects the model's starting budget rather than limiting retries. + provider_max_completion_tokens: int | None = None # Measured during onboarding; informs the run supervisor's projection # before live per-scenario costs exist. expected_cost_per_scenario_usd: float | None = None @@ -111,6 +116,36 @@ class ModelCard: "truncated it at exactly the ceiling." ), ), + "claude-sonnet-4-6": ModelCard( + litellm_id="claude-sonnet-4-6", + provider_max_completion_tokens=64_000, + notes="Provider output ceiling recorded by the serving metadata.", + ), + "claude-haiku-4-5-20251001": ModelCard( + litellm_id="claude-haiku-4-5-20251001", + provider_max_completion_tokens=64_000, + notes="Provider output ceiling recorded by the serving metadata.", + ), + "gemini/gemini-3.1-pro-preview": ModelCard( + litellm_id="gemini/gemini-3.1-pro-preview", + provider_max_completion_tokens=65_536, + notes="Provider output ceiling recorded by the serving metadata.", + ), + "gemini/gemini-3.1-flash-lite-preview": ModelCard( + litellm_id="gemini/gemini-3.1-flash-lite-preview", + provider_max_completion_tokens=65_536, + notes="Provider output ceiling recorded by the serving metadata.", + ), + "gemini/gemini-3.5-flash": ModelCard( + litellm_id="gemini/gemini-3.5-flash", + provider_max_completion_tokens=65_535, + notes="Provider output ceiling recorded by the serving metadata.", + ), + "gemini/gemini-3-flash-preview": ModelCard( + litellm_id="gemini/gemini-3-flash-preview", + provider_max_completion_tokens=65_535, + notes="Provider output ceiling recorded by the serving metadata.", + ), "xai/grok-4.5": ModelCard( litellm_id="xai/grok-4.5", answer_contract="tool", @@ -154,6 +189,7 @@ class ModelCard: litellm_id="gemini/gemini-3.6-flash", answer_contract="tool", thinking_budget=True, + provider_max_completion_tokens=65_536, expected_cost_per_scenario_usd=0.07, notes=( "Onboarded 2026-07-21: forced tool contract passed 3/3 and " @@ -263,6 +299,17 @@ def card_for(model_id: str) -> ModelCard | None: return MODEL_CARDS.get(model_id) +def completion_budget_ceiling_for(model_id: str) -> int: + """Return the documented hard provider cap or PolicyBench's 128k default.""" + card = card_for(model_id) + provider_max = card.provider_max_completion_tokens if card is not None else None + if provider_max is None: + return MAX_ESCALATED_COMPLETION_TOKENS + if provider_max <= 0: + raise ValueError("provider_max_completion_tokens must be positive") + return min(MAX_ESCALATED_COMPLETION_TOKENS, provider_max) + + def answer_contract_for(model_id: str) -> str: """Return the effective structured-answer contract for a model.""" card = card_for(model_id) diff --git a/policybench/onboard.py b/policybench/onboard.py index fa57a34..e9e0257 100644 --- a/policybench/onboard.py +++ b/policybench/onboard.py @@ -31,6 +31,7 @@ from litellm import completion, responses from policybench import eval_no_tools as harness +from policybench.completion_budget import completion_budget_from_kwargs from policybench.model_cards import ModelCard PROBE_FULL_VARIABLE_COUNT = 16 @@ -176,11 +177,7 @@ def _run_probe(name, scenario, variables, model_id, contract) -> ProbeResult: messages, kwargs, request_fn = _probe_request( scenario, variables, model_id, contract ) - budget = ( - kwargs.get("max_completion_tokens") - or kwargs.get("max_output_tokens") - or kwargs.get("max_tokens") - ) + budget = completion_budget_from_kwargs(kwargs) response = harness._run_request_with_wall_timeout(request_fn, kwargs) except Exception as error: return ProbeResult( diff --git a/policybench/runstore.py b/policybench/runstore.py index 8c64896..d3c2af3 100644 --- a/policybench/runstore.py +++ b/policybench/runstore.py @@ -85,8 +85,15 @@ #: Parse statuses recorded per prediction row. ``ok`` is a parsed value; #: ``missing`` is a row that exists (the response was attempted) but produced no -#: numeric prediction; ``replaced`` is a superseded row. -PREDICTION_PARSE_STATUSES = ("ok", "missing", "replaced", "error") +#: numeric prediction; ``budget_exhausted_at_ceiling`` is a length-terminated +#: miss after the escalation ladder; ``replaced`` is a superseded row. +PREDICTION_PARSE_STATUSES = ( + "ok", + "missing", + "budget_exhausted_at_ceiling", + "replaced", + "error", +) # --------------------------------------------------------------------------- # Column model for predictions.csv @@ -894,6 +901,8 @@ def _response_status_from_csv(record: dict[str, Any]) -> str: def _prediction_parse_status(record: dict[str, Any]) -> str: if not _is_missing(record.get("prediction")): return "ok" + if record.get("failure_source") == "budget_exhausted_at_ceiling": + return "budget_exhausted_at_ceiling" from policybench.eval_no_tools import is_infrastructure_error_text error = record.get("error") diff --git a/policybench/spend_ledger.py b/policybench/spend_ledger.py index eb5653c..3c36891 100644 --- a/policybench/spend_ledger.py +++ b/policybench/spend_ledger.py @@ -45,6 +45,13 @@ def read_spend_ledger(path: str | Path) -> list[dict]: return records +def count_budget_escalations(records: Iterable[dict]) -> int: + """Count requests that moved to a larger completion budget.""" + return sum( + record.get("escalated_from_budget_tokens") is not None for record in records + ) + + def _deduplicate_records(records: Iterable[dict]) -> list[dict]: indexed: dict[str, dict] = {} for record in records: diff --git a/policybench/supervisor.py b/policybench/supervisor.py index 2e498cb..f913528 100644 --- a/policybench/supervisor.py +++ b/policybench/supervisor.py @@ -36,10 +36,12 @@ PROMPT_CONTRACT_VERSION, answer_contract_for, card_for, + completion_budget_ceiling_for, explanation_chunk_size_for, ) from policybench.spend_ledger import ( SPEND_LEDGER_SUFFIX, + count_budget_escalations, read_spend_ledger, ) @@ -79,6 +81,7 @@ class RunState: stopped_reason: str | None = None started_at: float = 0.0 updated_at: float = 0.0 + budget_escalation_count: int = 0 def projected_total_usd(self) -> float | None: if not self.completed: @@ -163,6 +166,7 @@ def _treatment_fingerprint(self) -> dict: chunk_override=self.env.get("POLICYBENCH_CHUNK_OVERRIDE"), ), "prompt_contract_version": PROMPT_CONTRACT_VERSION, + "completion_budget_ceiling": completion_budget_ceiling_for(self.litellm_id), } def _validate_resume(self) -> dict | None: @@ -338,8 +342,18 @@ def _record(self, result: ScenarioResult) -> None: # -- heartbeat ---------------------------------------------------------- + def _budget_escalation_count_from_disk(self) -> int: + scenario_dir = self.run_dir / SCENARIO_DIR + if not scenario_dir.exists(): + return 0 + return sum( + count_budget_escalations(read_spend_ledger(path)) + for path in scenario_dir.glob(f"scenario_*.csv{SPEND_LEDGER_SUFFIX}") + ) + def write_heartbeat(self) -> None: self.state.updated_at = time.time() + self.state.budget_escalation_count = self._budget_escalation_count_from_disk() payload = { "model": self.state.model, "total": self.state.total, @@ -351,6 +365,7 @@ def write_heartbeat(self) -> None: "workers": self.state.workers, "stopped_reason": self.state.stopped_reason, "projection_warning": self.projection_warning, + "budget_escalation_count": self.state.budget_escalation_count, "treatment_fingerprint": self.treatment_fingerprint, "started_at": self.state.started_at, "updated_at": self.state.updated_at, diff --git a/tests/test_annotation_taxonomy.py b/tests/test_annotation_taxonomy.py index 28f1c46..1405bfb 100644 --- a/tests/test_annotation_taxonomy.py +++ b/tests/test_annotation_taxonomy.py @@ -1,4 +1,8 @@ -from policybench.annotation_taxonomy import infer_failure_category +from policybench.annotation_taxonomy import ( + FAILURE_SOURCE_VALUES, + infer_failure_category, + validate_failure_source, +) def test_missing_parsed_prediction_is_parse_contract_failure() -> None: @@ -18,3 +22,10 @@ def test_model_omitted_policy_amount_is_not_parse_contract_failure() -> None: assert category.failure_source == "llm_error" assert category.failure_subtype == "thresholds_rates" + + +def test_budget_exhaustion_is_a_valid_distinct_failure_source() -> None: + source = "budget_exhausted_at_ceiling" + + assert source in FAILURE_SOURCE_VALUES + assert validate_failure_source(source) == source diff --git a/tests/test_audit.py b/tests/test_audit.py index 95731b1..fc45b8a 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -14,6 +14,7 @@ ) from policybench.audit import ( AUDIT_OUTPUT_SCHEMA, + _row_failure_source, build_audit_cases, collect_audit, is_hedged, @@ -422,6 +423,57 @@ def test_parse_failure_only_case_skips_codex_and_is_deterministic(tmp_path: Path assert s0_rows.iloc[0]["failure_subtype"] == "missing_output" +def test_budget_exhaustion_source_survives_deterministic_missing_audit( + tmp_path: Path, +): + d = tmp_path / "us" + d.mkdir() + pd.DataFrame([{"scenario_id": "s0", "variable": "snap", "value": 0.0}]).to_csv( + d / "reference_outputs.csv", index=False + ) + pd.DataFrame( + [ + { + "model": "m1", + "scenario_id": "s0", + "variable": "snap", + "prediction": None, + "explanation": None, + "error": "budget_exhausted_at_ceiling: snap", + "failure_source": "budget_exhausted_at_ceiling", + } + ] + ).to_csv(d / "predictions.csv", index=False) + audit_dir = tmp_path / "audit" + + prepare_audit(d, audit_dir) + out = collect_audit(d, audit_dir) + + assert out["row"].iloc[0]["failure_source"] == "budget_exhausted_at_ceiling" + + +def test_classifier_cannot_invent_budget_exhaustion_source(): + missing_meta = {"missing_models": ["m1"], "recorded_failure_sources": {}} + parsed_meta = {"missing_models": [], "recorded_failure_sources": {}} + + assert ( + _row_failure_source( + missing_meta, + "m1", + "budget_exhausted_at_ceiling", + ) + == "parse_contract_failure" + ) + assert ( + _row_failure_source( + parsed_meta, + "m1", + "budget_exhausted_at_ceiling", + ) + == "llm_error" + ) + + def test_reprepare_drops_stale_verdict_when_case_changed(tmp_path: Path): """A verdict is invalidated when the case content (prompt) changes.""" d = tmp_path / "us" diff --git a/tests/test_eval_no_tools.py b/tests/test_eval_no_tools.py index 8e828c5..5700edd 100644 --- a/tests/test_eval_no_tools.py +++ b/tests/test_eval_no_tools.py @@ -75,6 +75,39 @@ def mini_scenario(): ) +def _chat_tool_response( + arguments: str | None, + *, + finish_reason: str, + prompt_tokens: int = 12, + completion_tokens: int = 3, + function_name: str = "submit_outputs", +): + tool_calls = [] + if arguments is not None: + tool_calls = [ + SimpleNamespace( + function=SimpleNamespace( + name=function_name, + arguments=arguments, + ) + ) + ] + message = SimpleNamespace( + content=None, + tool_calls=tool_calls, + function_call=None, + ) + return SimpleNamespace( + choices=[SimpleNamespace(message=message, finish_reason=finish_reason)], + usage=litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + @pytest.fixture def rich_scenario(): return Scenario( @@ -920,6 +953,548 @@ def test_run_single_no_tools_repairs_partial_batch_response( assert '"responses"' in result["raw_response"] +@patch("policybench.eval_no_tools.completion") +def test_length_empty_retries_same_request_with_doubled_budget( + mock_completion, + mini_scenario, + monkeypatch, + caplog, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_completion.side_effect = [ + _chat_tool_response(None, finish_reason="length", completion_tokens=256), + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500}}}', + finish_reason="stop", + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert result["budget_escalation_count"] == 1 + first_kwargs, second_kwargs = [ + call.kwargs for call in mock_completion.call_args_list + ] + assert first_kwargs["max_completion_tokens"] == 256 + assert second_kwargs["max_completion_tokens"] == 512 + assert second_kwargs["messages"] == first_kwargs["messages"] + assert "model=claude-opus-4-6" in caplog.text + assert "scenario=mini" in caplog.text + assert "from_budget=256 to_budget=512" in caplog.text + assert [ + record["completion_budget_tokens"] for record in result["spend_ledger"] + ] == [ + 256, + 512, + ] + assert result["spend_ledger"][1]["escalated_from_budget_tokens"] == 256 + + +@patch("policybench.eval_no_tools.completion") +def test_xai_length_empty_uses_max_tokens_across_multiple_escalations( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_completion.side_effect = [ + _chat_tool_response(None, finish_reason="length", completion_tokens=256), + _chat_tool_response(None, finish_reason="length", completion_tokens=512), + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500}}}', + finish_reason="stop", + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "xai/grok-4-1-fast-non-reasoning", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert [call.kwargs["max_tokens"] for call in mock_completion.call_args_list] == [ + 256, + 512, + 1024, + ] + assert result["budget_escalation_count"] == 2 + + +@patch("policybench.eval_no_tools.completion") +def test_length_partial_keeps_valid_cells_and_escalates_only_missing_cells( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + first_variable = "federal_income_tax_before_refundable_credits" + second_variable = "federal_refundable_credits" + mock_completion.side_effect = [ + _chat_tool_response( + '{"outputs":{"federal_income_tax_before_refundable_credits":' + '{"value":3500}}}', + finish_reason="length", + completion_tokens=256, + ), + _chat_tool_response( + '{"outputs":{"federal_income_tax_before_refundable_credits":' + '{"value":9999},"federal_refundable_credits":{"value":1200}}}', + finish_reason="stop", + ), + ] + + result = run_single_no_tools( + mini_scenario, + [first_variable, second_variable], + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["predictions"] == { + first_variable: 3500.0, + second_variable: 1200.0, + } + assert result["budget_escalation_count"] == 1 + retry_kwargs = mock_completion.call_args_list[1].kwargs + assert retry_kwargs["max_completion_tokens"] == 512 + retry_prompt = retry_kwargs["messages"][0]["content"] + assert f"- {second_variable}:" in retry_prompt + assert f"- {first_variable}:" not in retry_prompt + + +@patch("policybench.eval_no_tools.completion") +def test_length_missing_explanation_escalates_completion_budget( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_completion.side_effect = [ + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500}}}', + finish_reason="length", + completion_tokens=4096, + ), + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500,' + '"explanation":"Computed tax. value = 3500"}}}', + finish_reason="stop", + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=True, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert result["explanations"]["income_tax"].endswith("value = 3500") + assert [ + call.kwargs["max_completion_tokens"] for call in mock_completion.call_args_list + ] == [4096, 8192] + assert result["budget_escalation_count"] == 1 + + +@patch("policybench.eval_no_tools.completion") +def test_stop_finish_never_escalates( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_completion.return_value = _chat_tool_response(None, finish_reason="stop") + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] is None + assert result["budget_escalation_count"] == 0 + mock_completion.assert_called_once() + + +@patch("policybench.eval_no_tools.completion") +def test_length_exhaustion_at_ceiling_is_distinct_ordinary_miss( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr( + "policybench.eval_no_tools._completion_budget_ceiling", + lambda _model_id: 256, + raising=False, + ) + mock_completion.return_value = _chat_tool_response( + None, + finish_reason="length", + completion_tokens=256, + ) + + result = run_no_tools_eval( + [mini_scenario], + models={"Claude": "claude-opus-4-6"}, + programs=["income_tax"], + include_explanations=False, + ) + + row = result.iloc[0] + assert row["prediction"] is None + assert row["failure_source"] == "budget_exhausted_at_ceiling" + assert "budget_escalation_count" not in result.columns + mock_completion.assert_called_once() + + +@patch("policybench.eval_no_tools.completion") +def test_escalation_lands_exactly_on_provider_ceiling( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr( + "policybench.eval_no_tools._completion_budget_ceiling", + lambda _model_id: 300, + ) + mock_completion.side_effect = [ + _chat_tool_response(None, finish_reason="length", completion_tokens=256), + _chat_tool_response(None, finish_reason="length", completion_tokens=300), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert [ + call.kwargs["max_completion_tokens"] for call in mock_completion.call_args_list + ] == [256, 300] + assert result["failure_sources"] == {"income_tax": "budget_exhausted_at_ceiling"} + + +@patch("policybench.eval_no_tools.responses") +def test_responses_incomplete_max_output_tokens_escalates( + mock_responses, mini_scenario +): + mock_responses.side_effect = [ + SimpleNamespace( + status="incomplete", + incomplete_details=SimpleNamespace(reason="max_output_tokens"), + output_text="", + output=[], + usage=SimpleNamespace( + input_tokens=12, + output_tokens=256, + total_tokens=268, + ), + ), + SimpleNamespace( + status="completed", + incomplete_details=None, + output_text="", + output=[ + SimpleNamespace( + type="function_call", + name="submit_outputs", + arguments='{"outputs":{"income_tax":{"value":3500}}}', + ) + ], + usage=SimpleNamespace( + input_tokens=12, + output_tokens=3, + total_tokens=15, + ), + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "gpt-5.4", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert [ + call.kwargs["max_output_tokens"] for call in mock_responses.call_args_list + ] == [ + 256, + 512, + ] + + +@patch("policybench.eval_no_tools.responses") +def test_responses_bare_incomplete_at_requested_budget_escalates( + mock_responses, mini_scenario +): + mock_responses.side_effect = [ + SimpleNamespace( + status="incomplete", + incomplete_details=None, + output_text="", + output=[], + usage=SimpleNamespace( + input_tokens=12, + output_tokens=256, + total_tokens=268, + ), + ), + SimpleNamespace( + status="completed", + incomplete_details=None, + output_text="", + output=[ + SimpleNamespace( + type="function_call", + name="submit_outputs", + arguments='{"outputs":{"income_tax":{"value":3500}}}', + ) + ], + usage=SimpleNamespace( + input_tokens=12, + output_tokens=3, + total_tokens=15, + ), + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "gpt-5.4", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert [ + call.kwargs["max_output_tokens"] for call in mock_responses.call_args_list + ] == [256, 512] + + +@patch("policybench.eval_no_tools.responses") +def test_responses_bare_incomplete_below_requested_budget_does_not_escalate( + mock_responses, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_responses.return_value = SimpleNamespace( + status="incomplete", + incomplete_details=None, + output_text="", + output=[], + usage=SimpleNamespace( + input_tokens=12, + output_tokens=100, + total_tokens=112, + ), + ) + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "gpt-5.4", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["prediction"] is None + assert result["budget_escalation_count"] == 0 + mock_responses.assert_called_once() + + +@patch("policybench.eval_no_tools.completion") +def test_budget_escalation_count_lands_in_eval_manifest( + mock_completion, + mini_scenario, + tmp_path, +): + mock_completion.side_effect = [ + _chat_tool_response(None, finish_reason="length", completion_tokens=256), + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500}}}', + finish_reason="stop", + ), + ] + output_path = tmp_path / "predictions.csv" + + run_no_tools_eval( + [mini_scenario], + models={"Claude": "claude-opus-4-6"}, + programs=["income_tax"], + output_path=str(output_path), + include_explanations=False, + ) + + manifest = json.loads((tmp_path / "predictions.csv.meta.json").read_text()) + assert manifest["budget_escalation_count"] == 1 + + +@patch("policybench.eval_no_tools.completion") +def test_escalated_repair_error_preserves_partial_payload_and_count( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 1) + first_variable = "federal_income_tax_before_refundable_credits" + second_variable = "federal_refundable_credits" + mock_completion.side_effect = [ + _chat_tool_response(None, finish_reason="stop"), + _chat_tool_response( + '{"outputs":{"federal_income_tax_before_refundable_credits":' + '{"value":3500}}}', + finish_reason="length", + completion_tokens=256, + ), + RequestWallTimeoutError("escalated request timed out"), + ] + + result = run_single_no_tools( + mini_scenario, + [first_variable, second_variable], + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["predictions"] == { + first_variable: 3500.0, + second_variable: None, + } + assert result["budget_escalation_count"] == 1 + assert len(result["spend_ledger"]) == 3 + assert "RequestWallTimeoutError" in result["error"] + + +@patch("policybench.eval_no_tools.completion") +def test_escalated_initial_error_preserves_partial_payload_and_count( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + first_variable = "federal_income_tax_before_refundable_credits" + second_variable = "federal_refundable_credits" + mock_completion.side_effect = [ + _chat_tool_response( + '{"outputs":{"federal_income_tax_before_refundable_credits":' + '{"value":3500}}}', + finish_reason="length", + completion_tokens=256, + ), + RequestWallTimeoutError("escalated request timed out"), + ] + + result = run_single_no_tools( + mini_scenario, + [first_variable, second_variable], + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + assert result["predictions"] == { + first_variable: 3500.0, + second_variable: None, + } + assert result["budget_escalation_count"] == 1 + assert len(result["spend_ledger"]) == 2 + assert "RequestWallTimeoutError" in result["error"] + + +@patch("policybench.eval_no_tools.completion") +def test_escalated_initial_error_without_valid_cells_propagates( + mock_completion, + mini_scenario, +): + mock_completion.side_effect = [ + _chat_tool_response( + None, + finish_reason="length", + completion_tokens=256, + ), + RequestWallTimeoutError("escalated request timed out"), + ] + + with pytest.raises(RequestWallTimeoutError, match="escalated request timed out"): + run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=False, + _allow_chunking=False, + ) + + +@patch("policybench.eval_no_tools.completion") +def test_length_empty_explanation_followup_escalates_budget( + mock_completion, + mini_scenario, + monkeypatch, +): + monkeypatch.setattr("policybench.eval_no_tools.MAX_REPAIR_ROUNDS", 0) + mock_completion.side_effect = [ + _chat_tool_response( + '{"outputs":{"income_tax":{"value":3500}}}', + finish_reason="stop", + completion_tokens=100, + ), + _chat_tool_response( + None, + finish_reason="length", + completion_tokens=4096, + function_name="submit_explanations", + ), + _chat_tool_response( + '{"income_tax":"Computed tax. value = 3500"}', + finish_reason="stop", + function_name="submit_explanations", + ), + ] + + result = run_single_no_tools( + mini_scenario, + "income_tax", + "claude-opus-4-6", + include_explanations=True, + _allow_chunking=False, + ) + + assert result["prediction"] == 3500.0 + assert result["explanations"]["income_tax"].endswith("value = 3500") + assert [ + call.kwargs["max_completion_tokens"] for call in mock_completion.call_args_list + ] == [4096, 4096, 8192] + assert result["budget_escalation_count"] == 1 + assert [record["phase"] for record in result["spend_ledger"]] == [ + "initial", + "explanation_repair", + "explanation_repair", + ] + + @patch("policybench.eval_no_tools._request_predictions_once") def test_repair_does_not_overwrite_valid_prediction( mock_request_predictions, @@ -2358,6 +2933,11 @@ def test_run_no_tools_eval_writes_resume_metadata( == "outputs.{variable}.{value,explanation}" ) assert len(metadata["response_contract"]["prompt_template_sha256"]) == 64 + assert metadata["completion_budget_escalation"] == { + "default_ceiling": 128000, + "model_ceilings": {"gpt-5.4": 128000}, + "strategy": "double_on_length_with_missing_payload", + } @patch("policybench.eval_no_tools.run_single_no_tools") diff --git a/tests/test_export_full_run.py b/tests/test_export_full_run.py index 3cbc0d0..712631b 100644 --- a/tests/test_export_full_run.py +++ b/tests/test_export_full_run.py @@ -2,10 +2,12 @@ import pandas as pd +from policybench.case_annotations import wrong_prediction_rows from policybench.full_run_export import ( load_annotations, load_case_annotations, load_predictions, + merge_annotations, merge_case_annotations, ) @@ -232,6 +234,196 @@ def test_merge_case_annotations_attaches_notes_to_prediction_rows() -> None: assert merged["case_failure_subtypes"].tolist() == ["thresholds_rates"] +def test_merge_annotations_preserves_unannotated_budget_exhaustion() -> None: + predictions = pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "failure_source": "budget_exhausted_at_ceiling", + }, + { + "model": "model_b", + "scenario_id": "s001", + "variable": "income_tax", + "failure_source": None, + }, + ] + ) + annotations = pd.DataFrame( + [ + { + "model": "model_b", + "scenario_id": "s001", + "variable": "income_tax", + "annotation": "Used the wrong bracket.", + "failure_source": "llm_error", + "failure_subtype": "thresholds_rates", + } + ] + ) + + merged = merge_annotations(predictions, annotations) + + assert merged.loc[merged["model"] == "model_a", "failure_source"].iloc[0] == ( + "budget_exhausted_at_ceiling" + ) + + +def test_merge_annotations_cannot_overwrite_recorded_budget_exhaustion() -> None: + predictions = pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "failure_source": "budget_exhausted_at_ceiling", + } + ] + ) + annotations = pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "annotation": "Stale parse annotation.", + "failure_source": "parse_contract_failure", + "failure_subtype": "missing_output", + } + ] + ) + + merged = merge_annotations(predictions, annotations) + + assert merged["failure_source"].tolist() == ["budget_exhausted_at_ceiling"] + + +def test_merge_annotations_cannot_invent_budget_exhaustion() -> None: + predictions = pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "prediction": 100.0, + }, + { + "model": "model_b", + "scenario_id": "s001", + "variable": "income_tax", + "prediction": None, + }, + ] + ) + annotations = pd.DataFrame( + [ + { + "model": model, + "scenario_id": "s001", + "variable": "income_tax", + "annotation": "Incorrect reserved source.", + "failure_source": "budget_exhausted_at_ceiling", + "failure_subtype": "missing_output", + } + for model in ("model_a", "model_b") + ] + ) + + merged = merge_annotations(predictions, annotations) + + assert merged["failure_source"].tolist() == [ + "llm_error", + "parse_contract_failure", + ] + + +def test_wrong_prediction_rows_keeps_recorded_budget_source_over_annotation( + tmp_path: Path, +) -> None: + country_dir = tmp_path / "full_run" / "us" + country_dir.mkdir(parents=True) + pd.DataFrame( + [{"scenario_id": "s001", "variable": "income_tax", "value": 100.0}] + ).to_csv(country_dir / "reference_outputs.csv", index=False) + pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "prediction": None, + "failure_source": "budget_exhausted_at_ceiling", + } + ] + ).to_csv(country_dir / "predictions.csv", index=False) + annotations_dir = country_dir.parent / "annotations" + annotations_dir.mkdir() + pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "annotation": "Stale parse annotation.", + "failure_source": "parse_contract_failure", + "failure_subtype": "missing_output", + } + ] + ).to_csv(annotations_dir / "us_tax_annotations.csv", index=False) + + wrong = wrong_prediction_rows(country_dir) + + assert wrong["failure_source"].tolist() == ["budget_exhausted_at_ceiling"] + + +def test_wrong_prediction_rows_cannot_invent_budget_source(tmp_path: Path) -> None: + country_dir = tmp_path / "full_run" / "us" + country_dir.mkdir(parents=True) + pd.DataFrame( + [{"scenario_id": "s001", "variable": "income_tax", "value": 100.0}] + ).to_csv(country_dir / "reference_outputs.csv", index=False) + pd.DataFrame( + [ + { + "model": "model_a", + "scenario_id": "s001", + "variable": "income_tax", + "prediction": 200.0, + }, + { + "model": "model_b", + "scenario_id": "s001", + "variable": "income_tax", + "prediction": None, + }, + ] + ).to_csv(country_dir / "predictions.csv", index=False) + annotations_dir = country_dir.parent / "annotations" + annotations_dir.mkdir() + pd.DataFrame( + [ + { + "model": model, + "scenario_id": "s001", + "variable": "income_tax", + "annotation": "Incorrect reserved source.", + "failure_source": "budget_exhausted_at_ceiling", + "failure_subtype": "missing_output", + } + for model in ("model_a", "model_b") + ] + ).to_csv(annotations_dir / "us_tax_annotations.csv", index=False) + + wrong = wrong_prediction_rows(country_dir).sort_values("model") + + assert wrong["failure_source"].tolist() == [ + "llm_error", + "parse_contract_failure", + ] + + def test_available_countries_detects_only_populated_dirs(tmp_path): from policybench.full_run_export import _available_countries diff --git a/tests/test_model_cards.py b/tests/test_model_cards.py index 9df7ba0..1af06aa 100644 --- a/tests/test_model_cards.py +++ b/tests/test_model_cards.py @@ -11,11 +11,12 @@ from policybench.config import MODELS from policybench.eval_no_tools import ( _answer_contract_for_model, + _completion_budget_ceiling, _completion_controls, _request_timeout_seconds, _required_explanation_chunk_size, ) -from policybench.model_cards import card_for +from policybench.model_cards import ModelCard, card_for # model_id -> (contract, chunk_size, timeout_s, budget_for_16_vars_with_expl) EXPECTED = { @@ -92,3 +93,46 @@ def test_no_chunking_when_explanations_off(): ) def test_gpt_56_cards_record_measured_full_run_cost(model_id, expected_cost): assert card_for(model_id).expected_cost_per_scenario_usd == expected_cost + + +@pytest.mark.parametrize( + ("provider_max", "expected"), [(30_000, 30_000), (200_000, 128_000)] +) +def test_provider_max_limits_escalation_ceiling(monkeypatch, provider_max, expected): + card = ModelCard( + litellm_id="provider/model", + provider_max_completion_tokens=provider_max, + ) + monkeypatch.setattr("policybench.model_cards.card_for", lambda _model_id: card) + + assert _completion_budget_ceiling("provider/model") == expected + + +def test_provider_max_clamps_initial_completion_budget(monkeypatch): + card = ModelCard( + litellm_id="gpt-5-provider-capped", + provider_max_completion_tokens=2_000, + ) + monkeypatch.setattr("policybench.model_cards.card_for", lambda _model_id: card) + + assert _completion_controls( + "gpt-5-provider-capped", + include_explanations=True, + variables=["income_tax"], + ) == {"max_completion_tokens": 2_000} + + +@pytest.mark.parametrize( + ("model_id", "expected"), + [ + ("claude-sonnet-4-6", 64_000), + ("claude-haiku-4-5-20251001", 64_000), + ("gemini/gemini-3.1-pro-preview", 65_536), + ("gemini/gemini-3.1-flash-lite-preview", 65_536), + ("gemini/gemini-3.5-flash", 65_535), + ("gemini/gemini-3-flash-preview", 65_535), + ("gemini/gemini-3.6-flash", 65_536), + ], +) +def test_roster_model_provider_maximums_limit_escalation(model_id, expected): + assert _completion_budget_ceiling(model_id) == expected diff --git a/tests/test_runstore.py b/tests/test_runstore.py index 47032b2..a53f3fc 100644 --- a/tests/test_runstore.py +++ b/tests/test_runstore.py @@ -429,6 +429,29 @@ def test_status_counts_shape(tmp_path): store.close() +def test_status_counts_bucket_budget_exhaustion_separately(tmp_path): + store = RunStore(tmp_path / "run.db") + store.create_run("r1") + frame = _make_predictions_frame( + models=("model-a",), + scenarios=("scenario_000",), + variables=("snap", "ssi"), + ) + frame["run_id"] = "r1" + exhausted = frame["variable"] == "ssi" + frame.loc[exhausted, "prediction"] = None + frame.loc[exhausted, "failure_source"] = "budget_exhausted_at_ceiling" + + store.upsert_predictions(frame) + counts = store.status_counts("r1")["predictions_by_model_parse_status"] + + assert {(row["model"], row["parse_status"]): row["n"] for row in counts} == { + ("model-a", "budget_exhausted_at_ceiling"): 1, + ("model-a", "ok"): 1, + } + store.close() + + # --------------------------------------------------------------------------- # Import / export round-trip on synthetic data # --------------------------------------------------------------------------- diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py index 7115470..15a9b9e 100644 --- a/tests/test_supervisor.py +++ b/tests/test_supervisor.py @@ -15,7 +15,7 @@ import pytest from policybench.config import MODELS -from policybench.spend_ledger import spend_ledger_path +from policybench.spend_ledger import spend_ledger_path, upsert_spend_ledger from policybench.supervisor import ( ADAPTIVE_WINDOW, BUDGET_STOP_FRACTION, @@ -51,10 +51,12 @@ def stub_worker( cost_per_scenario: float = 0.1, fail_indices: set[int] | None = None, timeout_indices: set[int] | None = None, + budget_escalation_counts: dict[int, int] | None = None, ): """Replace _spawn with a no-op process and synthesize the scenario CSV.""" fail_indices = fail_indices or set() timeout_indices = timeout_indices or set() + budget_escalation_counts = budget_escalation_counts or {} def fake_spawn(index: int): if index not in fail_indices: @@ -67,6 +69,19 @@ def fake_spawn(index: int): "total_cost_usd": [cost_per_scenario / 2] * 2, } ).to_csv(out, index=False) + escalation_count = budget_escalation_counts.get(index, 0) + if escalation_count: + upsert_spend_ledger( + spend_ledger_path(out), + [ + { + "call_key": f"sync:{index}:{escalation_index}", + "escalated_from_budget_tokens": 256 * 2**escalation_index, + "completion_budget_tokens": 512 * 2**escalation_index, + } + for escalation_index in range(escalation_count) + ], + ) if index in timeout_indices: log = supervisor.scenario_csv(index).with_suffix(".log") log.parent.mkdir(parents=True, exist_ok=True) @@ -92,9 +107,33 @@ def test_happy_path_completes_all_and_combines(manifest, tmp_path, monkeypatch): "tool_choice_mode": "forced", "chunk_size": None, "prompt_contract_version": "2026-08-09-v2-scoring-contract", + "completion_budget_ceiling": 128000, } +def test_budget_escalation_counts_land_in_run_state( + manifest, + tmp_path, + monkeypatch, +): + supervisor = make_supervisor(manifest, tmp_path) + stub_worker( + supervisor, + monkeypatch, + budget_escalation_counts={0: 2, 3: 1}, + ) + + state = supervisor.run(poll_seconds=0.01) + + assert state.budget_escalation_count == 3 + heartbeat = json.loads((supervisor.run_dir / "run_state.json").read_text()) + assert heartbeat["budget_escalation_count"] == 3 + + resumed = make_supervisor(manifest, tmp_path) + resumed_state = resumed.run(poll_seconds=0) + assert resumed_state.budget_escalation_count == 3 + + def test_resume_skips_completed_scenarios(manifest, tmp_path, monkeypatch): initial = make_supervisor(manifest, tmp_path) stub_worker(initial, monkeypatch)