Escalate completion budgets on length-truncation instead of scoring the death

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

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(

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()

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
 
 

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)