Skip to content

Commit a09545a

Browse files
Akkariclaude
andcommitted
feat: cost tracking and estimation — per-call tracking, budget caps, pre-run estimates (Milestone #16)
- quorum/cost.py: CostTracker (thread-safe, thread-local file context), CallRecord, CostSummary, CostEstimate Pydantic models, BudgetExceededError, estimate_cost() with hardcoded rates for Claude/GPT/Gemini/Mistral, _FALLBACK_RATE for unknowns - providers/base.py: BaseProvider.__init__ accepts optional cost_tracker - providers/litellm_provider.py: extracts token usage + calls litellm.completion_cost() after each complete() call; handles None/exception from completion_cost gracefully - config.py: max_cost: float | None = None field on QuorumConfig - pipeline.py: CostTracker created per run_validation() and shared across run_batch_validation() workers; budget check after critics (single) and after each file (batch) via BudgetExceededError; cost added to run-manifest.json and batch-manifest.json with per-file breakdown; progressive running_cost_usd in batch manifests - cli.py: --max-cost USD flag, --yes/-y flag, pre-run cost estimate with conditional confirm prompt (skipped when < $0.50 or --yes), post-run cost summary (Cost: $X.XXXX, per-file breakdown) - tests/test_cost.py: 33 tests covering thread safety, budget caps, model rates, estimate accuracy, provider integration, graceful error handling Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 6e8c553 commit a09545a

7 files changed

Lines changed: 968 additions & 4 deletions

File tree

reference-implementation/quorum/cli.py

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,19 @@ def cli(verbose: bool) -> None:
125125
default=False,
126126
help="Disable learning memory for this run (do not read or write known_issues.json)",
127127
)
128+
@click.option(
129+
"--max-cost",
130+
default=None,
131+
type=float,
132+
metavar="USD",
133+
help="Budget cap in USD. Stops batch after each file if total spend exceeds this.",
134+
)
135+
@click.option(
136+
"--yes", "-y",
137+
is_flag=True,
138+
default=False,
139+
help="Skip pre-run cost estimate confirmation prompt.",
140+
)
128141
def run_cmd(
129142
target: str | None,
130143
pattern: str | None,
@@ -137,6 +150,8 @@ def run_cmd(
137150
verbose: bool,
138151
no_learning: bool,
139152
resume: Path | None,
153+
max_cost: float | None,
154+
yes: bool,
140155
) -> None:
141156
"""
142157
Validate artifacts against a rubric.
@@ -221,6 +236,10 @@ def run_cmd(
221236
if fix_loops is not None:
222237
quorum_config = quorum_config.with_overrides(max_fix_loops=fix_loops)
223238

239+
# Apply --max-cost override if provided
240+
if max_cost is not None:
241+
quorum_config = quorum_config.with_overrides(max_cost=max_cost)
242+
224243
# Resolve targets to determine single vs batch mode
225244
target_path = Path(target)
226245
is_batch = (
@@ -241,6 +260,9 @@ def run_cmd(
241260
if relationships:
242261
click.echo(f"Phase 2: cross-artifact validation enabled ({relationships})", err=True)
243262

263+
# Pre-run cost estimate
264+
_show_cost_estimate_and_confirm(files, quorum_config, yes)
265+
244266
batch_verdict, batch_dir = run_batch_validation(
245267
target=target,
246268
pattern=pattern,
@@ -253,6 +275,7 @@ def run_cmd(
253275
# Note: learning memory is not applied per-file in batch mode
254276

255277
print_batch_verdict(batch_verdict, batch_dir=batch_dir, verbose=verbose)
278+
_print_batch_cost_summary(batch_dir)
256279

257280
if batch_verdict.is_actionable:
258281
sys.exit(2)
@@ -290,6 +313,7 @@ def run_cmd(
290313
_print_learning_summary(run_dir)
291314

292315
print_verdict(verdict, run_dir=run_dir, verbose=verbose)
316+
_print_run_cost_summary(run_dir)
293317

294318
if verdict.is_actionable:
295319
sys.exit(2)
@@ -598,3 +622,87 @@ def _first_run_setup(force: bool = False) -> None:
598622
click.echo(f"✓ Configuration written to {config_path}")
599623
click.echo(f" Run: quorum run --target <your-file>")
600624
click.echo()
625+
626+
627+
def _show_cost_estimate_and_confirm(
628+
files: list,
629+
config: object,
630+
skip_prompt: bool,
631+
) -> None:
632+
"""
633+
Show a pre-run cost estimate and optionally prompt to continue.
634+
635+
Skips the prompt when:
636+
- --yes flag is set
637+
- estimate is below $0.50 (low-cost run, no need to interrupt)
638+
- not running in an interactive terminal
639+
"""
640+
try:
641+
from quorum.cost import estimate_cost
642+
estimate = estimate_cost(files, config)
643+
click.echo(
644+
f"Estimated cost: ~${estimate.estimated_usd:.2f} "
645+
f"({estimate.estimated_calls} critic calls across {estimate.files_count} files) "
646+
f"[approximate]",
647+
err=True,
648+
)
649+
650+
if skip_prompt or estimate.estimated_usd < 0.50 or not sys.stdin.isatty():
651+
return
652+
653+
if not click.confirm("Proceed with validation?", default=True):
654+
click.echo("Aborted.", err=True)
655+
sys.exit(0)
656+
except Exception:
657+
pass # Estimate is best-effort — never block a run
658+
659+
660+
def _print_run_cost_summary(run_dir: Path) -> None:
661+
"""Print cost summary from a single-file run manifest."""
662+
import json as _json
663+
manifest_path = run_dir / "run-manifest.json"
664+
if not manifest_path.exists():
665+
return
666+
try:
667+
data = _json.loads(manifest_path.read_text())
668+
cost = data.get("cost")
669+
if not cost:
670+
return
671+
_emit_cost_line(cost)
672+
except Exception:
673+
pass
674+
675+
676+
def _print_batch_cost_summary(batch_dir: Path) -> None:
677+
"""Print cost summary from a batch manifest."""
678+
import json as _json
679+
manifest_path = batch_dir / "batch-manifest.json"
680+
if not manifest_path.exists():
681+
return
682+
try:
683+
data = _json.loads(manifest_path.read_text())
684+
cost = data.get("cost")
685+
if not cost:
686+
return
687+
_emit_cost_line(cost)
688+
except Exception:
689+
pass
690+
691+
692+
def _emit_cost_line(cost: dict) -> None:
693+
"""Emit the formatted cost line(s) to stderr."""
694+
total = cost.get("total_usd", 0.0)
695+
prompt_t = cost.get("prompt_tokens", 0)
696+
completion_t = cost.get("completion_tokens", 0)
697+
calls = cost.get("calls", 0)
698+
click.echo(
699+
f"Cost: ${total:.4f} "
700+
f"({prompt_t:,} prompt + {completion_t:,} completion tokens across {calls} calls)",
701+
err=True,
702+
)
703+
per_file = cost.get("per_file", {})
704+
if per_file:
705+
# Sort by cost descending, show top 5
706+
sorted_files = sorted(per_file.items(), key=lambda x: -x[1])[:5]
707+
parts = " | ".join(f"{Path(fp).name} ${c:.4f}" for fp, c in sorted_files)
708+
click.echo(f"Per-file: {parts}", err=True)

reference-implementation/quorum/config.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,10 @@ class QuorumConfig(BaseModel):
5959
default=True,
6060
description="Run deterministic pre-screen checks before LLM critics",
6161
)
62+
max_cost: float | None = Field(
63+
default=None,
64+
description="Maximum allowed LLM spend in USD. Stops batch after each file if exceeded.",
65+
)
6266

6367
@field_validator("critics")
6468
@classmethod

0 commit comments

Comments
 (0)