Skip to content

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

Description

@keli-wen

✨ Feature Summary

Surface token usage (and optionally cost) for each QuantMind LLM run, and aggregate it across a batch_run, so callers can see what a build / extract / retrieve consumed. Nothing exposes this today.

🎯 Motivation

QuantMind discards the RunResult.context_wrapper.usage that the Agents SDK already computes: the flow and retrieval callables return domain objects (PaperStructureTree, list[RetrievalEvidence], ...) and drop the usage. The max_*_tokens cfg fields are generation / input bounds, not accounting, and BaseFlowCfg.max_total_input_tokens / max_total_cost_usd are declared but unenforced. This is most painful with batch_run, which fans out many papers with zero cost visibility. Getting a real number today requires manually patching Runner.run to read context_wrapper.usage.

πŸ“‹ Detailed Description

Aggregate context_wrapper.usage (input / output / total tokens and request count) per run and across a batch, and make it available to callers without forcing them to hook the SDK. Keep model pricing out of the library (report tokens; let callers apply per-model prices) or accept an optional price table. Optionally enforce the existing max_total_input_tokens / max_total_cost_usd guardrails from the same accounting.

πŸ”§ Proposed Implementation

A small usage-accounting run-hook (or a returned usage summary) layered on the existing flows/_runner observability seam, accumulating SDK usage per run and per batch.

API Design

tree = await PaperFlow(cfg).build(input)
usage  # UsageSummary(requests, input_tokens, output_tokens, total_tokens) via a run-hook / accessor

batch = await batch_run(flow.build, inputs)
batch.usage  # aggregated UsageSummary across the fan-out

🎨 User Experience

A caller reads input_tokens / output_tokens / requests after any run and after a batch, and optionally sets max_total_input_tokens / max_total_cost_usd to fail closed when a budget is exceeded.

πŸ“Š Use Cases

  1. Cost visibility for batch ingestion: a batch_run over N papers reports aggregate tokens so an operator can estimate spend.
  2. Budget guardrails: enforce max_total_input_tokens / max_total_cost_usd per run or per batch.
  3. Model comparison: compare token / turn counts across models (an agentic retrieve converging in 3 turns vs 10) to pick a cheaper default.

πŸ”— Related Issues

Implementation Considerations

Pricing changes over time and varies by provider, so the library should avoid hard-coding prices. Usage is already produced by the SDK, so this is aggregation and surfacing, not new model plumbing.

Breaking Changes

  • This feature would introduce breaking changes
  • This feature is backward compatible

Dependencies

  • Requires new dependencies
  • Uses existing dependencies only

Checklist

  • I have searched existing issues to avoid duplicates
  • I have provided a clear and detailed description
  • I have explained the motivation and use cases
  • I have considered the implementation approach
  • I have thought about potential breaking changes

Metadata

Metadata

Assignees

No one assigned

    Labels

    area: configsTyped inputs and configuration contractsarea: flowsPublic operation implementations under quantmind/flows/area: mindMemory, tools, MCP integration, and the cognitive layertype: featureAdds a new capability or observable behavior

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions