feat(flows): aggregate token/cost usage in batch_run and enforce budgets#132
Open
dhruvraajeev wants to merge 1 commit into
Open
feat(flows): aggregate token/cost usage in batch_run and enforce budgets#132dhruvraajeev wants to merge 1 commit into
dhruvraajeev wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Populates the previously dead
BatchResult.tokens_total/cost_estimate_usdfields and enforces the declared-but-unenforcedBaseFlowCfg.max_total_input_tokens/max_total_cost_usdguardrails. This addresses the batch case called out in #127 as most painful —batch_runfanning out many papers with zero cost visibility.Approach
Usage is computed by the SDK as
RunResult.context_wrapper.usage, butrun_with_observabilityreturns onlyfinal_outputand drops the rest, and flows return domain objects — sobatch_runnever sees it. Rather than change every flow's return type (which would break theCallable[..., Awaitable[OutputT]]contractbatch_rundepends on and every flow's public signature), this surfaces usage across that existing seam with acontextvars-scoped accumulator:quantmind/flows/_usage.pyholdsUsageSummary,PriceRate, ausage_scope()context manager, andrecord_usage().run_with_observabilitycallsrecord_usage(result.context_wrapper.usage)— a no-op unless a scope is active, so every existing caller is unchanged.batch_runopens oneusage_scopeper input.asynciocopies the context (and the mutable accumulator reference) into each child task, so usage from nestedgather/wait_forfan-outs (e.g._paper_summary, which calls the seam twice per input) folds into the one scope. All mutation is betweenawaitpoints 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 crossesmax_total_input_tokensormax_total_cost_usd,batch_runstops launching new work and records aBudgetExceededErrorfor 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/retrievalusesRunner.rundirectly rather than the flows seam, so its usage is not captured here. Left as a follow-up.cost_estimate_usdstays0.0and only the token guardrail applies.Tests and verification
tests/flows/test_usage.py: scope isolation across concurrent tasks, accumulation across nestedgather, and no-op outside a scope.tests/flows/test_batch.py: aggregation across inputs, priced cost, and token / cost budget trips (skip and raise modes). Existingtokens_total == {}/cost == 0.0assertions still hold for flows that record no usage.examples/flows/batch_usage.pyand an updateddocs/README.mdcatalog row.bash scripts/verify.shpasses 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).