Skip to content

feat(flows): aggregate token/cost usage in batch_run and enforce budgets#132

Open
dhruvraajeev wants to merge 1 commit into
LLMQuant:masterfrom
dhruvraajeev:feat/batch-usage-accounting
Open

feat(flows): aggregate token/cost usage in batch_run and enforce budgets#132
dhruvraajeev wants to merge 1 commit into
LLMQuant:masterfrom
dhruvraajeev:feat/batch-usage-accounting

Conversation

@dhruvraajeev

Copy link
Copy Markdown

Summary

Populates the previously dead BatchResult.tokens_total / cost_estimate_usd fields and enforces the declared-but-unenforced BaseFlowCfg.max_total_input_tokens / max_total_cost_usd guardrails. This addresses the batch case called out in #127 as most painful — batch_run fanning out many papers with zero cost visibility.

Approach

Usage is computed by the SDK as RunResult.context_wrapper.usage, but run_with_observability returns only final_output and drops the rest, and flows return domain objects — so batch_run never sees it. Rather than change every flow's return type (which would break the Callable[..., Awaitable[OutputT]] contract batch_run depends on and every flow's public signature), this surfaces usage across that existing seam with a contextvars-scoped accumulator:

  • New quantmind/flows/_usage.py holds UsageSummary, PriceRate, a usage_scope() context manager, and record_usage().
  • run_with_observability calls record_usage(result.context_wrapper.usage) — a no-op unless a scope is active, so every existing caller is unchanged.
  • batch_run opens one usage_scope per input. asyncio copies the context (and the mutable accumulator reference) into each child task, so usage from nested gather / wait_for fan-outs (e.g. _paper_summary, which calls the seam twice per input) folds into the one scope. All mutation is between await points on the single-threaded loop, so no lock is needed.

Cost is caller-priced through an optional prices: dict[str, PriceRate] table (per-1M-token rates); the library ships no model prices. When a running total crosses max_total_input_tokens or max_total_cost_usd, batch_run stops launching new work and records a BudgetExceededError for skipped inputs (post-hoc gate — under concurrency spend can't be pre-checked; the upgrade path is a pre-flight token estimate, noted in a code comment).

Scope

  • mind/retrieval uses Runner.run directly rather than the flows seam, so its usage is not captured here. Left as a follow-up.
  • Cost is reported only when a price table is supplied; otherwise cost_estimate_usd stays 0.0 and only the token guardrail applies.

Tests and verification

  • New tests/flows/test_usage.py: scope isolation across concurrent tasks, accumulation across nested gather, and no-op outside a scope.
  • Extended tests/flows/test_batch.py: aggregation across inputs, priced cost, and token / cost budget trips (skip and raise modes). Existing tokens_total == {} / cost == 0.0 assertions still hold for flows that record no usage.
  • New examples/flows/batch_usage.py and an updated docs/README.md catalog row.
  • bash scripts/verify.sh passes locally: 403 passed, coverage 85.56% (floor 75%), with ruff, basedpyright, and import-linter all green.

Closes #127 (batch usage + guardrails slice; per-run tracing in #130 and the dashboard hook in #83 are complementary surfaces).

Populate the previously dead BatchResult.tokens_total and
cost_estimate_usd fields and enforce the declared-but-unenforced
max_total_input_tokens / max_total_cost_usd guardrails.

Usage is surfaced across the run_with_observability seam via a
contextvars-scoped accumulator (quantmind/flows/_usage.py): the runner
records each run's SDK usage into the active scope, and batch_run opens
one scope per input, so per-input usage aggregates without changing any
flow's return type. Cost is caller-priced via an optional prices table
so the library ships no model prices. When a running total crosses a
budget, batch_run stops launching new work and marks skipped inputs with
BudgetExceededError.

Adds tests/flows/test_usage.py, extends tests/flows/test_batch.py, adds
examples/flows/batch_usage.py, and updates the docs/README.md catalog.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(flows): surface token-usage and cost for runs and batch_run

1 participant