From ca105f9a5068c91d608b45aa6338283380e712e6 Mon Sep 17 00:00:00 2001 From: Murali Chillakuru Date: Thu, 30 Jul 2026 11:30:45 -0400 Subject: [PATCH 1/3] feat(results): report scored coverage alongside every rate Rows that could not be judged are dropped from every rate denominator, and nothing showed how many. That bias is directional rather than random: a provider content filter rejects the most adversarial transcripts, which are exactly the ones most likely to contain a real violation, so a run whose worst rows were filtered reports a low violation rate and reads as a pass. Adds compute_coverage() and format_coverage(). Coverage is written into metrics.json and printed immediately above the rates in the run headline, because a rate is only meaningful next to the denominator it was computed over. Above 10 percent excluded, a warning states the direction of the bias - the true rate is likely higher, not merely uncertain. infer_judge_status collapses every non-ok status to judge_failed, so the breakdown reads the raw judge_status field instead; otherwise filter_skipped and a genuine judge error are indistinguishable, and they mean different things. A test covers that. Also: a failed metrics.json write was logged at DEBUG, which lost the whole cost record while the run still reported success. It now warns and records metrics_write_failed on the manifest, so the gap is discoverable after the fact. Also: _detach_executor_from_atexit read executor._threads unguarded and reported failures at DEBUG. It runs on the graceful-teardown path, so a silent failure there means worker threads block process exit with nothing explaining why. The attribute is now guarded and failures warn. The Python 3.13+ absence of threading._shutdown_locks is expected and stays at DEBUG. --- assert_ai/core/config_model.py | 5 ++ assert_ai/core/runtime_safety.py | 40 ++++++++-- assert_ai/results.py | 76 ++++++++++++++++++ assert_ai/runner.py | 32 +++++++- tests/test_coverage_reporting.py | 133 +++++++++++++++++++++++++++++++ 5 files changed, 278 insertions(+), 8 deletions(-) create mode 100644 tests/test_coverage_reporting.py diff --git a/assert_ai/core/config_model.py b/assert_ai/core/config_model.py index 7ef40b56..19fa1518 100644 --- a/assert_ai/core/config_model.py +++ b/assert_ai/core/config_model.py @@ -278,6 +278,11 @@ class RunManifest: progress: dict[str, Any] | None = None artifact_versions: dict[str, dict[str, Any]] = field(default_factory=dict) stage_timings: dict[str, dict[str, Any]] = field(default_factory=dict) + # Set when metrics.json could not be written. Without this a failed write + # loses the entire cost record while the run still reports success, so the + # gap has to be recorded somewhere durable rather than only in the log. + metrics_write_failed: bool | None = None + metrics_write_error: str | None = None def to_dict(self) -> dict[str, Any]: empty_collection_keys = {"artifact_versions", "stage_timings"} diff --git a/assert_ai/core/runtime_safety.py b/assert_ai/core/runtime_safety.py index fc9902b7..4cd09304 100644 --- a/assert_ai/core/runtime_safety.py +++ b/assert_ai/core/runtime_safety.py @@ -339,23 +339,50 @@ def _detach_executor_from_atexit(executor: ThreadPoolExecutor) -> None: is closed" tracebacks as their resources GC against the closed loop) until the OS reaps them at process exit; we cannot stop them cleanly from Python. + + Every attribute reached here is private to CPython or to + ``concurrent.futures``, so each access is guarded and a failure is reported + at WARNING rather than DEBUG. This runs on the graceful-teardown path, and a + silent failure here means worker threads block process exit with nothing + explaining why. """ + worker_threads = getattr(executor, "_threads", None) + if worker_threads is None: + log.warning( + "Could not read the executor's worker threads (ThreadPoolExecutor " + "internals changed in this Python build), so they were not detached " + "from interpreter shutdown. Leaked worker threads may delay process " + "exit." + ) + return + try: threads_queues = getattr(_cft, "_threads_queues", None) if threads_queues is not None: - for t in list(executor._threads): + for t in list(worker_threads): threads_queues.pop(t, None) + else: + log.warning( + "concurrent.futures._threads_queues is unavailable; the " + "executor was not detached from its atexit handler and may " + "delay process exit." + ) except Exception: # noqa: BLE001 - log.debug( - "Failed to detach executor from concurrent.futures atexit", + log.warning( + "Failed to detach executor from the concurrent.futures atexit " + "handler; leaked worker threads may delay process exit.", exc_info=True, ) try: shutdown_locks = getattr(threading, "_shutdown_locks", None) shutdown_locks_lock = getattr(threading, "_shutdown_locks_lock", None) if shutdown_locks is None: + # Expected on Python 3.13+, where this join moved behind C-level + # thread handles and _AbandonableThreadPoolExecutor avoids it by + # creating daemon workers instead. Not a problem worth warning about. + log.debug("threading._shutdown_locks is unavailable; nothing to detach") return - for t in list(executor._threads): + for t in list(worker_threads): lock = getattr(t, "_tstate_lock", None) if lock is None: continue @@ -365,8 +392,9 @@ def _detach_executor_from_atexit(executor: ThreadPoolExecutor) -> None: else: shutdown_locks.discard(lock) except Exception: # noqa: BLE001 - log.debug( - "Failed to detach executor from threading._shutdown", + log.warning( + "Failed to detach executor from threading._shutdown; leaked worker " + "threads may delay process exit.", exc_info=True, ) diff --git a/assert_ai/results.py b/assert_ai/results.py index 8d8bbb3e..5e077274 100644 --- a/assert_ai/results.py +++ b/assert_ai/results.py @@ -5,6 +5,8 @@ from __future__ import annotations +import logging +from collections import Counter from pathlib import Path from statistics import mean, median from typing import Any, Iterable @@ -17,6 +19,71 @@ is_valid_event_flag, ) +log = logging.getLogger(__name__) + +# Above this share of unscored rows, the reported rates describe a small enough +# slice of the suite that quoting them without the coverage is misleading. +COVERAGE_WARN_THRESHOLD = 0.10 + + +def compute_coverage(rows: list[dict[str, Any]]) -> dict[str, Any]: + """Summarise how much of ``rows`` actually produced a usable verdict. + + Rows that could not be judged are excluded from every rate denominator. That + exclusion is not random: a provider content filter rejects the *most* + adversarial transcripts, which are the ones most likely to contain a real + violation. So a run whose worst rows were dropped reports a low violation + rate and reads as a pass. Reporting the denominator alongside the rate is + what makes that visible. + + ``infer_judge_status`` collapses every non-ok status to ``judge_failed``, so + the per-status breakdown reads the raw ``judge_status`` field to keep + ``filter_skipped`` distinguishable from a judge error. + """ + total = len(rows) + by_status: Counter[str] = Counter() + for row in rows: + inferred = infer_judge_status(row) + if inferred == "ok": + by_status["ok"] += 1 + continue + raw = row.get("judge_status") + if isinstance(raw, str) and raw and raw != "ok": + by_status[raw] += 1 + else: + by_status[inferred] += 1 + + scored = by_status.get("ok", 0) + excluded = total - scored + return { + "total": total, + "scored": scored, + "excluded": excluded, + "scored_rate": (scored / total) if total else 0.0, + "excluded_rate": (excluded / total) if total else 0.0, + "by_status": dict(by_status), + "below_threshold": bool(total) and (excluded / total) > COVERAGE_WARN_THRESHOLD, + } + + +def format_coverage(coverage: dict[str, Any]) -> str: + """Render coverage as a single line to print directly above the rates.""" + total = coverage.get("total", 0) + scored = coverage.get("scored", 0) + rate = coverage.get("scored_rate", 0.0) * 100.0 + line = f"Scored {scored}/{total} ({rate:.1f}%)" + excluded_statuses = { + status: count + for status, count in (coverage.get("by_status") or {}).items() + if status != "ok" and count + } + if excluded_statuses: + detail = " · ".join( + f"{count} {status}" for status, count in sorted(excluded_statuses.items()) + ) + line += f" ! {detail}" + return line + def current_stage_status(manifest: dict[str, Any] | None) -> tuple[str, str]: """Return overall manifest status and the currently running stage, if any.""" @@ -283,6 +350,14 @@ def _compute_test_set_metrics( scored_rows = [row for row in rows if infer_judge_status(row) == "ok"] judge_failures = len(rows) - len(scored_rows) + coverage = compute_coverage(rows) + if coverage["below_threshold"]: + log.warning( + "%s - rates below describe only the scored rows. Excluded rows are " + "not a random sample: content filters reject the most adversarial " + "transcripts, so the true rate is likely higher than reported.", + format_coverage(coverage), + ) dimensions = { dim: compute_dimension_summary(scored_rows, dim) for dim in detect_dimensions(scored_rows) @@ -291,6 +366,7 @@ def _compute_test_set_metrics( metrics: dict[str, Any] = { "total": len(rows), "scored_total": len(scored_rows), + "coverage": coverage, "judge_failures": judge_failures, "judge_failure_rate": judge_failures / len(rows), "policy_violation_rate": dimension_rate({"dimensions": dimensions}, "policy_violation"), diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 9e98264e..9c5a1743 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -521,8 +521,10 @@ def _log_run_headline(run_root: Path) -> None: # Imported lazily to avoid a hard dependency for callers that import the # runner without ever invoking it (e.g. test scaffolding). from assert_ai.results import ( + compute_coverage, compute_prompt_metrics, compute_scenario_metrics, + format_coverage, ) from assert_ai.core.io import load_jsonl @@ -551,6 +553,17 @@ def _log_run_headline(run_root: Path) -> None: log.info(f" Judge: {judge}") log.info(f" Total: {total} ({scored} scored)") + # Coverage sits directly above the rates, because a rate is only meaningful + # alongside the denominator it was computed over. + combined_coverage = compute_coverage(score_rows) + log.info(f" {format_coverage(combined_coverage)}") + if combined_coverage["below_threshold"]: + log.warning( + " Excluded rows are not a random sample - content filters reject the " + "most adversarial transcripts, so the true rate is likely higher than " + "the rates below." + ) + def _fmt_rate(value: Any) -> str: if value is None or not isinstance(value, (int, float)): return "—" @@ -936,6 +949,7 @@ def _run_stages_inner( total_elapsed = time.monotonic() - pipeline_start metrics_written = False + metrics_write_error: str | None = None if run_root is not None and stage_usage: try: metrics_path = run_root / "metrics.json" @@ -952,8 +966,17 @@ def _run_stages_inner( f"{_format_token_count(totals['output_tokens'])} out · " f"{cache_pct:.1f}% cached" ) - except Exception: # noqa: BLE001 - log.debug("Failed to write metrics.json", exc_info=True) + except Exception as exc: # noqa: BLE001 + # This was logged at DEBUG, so a failed write lost the entire cost + # record while the run still reported success. The failure is now + # visible in the log and recorded in the manifest. + metrics_write_error = f"{type(exc).__name__}: {exc}" + log.warning( + "Failed to write metrics.json (%s). Token usage and cost " + "accounting for this run were not recorded.", + metrics_write_error, + ) + log.debug("metrics.json write traceback", exc_info=True) if failed_stage is None: log.info(f"Pipeline completed ({total_elapsed:.1f}s)") @@ -982,6 +1005,11 @@ def _run_stages_inner( manifest.ended_at = datetime.now(timezone.utc).isoformat() manifest.status = "completed" if failed_stage is None else "failed" + if metrics_write_error is not None: + # Recorded on the manifest so the gap is discoverable later, not only in + # whatever log stream happened to be attached at the time. + manifest.metrics_write_failed = True + manifest.metrics_write_error = metrics_write_error _record_run_artifacts(manifest, ctx, run_root) _write_manifest(manifest, run_root) return 0 if failed_stage is None else 1 diff --git a/tests/test_coverage_reporting.py b/tests/test_coverage_reporting.py new file mode 100644 index 00000000..e0418f2a --- /dev/null +++ b/tests/test_coverage_reporting.py @@ -0,0 +1,133 @@ +"""Tests for coverage reporting on judged rows. + +Rows that could not be judged are dropped from every rate denominator. The +exclusion is directional, not random: a provider content filter rejects the most +adversarial transcripts, which are the ones most likely to contain a real +violation. Reporting the rate without the denominator therefore biases it +optimistically. +""" + +from __future__ import annotations + +from assert_ai.results import ( + COVERAGE_WARN_THRESHOLD, + compute_coverage, + format_coverage, +) + + +def _ok_row(**extra): + row = { + "judge_status": "ok", + "verdict": { + "dimensions": {"policy_violation": False, "overrefusal": False}, + "node_judgments": [ + {"node_name": "b1", "violated": False, "confidence": "high"} + ], + }, + "score_keys": ["policy_violation", "overrefusal"], + } + row.update(extra) + return row + + +def _excluded_row(status: str): + return { + "judge_status": status, + "verdict": {"error": status}, + "score_keys": ["policy_violation", "overrefusal"], + } + + +class TestComputeCoverage: + def test_all_scored(self): + cov = compute_coverage([_ok_row(), _ok_row()]) + assert cov["total"] == 2 + assert cov["scored"] == 2 + assert cov["excluded"] == 0 + assert cov["scored_rate"] == 1.0 + assert cov["below_threshold"] is False + + def test_empty_rows(self): + cov = compute_coverage([]) + assert cov["total"] == 0 + assert cov["scored"] == 0 + assert cov["scored_rate"] == 0.0 + assert cov["below_threshold"] is False + + def test_denominator_excludes_unjudged_rows(self): + rows = [_ok_row() for _ in range(7)] + [ + _excluded_row("filter_skipped") for _ in range(3) + ] + cov = compute_coverage(rows) + assert cov["total"] == 10 + assert cov["scored"] == 7 + assert cov["excluded"] == 3 + assert cov["scored_rate"] == 0.7 + + def test_filter_skipped_is_distinguished_from_judge_failure(self): + """infer_judge_status flattens both to judge_failed; the breakdown must not.""" + rows = [ + _ok_row(), + _excluded_row("filter_skipped"), + _excluded_row("judge_failed"), + ] + cov = compute_coverage(rows) + assert cov["by_status"]["filter_skipped"] == 1 + assert cov["by_status"]["judge_failed"] == 1 + assert cov["by_status"]["ok"] == 1 + + def test_threshold_not_tripped_just_below(self): + rows = [_ok_row() for _ in range(91)] + [_excluded_row("judge_failed") for _ in range(9)] + cov = compute_coverage(rows) + assert cov["excluded_rate"] < COVERAGE_WARN_THRESHOLD + assert cov["below_threshold"] is False + + def test_threshold_tripped_just_above(self): + rows = [_ok_row() for _ in range(89)] + [_excluded_row("judge_failed") for _ in range(11)] + cov = compute_coverage(rows) + assert cov["excluded_rate"] > COVERAGE_WARN_THRESHOLD + assert cov["below_threshold"] is True + + def test_row_claiming_ok_without_a_verdict_is_not_counted_as_scored(self): + rows = [{"judge_status": "ok", "verdict": {}, "score_keys": ["policy_violation"]}] + cov = compute_coverage(rows) + assert cov["scored"] == 0 + assert cov["by_status"].get("judge_failed") == 1 + +class TestFormatCoverage: + def test_reports_scored_over_total(self): + line = format_coverage(compute_coverage([_ok_row(), _excluded_row("judge_failed")])) + assert "1/2" in line + assert "50.0%" in line + + def test_lists_excluded_statuses(self): + rows = [_ok_row()] + [_excluded_row("filter_skipped") for _ in range(2)] + line = format_coverage(compute_coverage(rows)) + assert "2 filter_skipped" in line + + def test_clean_run_has_no_exclusion_detail(self): + line = format_coverage(compute_coverage([_ok_row()])) + assert "filter_skipped" not in line + assert "judge_failed" not in line + + +class TestMetricsIncludeCoverage: + def test_coverage_present_in_computed_metrics(self): + from assert_ai.results import compute_prompt_metrics + + rows = [_ok_row() for _ in range(3)] + [_excluded_row("filter_skipped")] + metrics = compute_prompt_metrics(rows) + assert metrics is not None + assert metrics["coverage"]["total"] == 4 + assert metrics["coverage"]["scored"] == 3 + + def test_low_coverage_logs_a_warning(self, caplog): + from assert_ai.results import compute_prompt_metrics + + rows = [_ok_row() for _ in range(5)] + [ + _excluded_row("filter_skipped") for _ in range(5) + ] + with caplog.at_level("WARNING"): + compute_prompt_metrics(rows) + assert "true rate is likely higher" in caplog.text From 8eeeb45ad2cc3275e46ffb82266d9b43673117b0 Mon Sep 17 00:00:00 2001 From: Murali Chillakuru Date: Thu, 30 Jul 2026 11:36:01 -0400 Subject: [PATCH 2/3] feat(results): report chance-corrected judge agreement and a judge fingerprint Multi-judge runs recorded a raw percent-agreement figure for the first dimension of each row. Raw agreement is not a reliability measure: violation rates are skewed, and on a skewed base rate two judges voting independently agree most of the time by chance alone, so a high figure can describe almost no real reliability. Adds fleiss_kappa() to analysis/stats.py with no new dependency, verified against the published Fleiss (1971) worked example to four decimal places rather than only against its own edge cases. None is treated as a real category, because a judge marking a dimension not-applicable took a position rather than leaving a gap. Kappa is computed at run level, not per row: it needs many items to estimate the marginal category distribution, so a per-row value would be degenerate. compute_judge_agreement pools votes across rows and reports per dimension, warning below the Landis and Koch 0.60 benchmark. The existing per-row agreement field is left as it was. Adds compute_judge_fingerprint() over judge model, prompt hash, dimension set and judge count, plus warn_if_judge_changed(). Swapping a judge model to cut cost moves rates on an unchanged target, and without this the shift is attributed to the target. --- assert_ai/analysis/stats.py | 77 ++++++++++++++++ assert_ai/results.py | 142 ++++++++++++++++++++++++++++++ tests/test_judge_agreement.py | 161 ++++++++++++++++++++++++++++++++++ 3 files changed, 380 insertions(+) create mode 100644 tests/test_judge_agreement.py diff --git a/assert_ai/analysis/stats.py b/assert_ai/analysis/stats.py index 0494e7c6..cc21f3ce 100644 --- a/assert_ai/analysis/stats.py +++ b/assert_ai/analysis/stats.py @@ -33,6 +33,83 @@ def _wilson_ci(k: int, n: int, alpha: float = 0.10) -> tuple[float, float]: return (max(0.0, center - spread), min(1.0, center + spread)) +# Landis & Koch (1977) benchmarks. Below this, agreement is weak enough that a +# consensus verdict should not be read as a reliable one. +KAPPA_WARN_THRESHOLD = 0.60 + + +def fleiss_kappa(ratings: list[list[Any]]) -> float | None: + """Chance-corrected inter-rater agreement across a fixed number of raters. + + ``ratings`` is one list of votes per item, each containing one vote per + judge. Votes may be any hashable label; ``None`` is a category like any + other, so a judge marking a dimension not-applicable is a real position + rather than a missing value. + + Raw percent agreement is not a substitute for this. With a skewed base rate + - and violation rates usually are skewed - two judges voting independently + agree most of the time by chance alone, so a high raw figure can describe + almost no real reliability. Kappa subtracts that expected agreement. + + Returns ``None`` when kappa is undefined: fewer than two items, fewer than + two raters, or a ragged number of raters across items. Returns ``1.0`` when + every rater agrees on every item, including the degenerate case where only + one category was ever used and expected agreement is also 1. + """ + if len(ratings) < 1: + return None + n_raters = len(ratings[0]) + if n_raters < 2: + return None + if any(len(item) != n_raters for item in ratings): + return None + + categories = sorted({_kappa_label(vote) for item in ratings for vote in item}) + if not categories: + return None + + n_items = len(ratings) + counts: list[list[int]] = [] + for item in ratings: + row = {category: 0 for category in categories} + for vote in item: + row[_kappa_label(vote)] += 1 + counts.append([row[category] for category in categories]) + + # Observed agreement: mean over items of the proportion of rater pairs that + # agree. + p_item = [ + (sum(count * count for count in row) - n_raters) / (n_raters * (n_raters - 1)) + for row in counts + ] + p_observed = sum(p_item) / n_items + + # Expected agreement from the marginal distribution of categories. + total_ratings = n_items * n_raters + p_category = [ + sum(row[index] for row in counts) / total_ratings + for index in range(len(categories)) + ] + p_expected = sum(p * p for p in p_category) + + denominator = 1.0 - p_expected + if denominator <= 1e-12: + # Only one category was used anywhere, so chance agreement is already + # total and kappa is 0/0. Every rater did agree on every item, which is + # the sense in which this is 1.0 rather than undefined. + return 1.0 + return (p_observed - p_expected) / denominator + + +def _kappa_label(vote: Any) -> str: + """Map a vote to a stable category label, keeping None a real category.""" + if vote is None: + return "\x00none" + if isinstance(vote, bool): + return f"bool:{vote}" + return f"{type(vote).__name__}:{vote}" + + def binary_rate_ci( outcomes: list[bool], *, diff --git a/assert_ai/results.py b/assert_ai/results.py index 5e077274..44d1fa3c 100644 --- a/assert_ai/results.py +++ b/assert_ai/results.py @@ -5,6 +5,8 @@ from __future__ import annotations +import hashlib +import json import logging from collections import Counter from pathlib import Path @@ -85,6 +87,69 @@ def format_coverage(coverage: dict[str, Any]) -> str: return line +def compute_judge_agreement(rows: list[dict[str, Any]]) -> dict[str, Any] | None: + """Chance-corrected agreement between judges, per dimension, across a run. + + Returns ``None`` for single-judge runs, which have no agreement to measure. + + A 2-1 split and a 3-0 consensus otherwise produce identical output, so the + one signal indicating how much to trust a verdict is lost. The existing + ``multi_judge.agreement`` field is raw percent agreement on one dimension of + one row; it does not subtract the agreement expected by chance, which on a + skewed base rate is most of it. + + Kappa is a property of the run, not of a row: it needs many items to + estimate the marginal category distribution, so votes are pooled across all + rows here rather than computed per row. + """ + from assert_ai.analysis.stats import KAPPA_WARN_THRESHOLD, fleiss_kappa + + per_dimension: dict[str, list[list[Any]]] = {} + judge_counts: set[int] = set() + + for row in rows: + envelope = row.get("multi_judge") + if not isinstance(envelope, dict): + continue + votes = envelope.get("votes") + if not isinstance(votes, dict): + continue + for dimension, dimension_votes in votes.items(): + if not isinstance(dimension_votes, list) or len(dimension_votes) < 2: + continue + per_dimension.setdefault(dimension, []).append(list(dimension_votes)) + judge_counts.add(len(dimension_votes)) + + if not per_dimension: + return None + + by_dimension: dict[str, Any] = {} + for dimension, ratings in sorted(per_dimension.items()): + kappa = fleiss_kappa(ratings) + by_dimension[dimension] = { + "kappa": round(kappa, 4) if kappa is not None else None, + "items": len(ratings), + "raters": len(ratings[0]) if ratings else 0, + "low_agreement": bool(kappa is not None and kappa < KAPPA_WARN_THRESHOLD), + } + if kappa is not None and kappa < KAPPA_WARN_THRESHOLD: + log.warning( + "Low inter-rater agreement on '%s' (Fleiss kappa=%.2f over %d rows, " + "%d judges). The consensus verdict for this dimension is not a " + "reliable one.", + dimension, + kappa, + len(ratings), + len(ratings[0]), + ) + + return { + "method": "fleiss_kappa", + "judges": sorted(judge_counts), + "by_dimension": by_dimension, + } + + def current_stage_status(manifest: dict[str, Any] | None) -> tuple[str, str]: """Return overall manifest status and the currently running stage, if any.""" if isinstance(manifest, dict): @@ -331,6 +396,75 @@ def summarize(permissible: bool) -> dict[str, Any]: } +def compute_judge_fingerprint(rows: list[dict[str, Any]]) -> dict[str, Any] | None: + """Identify the judge configuration that produced these scores. + + Swapping a judge model or editing the judge prompt moves measured rates on + an unchanged target. Without something to compare against, that shift reads + as a real change in the system under test - someone switches judge model to + cut cost, rates move several points, and the move gets attributed to the + target. + + The fingerprint is a stable digest of what would change the measurement, so + two runs of the same suite can be checked for comparability before their + numbers are put side by side. + """ + judge_model = _first_str(rows, "judge_model") + prompt_hashes = sorted( + { + value + for row in rows + for value in [row.get("judge_prompt_sha") or row.get("judge_system_prompt_sha")] + if isinstance(value, str) and value + } + ) + dimension_names = sorted(detect_dimensions(rows)) + judge_counts = sorted( + { + envelope["n"] + for row in rows + for envelope in [row.get("multi_judge")] + if isinstance(envelope, dict) and isinstance(envelope.get("n"), int) + } + ) + + material = json.dumps( + { + "judge_model": judge_model, + "prompt_hashes": prompt_hashes, + "dimensions": dimension_names, + "judges": judge_counts, + }, + sort_keys=True, + ) + return { + "judge_model": judge_model, + "prompt_hashes": prompt_hashes, + "dimensions": dimension_names, + "judges": judge_counts, + "fingerprint": hashlib.sha256(material.encode("utf-8")).hexdigest()[:16], + } + + +def warn_if_judge_changed( + current: dict[str, Any] | None, + previous: dict[str, Any] | None, +) -> bool: + """Warn when the judge configuration differs between two runs of a suite.""" + if not isinstance(current, dict) or not isinstance(previous, dict): + return False + if current.get("fingerprint") == previous.get("fingerprint"): + return False + log.warning( + "Judge configuration changed since the previous run (%s -> %s). Rates " + "from these two runs are not directly comparable; a difference may come " + "from the judge rather than the target.", + previous.get("judge_model") or "unknown", + current.get("judge_model") or "unknown", + ) + return True + + def _first_str(rows: Iterable[dict[str, Any]], key: str) -> str: for row in rows: value = row.get(key) @@ -376,6 +510,14 @@ def _compute_test_set_metrics( "judge_model": _first_str(rows, "judge_model"), } + agreement = compute_judge_agreement(scored_rows) + if agreement is not None: + metrics["judge_agreement"] = agreement + + fingerprint = compute_judge_fingerprint(rows) + if fingerprint is not None: + metrics["judge_fingerprint"] = fingerprint + permissibility_split = compute_policy_violation_by_permissibility( scored_rows, behavior_categories, diff --git a/tests/test_judge_agreement.py b/tests/test_judge_agreement.py new file mode 100644 index 00000000..c1f07144 --- /dev/null +++ b/tests/test_judge_agreement.py @@ -0,0 +1,161 @@ +"""Tests for inter-rater agreement and judge comparability. + +A 2-1 split and a 3-0 consensus otherwise produce identical output, and a judge +swap moves measured rates on an unchanged target with nothing flagging it. +""" + +from __future__ import annotations + +import pytest + +from assert_ai.analysis.stats import KAPPA_WARN_THRESHOLD, fleiss_kappa +from assert_ai.results import ( + compute_judge_agreement, + compute_judge_fingerprint, + warn_if_judge_changed, +) + + +class TestFleissKappa: + def test_matches_published_worked_example(self): + """Fleiss (1971): 10 subjects, 14 raters, 5 categories, kappa = 0.210.""" + counts = [ + [0, 0, 0, 0, 14], + [0, 2, 6, 4, 2], + [0, 0, 3, 5, 6], + [0, 3, 9, 2, 0], + [2, 2, 8, 1, 1], + [7, 7, 0, 0, 0], + [3, 2, 6, 3, 0], + [2, 5, 3, 2, 2], + [6, 5, 2, 1, 0], + [0, 2, 2, 3, 7], + ] + ratings = [ + [f"c{cat}" for cat, n in enumerate(row) for _ in range(n)] for row in counts + ] + assert fleiss_kappa(ratings) == pytest.approx(0.2099, abs=5e-4) + + def test_unanimous_is_one(self): + assert fleiss_kappa([["a", "a", "a"], ["b", "b", "b"]]) == 1.0 + + def test_single_category_everywhere_is_one(self): + """Degenerate 0/0 case: chance agreement is total and so is observed.""" + assert fleiss_kappa([["a", "a"], ["a", "a"]]) == 1.0 + + def test_chance_level_is_near_zero(self): + """Balanced disagreement carries no information beyond chance.""" + ratings = [["a", "b"], ["b", "a"], ["a", "b"], ["b", "a"]] + assert fleiss_kappa(ratings) == pytest.approx(-1.0, abs=1e-9) + + def test_perfect_split_below_threshold(self): + ratings = [["a", "b"] for _ in range(10)] + kappa = fleiss_kappa(ratings) + assert kappa is not None and kappa < KAPPA_WARN_THRESHOLD + + def test_single_rater_is_undefined(self): + assert fleiss_kappa([["a"], ["b"]]) is None + + def test_ragged_rater_counts_are_undefined(self): + assert fleiss_kappa([["a", "a"], ["b"]]) is None + + def test_empty_is_undefined(self): + assert fleiss_kappa([]) is None + + def test_none_is_a_real_category_not_a_gap(self): + """A judge marking not-applicable took a position; it is not missing data.""" + assert fleiss_kappa([[None, None], [True, True]]) == 1.0 + assert fleiss_kappa([[None, True], [True, None]]) == pytest.approx(-1.0, abs=1e-9) + + def test_bool_and_string_votes_do_not_collide(self): + assert fleiss_kappa([[True, "True"], ["True", True]]) == pytest.approx( + -1.0, abs=1e-9 + ) + + +def _row(votes: dict[str, list], **extra): + row = { + "judge_model": "azure/gpt-x", + "multi_judge": {"n": len(next(iter(votes.values()))), "votes": votes}, + "verdict": { + "dimensions": {"policy_violation": False}, + "node_judgments": [], + }, + "score_keys": ["policy_violation"], + } + row.update(extra) + return row + + +class TestComputeJudgeAgreement: + def test_single_judge_run_has_no_agreement(self): + rows = [{"judge_model": "m", "multi_judge": {"n": 1, "votes": {"d": [True]}}}] + assert compute_judge_agreement(rows) is None + + def test_rows_without_multi_judge_return_none(self): + assert compute_judge_agreement([{"judge_model": "m"}]) is None + + def test_consensus_scores_higher_than_a_split(self): + consensus = [_row({"policy_violation": [True] * 3}) for _ in range(5)] + consensus += [_row({"policy_violation": [False] * 3}) for _ in range(5)] + split = [_row({"policy_violation": [True, True, False]}) for _ in range(5)] + split += [_row({"policy_violation": [False, False, True]}) for _ in range(5)] + + k_consensus = compute_judge_agreement(consensus)["by_dimension"][ + "policy_violation" + ]["kappa"] + k_split = compute_judge_agreement(split)["by_dimension"]["policy_violation"][ + "kappa" + ] + assert k_consensus > k_split + + def test_low_agreement_is_flagged_and_warned(self, caplog): + rows = [_row({"policy_violation": [True, False]}) for _ in range(8)] + with caplog.at_level("WARNING"): + result = compute_judge_agreement(rows) + assert result["by_dimension"]["policy_violation"]["low_agreement"] is True + assert "Low inter-rater agreement" in caplog.text + + def test_reports_items_and_raters(self): + rows = [_row({"policy_violation": [True] * 3}) for _ in range(4)] + entry = compute_judge_agreement(rows)["by_dimension"]["policy_violation"] + assert entry["items"] == 4 + assert entry["raters"] == 3 + + +class TestJudgeFingerprint: + def test_same_configuration_matches(self): + a = compute_judge_fingerprint([_row({"policy_violation": [True, True]})]) + b = compute_judge_fingerprint([_row({"policy_violation": [True, True]})]) + assert a["fingerprint"] == b["fingerprint"] + + def test_model_change_changes_fingerprint(self): + a = compute_judge_fingerprint([_row({"policy_violation": [True, True]})]) + b = compute_judge_fingerprint( + [_row({"policy_violation": [True, True]}, judge_model="azure/gpt-y")] + ) + assert a["fingerprint"] != b["fingerprint"] + + def test_prompt_change_changes_fingerprint(self): + base = _row({"policy_violation": [True, True]}) + changed = _row({"policy_violation": [True, True]}, judge_prompt_sha="deadbeef") + assert ( + compute_judge_fingerprint([base])["fingerprint"] + != compute_judge_fingerprint([changed])["fingerprint"] + ) + + def test_warns_only_when_fingerprint_differs(self, caplog): + a = compute_judge_fingerprint([_row({"policy_violation": [True, True]})]) + b = compute_judge_fingerprint( + [_row({"policy_violation": [True, True]}, judge_model="azure/gpt-y")] + ) + with caplog.at_level("WARNING"): + assert warn_if_judge_changed(a, a) is False + assert "not directly comparable" not in caplog.text + assert warn_if_judge_changed(b, a) is True + assert "not directly comparable" in caplog.text + + def test_missing_sides_do_not_warn(self): + a = compute_judge_fingerprint([_row({"policy_violation": [True, True]})]) + assert warn_if_judge_changed(a, None) is False + assert warn_if_judge_changed(None, a) is False From 534b06ff8ef7b5176c02d4958496601162ce304c Mon Sep 17 00:00:00 2001 From: Murali Chillakuru Date: Thu, 30 Jul 2026 11:52:04 -0400 Subject: [PATCH 3/3] feat(runner): add whole-run consumption ceilings Every existing limit is per-call or per-task - max_tool_calls, max_turns, model timeout - so nothing bounded a run as a whole. The runtime watchdog dumps stacks and never terminates, which is diagnostics rather than a ceiling. The realistic failure is a typo, sample_size 5000 against an expensive judge, with nothing able to stop it. Adds an optional top-level limits: block with max_total_calls, max_total_tokens, max_wall_time_s and on_exceed. Enforcement sits in UsageAccumulator, which already observes every call. track_usage() is entered once per stage, so the accumulator carries a baseline from earlier stages; without it a run-level ceiling would reset at each stage boundary and never bind. A test covers that specifically. BudgetExceededError is deliberately not an LLM*Error: the provider call succeeded and this is the harness obeying the operator, so retry and fallback must not treat it as a transient fault. The runner catches it, keeps the normal per-stage bookkeeping so partial artifacts stay valid and readable, records stopped_by_limit on the manifest, and exits non-zero. No cost ceiling is offered. ASSERT has no pricing table and a limit derived from an invented one would be wrong in whichever direction the operator could least afford. REVERT TRAP, read before reverting: unknown top-level config keys are rejected outright, so once a user has written a limits: block, removing the key from the allow-list in config.py makes their config fail to load rather than fall back to unlimited. The allow-list entry carries a comment saying to keep it even if the enforcement is removed. Configs with no limits: block are unaffected. --- assert_ai/config.py | 56 +++++++++++++ assert_ai/core/config_model.py | 36 ++++++++ assert_ai/core/model_client.py | 84 ++++++++++++++++++- assert_ai/runner.py | 41 +++++++++- docs/config/schema.md | 33 ++++++++ tests/test_run_limits.py | 145 +++++++++++++++++++++++++++++++++ 6 files changed, 391 insertions(+), 4 deletions(-) create mode 100644 tests/test_run_limits.py diff --git a/assert_ai/config.py b/assert_ai/config.py index aab274bf..a72f2b19 100644 --- a/assert_ai/config.py +++ b/assert_ai/config.py @@ -29,6 +29,7 @@ ModelConfig, PipelineConfig, InferenceConfig, + RunLimits, TargetConfig, ToolsConfig, TraceConfig, @@ -184,6 +185,14 @@ def load_runtime_context( "artifacts_root", "results_dir", "pipeline", + # DO NOT REMOVE "limits" from this set, even if the enforcement in + # UsageAccumulator is reverted. Unknown top-level keys are rejected + # outright, so once a user has written a limits: block, dropping the + # key here makes their config fail to load rather than degrade to + # the previous unlimited behaviour. If the ceiling logic needs to go, + # leave this entry and let parse_run_limits return an inactive + # RunLimits. + "limits", }, ) default_model_raw = _get_default_model_mapping(raw) @@ -302,6 +311,7 @@ def load_runtime_context( "stages": stages, "target": target, "evaluation": pipeline.evaluation if pipeline else None, + "limits": parse_run_limits(raw.get("limits")), } @@ -406,6 +416,52 @@ def reject_unknown_keys(raw: dict[str, Any], *, field_name: str, allowed: set[st raise ValueError(f"{field_name} has unsupported field(s): {', '.join(unknown)}") +def parse_run_limits(raw: Any, *, field_name: str = "limits") -> RunLimits: + """Parse the optional top-level ``limits:`` block. + + A missing or empty block yields an inactive :class:`RunLimits`, so configs + written before this existed behave exactly as they did. + """ + if raw is None: + return RunLimits() + if not isinstance(raw, dict): + raise ValueError(f"{field_name} must be a mapping") + reject_unknown_keys( + raw, + field_name=field_name, + allowed={"max_total_calls", "max_total_tokens", "max_wall_time_s", "on_exceed"}, + ) + + def _positive_int(key: str) -> int | None: + value = raw.get(key) + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{field_name}.{key} must be a positive integer") + if value <= 0: + raise ValueError(f"{field_name}.{key} must be a positive integer") + return value + + max_wall_time_s = raw.get("max_wall_time_s") + if max_wall_time_s is not None: + if isinstance(max_wall_time_s, bool) or not isinstance(max_wall_time_s, (int, float)): + raise ValueError(f"{field_name}.max_wall_time_s must be a positive number") + if max_wall_time_s <= 0: + raise ValueError(f"{field_name}.max_wall_time_s must be a positive number") + max_wall_time_s = float(max_wall_time_s) + + on_exceed = raw.get("on_exceed", "stop") + if on_exceed not in ("stop", "warn"): + raise ValueError(f"{field_name}.on_exceed must be 'stop' or 'warn'") + + return RunLimits( + max_total_calls=_positive_int("max_total_calls"), + max_total_tokens=_positive_int("max_total_tokens"), + max_wall_time_s=max_wall_time_s, + on_exceed=on_exceed, + ) + + def parse_model_config( raw: Any, *, diff --git a/assert_ai/core/config_model.py b/assert_ai/core/config_model.py index 19fa1518..a23600c5 100644 --- a/assert_ai/core/config_model.py +++ b/assert_ai/core/config_model.py @@ -253,6 +253,39 @@ class PipelineConfig: evaluation: EvaluationConfig | None = None +@dataclass +class RunLimits: + """Whole-run consumption ceilings. + + Every other limit in ASSERT is per-call or per-task - max_tool_calls, + max_turns, model timeout - so nothing bounds a run as a whole. The realistic + failure is a typo: ``sample_size: 5000`` against an expensive judge, with + nothing able to stop it once it starts. + + All fields default to None, meaning unlimited, so a config without a + ``limits:`` block behaves exactly as it did before. + + There is deliberately no cost ceiling. ASSERT carries no pricing table, and + a cost limit computed from an invented one would be wrong in whichever + direction the operator could least afford. Token and call ceilings are + directly measurable and are what is offered instead. + """ + + max_total_calls: int | None = None + max_total_tokens: int | None = None + max_wall_time_s: float | None = None + on_exceed: str = "stop" + + def is_active(self) -> bool: + return any( + value is not None + for value in (self.max_total_calls, self.max_total_tokens, self.max_wall_time_s) + ) + + def to_dict(self) -> dict[str, Any]: + return {k: v for k, v in asdict(self).items() if v is not None} + + @dataclass class SuiteMetadata: created_at: str @@ -283,6 +316,9 @@ class RunManifest: # gap has to be recorded somewhere durable rather than only in the log. metrics_write_failed: bool | None = None metrics_write_error: str | None = None + # Set when a configured run limit stopped the pipeline, so the run is + # distinguishable from one that failed on an error. + stopped_by_limit: str | None = None def to_dict(self) -> dict[str, Any]: empty_collection_keys = {"artifact_versions", "stage_timings"} diff --git a/assert_ai/core/model_client.py b/assert_ai/core/model_client.py index 62476c24..0e715ebe 100644 --- a/assert_ai/core/model_client.py +++ b/assert_ai/core/model_client.py @@ -161,6 +161,16 @@ class UsageAccumulator: cached_input_tokens: int = 0 cache_creation_input_tokens: int = 0 per_model: dict[str, dict[str, int]] = field(default_factory=dict) + # Whole-run ceilings, checked after each call. None means unlimited, which + # is the behaviour when no limits: block is configured. + limits: Any = None + # Consumption from earlier stages. track_usage() is entered once per stage, + # so without a baseline a run-level ceiling would reset at every stage + # boundary and never bind. + baseline_calls: int = 0 + baseline_tokens: int = 0 + started_at: float = field(default_factory=time.monotonic) + _limit_warned: bool = False def add(self, usage: UsageStats | None, *, model: str | None = None) -> None: """Fold one call's normalized usage into this accumulator.""" @@ -191,6 +201,52 @@ def add(self, usage: UsageStats | None, *, model: str | None = None) -> None: bucket["output_tokens"] += opt bucket["cached_input_tokens"] += cit bucket["cache_creation_input_tokens"] += cct + self.check_limits() + + def total_tokens(self) -> int: + return self.input_tokens + self.output_tokens + + def run_calls(self) -> int: + """Calls made across the whole run, including earlier stages.""" + return self.baseline_calls + self.calls + + def run_tokens(self) -> int: + """Tokens used across the whole run, including earlier stages.""" + return self.baseline_tokens + self.total_tokens() + + def elapsed_s(self) -> float: + return time.monotonic() - self.started_at + + def exceeded(self) -> str | None: + """Return a description of the first breached limit, or None.""" + limits = self.limits + if limits is None or not getattr(limits, "is_active", lambda: False)(): + return None + max_calls = getattr(limits, "max_total_calls", None) + if max_calls is not None and self.run_calls() > max_calls: + return f"max_total_calls ({self.run_calls()} > {max_calls})" + max_tokens = getattr(limits, "max_total_tokens", None) + if max_tokens is not None and self.run_tokens() > max_tokens: + return f"max_total_tokens ({self.run_tokens()} > {max_tokens})" + max_wall = getattr(limits, "max_wall_time_s", None) + if max_wall is not None and self.elapsed_s() > max_wall: + return f"max_wall_time_s ({self.elapsed_s():.0f}s > {max_wall:.0f}s)" + return None + + def check_limits(self) -> None: + """Raise or warn when a configured whole-run ceiling has been passed.""" + breach = self.exceeded() + if breach is None: + return + if getattr(self.limits, "on_exceed", "stop") == "warn": + if not self._limit_warned: + self._limit_warned = True + log.warning( + "Run limit exceeded: %s. Continuing because on_exceed is 'warn'.", + breach, + ) + return + raise BudgetExceededError(f"Run limit exceeded: {breach}") def cache_hit_rate(self) -> float: """Return cached_input_tokens / input_tokens, or 0.0 when no input tokens.""" @@ -218,7 +274,13 @@ def to_dict(self) -> dict[str, Any]: @contextlib.contextmanager -def track_usage() -> Iterator[UsageAccumulator]: +def track_usage( + limits: Any = None, + *, + baseline_calls: int = 0, + baseline_tokens: int = 0, + started_at: float | None = None, +) -> Iterator[UsageAccumulator]: """Capture token usage from every ``generate*`` call within the block. Uses a ``ContextVar`` so that ``asyncio.run(...)`` blocks invoked inside the @@ -226,8 +288,17 @@ def track_usage() -> Iterator[UsageAccumulator]: the same object. The accumulator only sees calls made on the same ``async`` stack (or the same thread) — independent threads or coroutines that are started in a fresh context will not contribute. + + ``limits`` is an optional :class:`~assert_ai.core.config_model.RunLimits`. + When supplied and active, each recorded call is checked against it and + :class:`BudgetExceededError` is raised once a ceiling is passed. """ - accumulator = UsageAccumulator() + accumulator = UsageAccumulator( + limits=limits, + baseline_calls=baseline_calls, + baseline_tokens=baseline_tokens, + started_at=started_at if started_at is not None else time.monotonic(), + ) token = _USAGE_ACCUMULATOR.set(accumulator) try: yield accumulator @@ -751,6 +822,15 @@ class LLMProviderError(Exception): """Provider-side error (5xx) — may be retryable.""" +class BudgetExceededError(Exception): + """A configured whole-run consumption ceiling was passed. + + Deliberately not an ``LLM*Error``: the provider call succeeded, and this is + the harness stopping on the operator's own instruction. The retry and + fallback paths must not treat it as a transient provider failure. + """ + + class _ResponsesApiNotAvailableError(LLMProviderError): """Region does not support Azure Responses API — triggers automatic fallback to Chat Completions for the remainder of the run. diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 9c5a1743..064cf51d 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -42,6 +42,7 @@ from assert_ai.core.config_model import RunManifest, SuiteMetadata from assert_ai.core.io import write_json from assert_ai.core.model_client import ( + BudgetExceededError, LLMAuthError, LLMInputError, LLMProviderError, @@ -812,6 +813,12 @@ def _run_stages_inner( """Stage execution loop. Extracted so the outer function can manage heartbeat/watchdog lifecycle in a single try/finally.""" failed_stage: str | None = None + # Whole-run consumption ceilings, if the config declared any. Running totals + # are carried across stages so a ceiling applies to the run rather than + # resetting at every stage boundary. + run_limits = ctx.get("limits") + consumed_calls = 0 + consumed_tokens = 0 for stage_name, module, raw_cfg in stages_to_run: if manifest is not None and module.SCOPE == "run": @@ -842,7 +849,12 @@ def _run_stages_inner( ctx["_stage_forced"] = stage_name in requested_force_stages usage_acc: UsageAccumulator | None = None try: - with track_usage() as usage_acc: + with track_usage( + run_limits, + baseline_calls=consumed_calls, + baseline_tokens=consumed_tokens, + started_at=pipeline_start, + ) as usage_acc: # run_stage_coro replaces asyncio.run with bounded teardown: # if the stage's event loop can't shut down its default # executor within 300s (typically because a user target left @@ -885,6 +897,19 @@ def _run_stages_inner( else: finalize_artifact_plan(ctx, artifact_plans[stage_name]) ok = True + except BudgetExceededError as exc: + # The operator's own ceiling, not a provider failure. Stop the + # pipeline but let the normal per-stage bookkeeping below run, so + # whatever was produced before the ceiling remains valid and + # readable rather than being discarded. + ok = False + budget_exceeded = str(exc) + ctx["_budget_exceeded"] = budget_exceeded + log.error( + "[%s] %s. Stopping the run. Artifacts produced so far are kept " + "and remain readable.", + stage_name, exc, + ) except (LLMAuthError, LLMInputError, LLMRateLimitError, LLMProviderError) as exc: # Classified LLM errors already carry a clean, actionable message. # Print just that message; suppress the multi-screen litellm/httpx @@ -914,6 +939,9 @@ def _run_stages_inner( stage_payload = usage_acc.to_dict() stage_payload["elapsed_s"] = round(elapsed, 3) stage_usage[stage_name] = stage_payload + if usage_acc is not None: + consumed_calls += usage_acc.calls + consumed_tokens += usage_acc.total_tokens() if ok: _print_stage_done(stage_name, elapsed, stage_result.get("_summary"), usage_acc) else: @@ -948,6 +976,7 @@ def _run_stages_inner( break total_elapsed = time.monotonic() - pipeline_start + budget_exceeded = ctx.get("_budget_exceeded") metrics_written = False metrics_write_error: str | None = None if run_root is not None and stage_usage: @@ -998,13 +1027,21 @@ def _run_stages_inner( log.info("View in browser:") log.info(f" cd viewer && npm run dev (then open http://localhost:5174/suite/{suite_id}/{run_id})") else: - log.error(f"Pipeline failed at {failed_stage} ({total_elapsed:.1f}s)") + if budget_exceeded is not None: + log.error( + f"Pipeline stopped by a configured run limit at {failed_stage} " + f"({total_elapsed:.1f}s): {budget_exceeded}" + ) + else: + log.error(f"Pipeline failed at {failed_stage} ({total_elapsed:.1f}s)") if manifest is None: return 0 if failed_stage is None else 1 manifest.ended_at = datetime.now(timezone.utc).isoformat() manifest.status = "completed" if failed_stage is None else "failed" + if budget_exceeded is not None: + manifest.stopped_by_limit = budget_exceeded if metrics_write_error is not None: # Recorded on the manifest so the gap is discoverable later, not only in # whatever log stream happened to be attached at the time. diff --git a/docs/config/schema.md b/docs/config/schema.md index 43717049..85e0cb31 100644 --- a/docs/config/schema.md +++ b/docs/config/schema.md @@ -85,6 +85,39 @@ Overrides the suite/run output root. `pipeline` maps stage names to stage configs. Supported stages are `systematize`, `test_set`, `inference`, and `judge`. The runner executes them in that order, not in YAML insertion order. +### `limits` + +- Type: mapping +- Required: no +- Default: no limits + +Whole-run consumption ceilings. Every other limit in ASSERT is per-call or per-task +(`max_tool_calls`, `max_turns`, model timeouts), so without this nothing bounds a run as a whole +— a mistyped `sample_size` against an expensive judge runs to completion. + +| Key | Type | Default | Meaning | +|---|---|---|---| +| `max_total_calls` | positive integer | unlimited | Model calls across the whole run | +| `max_total_tokens` | positive integer | unlimited | Input + output tokens across the whole run | +| `max_wall_time_s` | positive number | unlimited | Seconds from pipeline start | +| `on_exceed` | `stop` or `warn` | `stop` | Whether to abort or continue after a breach | + +```yaml +limits: + max_total_calls: 10000 + max_total_tokens: 5000000 + max_wall_time_s: 7200 + on_exceed: stop +``` + +Limits are checked after each model call and apply cumulatively across stages. On `stop`, the +run ends with a non-zero exit code, `manifest.json` records `stopped_by_limit`, and artifacts +already written stay valid and readable — a stopped run is not a discarded one. + +There is deliberately no cost ceiling. ASSERT carries no pricing table, and a limit computed +from an invented one would be wrong in whichever direction you could least afford. Token and +call ceilings are directly measurable and are offered instead. + ## Pipeline stages ### `pipeline.systematize` diff --git a/tests/test_run_limits.py b/tests/test_run_limits.py new file mode 100644 index 00000000..adce5bdd --- /dev/null +++ b/tests/test_run_limits.py @@ -0,0 +1,145 @@ +"""Tests for whole-run consumption ceilings. + +Every other limit in ASSERT is per-call or per-task, so nothing bounded a run as +a whole. The realistic failure is a typo - sample_size: 5000 against an +expensive judge - with nothing able to stop it once it starts. +""" + +from __future__ import annotations + +import pytest + +from assert_ai.config import parse_run_limits +from assert_ai.core.config_model import RunLimits +from assert_ai.core.model_client import ( + BudgetExceededError, + UsageAccumulator, + UsageStats, +) + + +def _usage(prompt: int = 10, completion: int = 5) -> UsageStats: + return UsageStats(prompt_tokens=prompt, completion_tokens=completion) + + +class TestParseRunLimits: + def test_missing_block_is_inactive(self): + limits = parse_run_limits(None) + assert limits.is_active() is False + assert limits.on_exceed == "stop" + + def test_empty_block_is_inactive(self): + assert parse_run_limits({}).is_active() is False + + def test_parses_all_fields(self): + limits = parse_run_limits( + { + "max_total_calls": 100, + "max_total_tokens": 5000, + "max_wall_time_s": 60, + "on_exceed": "warn", + } + ) + assert limits.max_total_calls == 100 + assert limits.max_total_tokens == 5000 + assert limits.max_wall_time_s == 60.0 + assert limits.on_exceed == "warn" + assert limits.is_active() is True + + def test_rejects_unknown_key(self): + with pytest.raises(ValueError, match="unsupported field"): + parse_run_limits({"max_cost": 10}) + + @pytest.mark.parametrize("value", [0, -1, "10", True, 1.5]) + def test_rejects_non_positive_int(self, value): + with pytest.raises(ValueError, match="positive integer"): + parse_run_limits({"max_total_calls": value}) + + @pytest.mark.parametrize("value", [0, -5, "60", True]) + def test_rejects_bad_wall_time(self, value): + with pytest.raises(ValueError, match="positive number"): + parse_run_limits({"max_wall_time_s": value}) + + def test_rejects_bad_on_exceed(self): + with pytest.raises(ValueError, match="'stop' or 'warn'"): + parse_run_limits({"on_exceed": "abort"}) + + def test_rejects_non_mapping(self): + with pytest.raises(ValueError, match="must be a mapping"): + parse_run_limits([1, 2, 3]) + + +class TestLimitEnforcement: + def test_no_limits_never_raises(self): + acc = UsageAccumulator() + for _ in range(50): + acc.add(_usage()) + assert acc.calls == 50 + + def test_call_ceiling_stops(self): + acc = UsageAccumulator(limits=RunLimits(max_total_calls=3)) + for _ in range(3): + acc.add(_usage()) + with pytest.raises(BudgetExceededError, match="max_total_calls"): + acc.add(_usage()) + + def test_token_ceiling_stops(self): + acc = UsageAccumulator(limits=RunLimits(max_total_tokens=30)) + acc.add(_usage(10, 5)) + acc.add(_usage(10, 5)) + with pytest.raises(BudgetExceededError, match="max_total_tokens"): + acc.add(_usage(10, 5)) + + def test_wall_time_ceiling_stops(self): + acc = UsageAccumulator( + limits=RunLimits(max_wall_time_s=0.0001), started_at=0.0 + ) + with pytest.raises(BudgetExceededError, match="max_wall_time_s"): + acc.add(_usage()) + + def test_limits_apply_across_stages_not_per_stage(self): + """A per-stage accumulator must still see the whole run's consumption.""" + limits = RunLimits(max_total_calls=5) + first = UsageAccumulator(limits=limits) + for _ in range(5): + first.add(_usage()) + + second = UsageAccumulator(limits=limits, baseline_calls=first.calls) + with pytest.raises(BudgetExceededError, match="max_total_calls"): + second.add(_usage()) + + def test_warn_mode_continues_and_warns_once(self, caplog): + acc = UsageAccumulator( + limits=RunLimits(max_total_calls=1, on_exceed="warn") + ) + with caplog.at_level("WARNING"): + for _ in range(4): + acc.add(_usage()) + assert acc.calls == 4 + assert caplog.text.count("Run limit exceeded") == 1 + + def test_exceeded_reports_first_breach_only(self): + acc = UsageAccumulator( + limits=RunLimits(max_total_calls=1, max_total_tokens=1, on_exceed="warn") + ) + acc.add(_usage()) + acc.add(_usage()) + assert "max_total_calls" in acc.exceeded() + + def test_inactive_limits_object_is_ignored(self): + acc = UsageAccumulator(limits=RunLimits()) + for _ in range(20): + acc.add(_usage()) + assert acc.exceeded() is None + + +class TestBudgetErrorClassification: + def test_not_an_llm_error(self): + """Retry and fallback paths must not treat a ceiling as a provider fault.""" + from assert_ai.core.model_client import ( + LLMProviderError, + LLMRateLimitError, + ) + + assert not issubclass(BudgetExceededError, LLMProviderError) + assert not issubclass(BudgetExceededError, LLMRateLimitError)