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

docs/README.md

@@ -18,7 +18,7 @@ harness.
 | Paper structure build | `quantmind.flows.PaperFlow` | `PaperFlow(PaperStructureCfg)`; `build()`: `PaperInput` | `PaperStructureTree` (self-contained) | [Build and retrieve](../examples/mind/paper_structure_retrieval.py) | [Structure retrieval design](../contexts/design/mind/retrieval.md) |
 | Reasoning-based retrieval (agentic) | `quantmind.mind.AgenticRetriever` | `AgenticRetriever(RetrievalCfg)`; `retrieve()`: one `StructureTree` + question (no library) | `list[RetrievalEvidence]` | [Build and retrieve](../examples/mind/paper_structure_retrieval.py) | [Structure retrieval design](../contexts/design/mind/retrieval.md) |
 | News collection | `quantmind.flows.collect_news` | `NewsWindow`, `NewsCollectionCfg` | `NewsBatch` from `quantmind.preprocess` | [Collect news](../examples/flows/collect_news.py) | [News collection design](../contexts/design/flow/news.md) |
-| Bounded fan-out | `quantmind.flows.batch_run` | Operation inputs and shared config | `BatchResult` | [README usage](../README.md#-usage-examples) | API docstrings |
+| Bounded fan-out | `quantmind.flows.batch_run` | Operation inputs, shared config, and optional `prices` table | `BatchResult` (with aggregate `tokens_total` / `cost_estimate_usd` and budget guardrails) | [Batch usage and budgets](../examples/flows/batch_usage.py) | API docstrings |
 | Local semantic search | `quantmind.library.LocalKnowledgeLibrary` | `BaseKnowledge` or `PaperFlowResult`, `SemanticQuery` | `list[SemanticHit]` | [Library example](../examples/library/README.md) | [Library guide](library.md) |
 | Page-aware document RAG | `quantmind.rag.chunk_parsed_document`, `quantmind.rag.retrieve_parsed_document` | `ParsedDocument`, splitter config, and query | `tuple[ParsedDocumentHit, ...]` | [Paper RAG](../examples/rag/paper.py) | [Document RAG design](../contexts/design/rag/document.md) |
 

examples/flows/batch_usage.py

@@ -0,0 +1,52 @@
+"""Fan out over papers, then read aggregate token usage and priced cost.
+
+``batch_run`` now reports the SDK token usage it consumed and, when given a
+price table, an estimated USD cost — and enforces the ``cfg`` budget
+guardrails, marking any input skipped after the budget trips with a
+``BudgetExceededError``. Requires ``OPENAI_API_KEY`` (like the other flow
+examples).
+"""
+
+import asyncio
+
+from quantmind.configs import PaperFlowCfg
+from quantmind.configs.paper import ArxivIdentifier
+from quantmind.flows import (
+    BudgetExceededError,
+    PriceRate,
+    batch_run,
+    paper_flow,
+)
+
+
+async def main() -> None:
+    """Build several papers under one shared budget and report spend."""
+    cfg = PaperFlowCfg(
+        model="gpt-4o-mini",
+        max_total_input_tokens=200_000,  # stop launching once we cross this
+    )
+    # Caller-supplied pricing (USD per 1M tokens); the library ships none.
+    prices = {
+        "gpt-4o-mini": PriceRate(input_usd_per_1m=0.15, output_usd_per_1m=0.60),
+    }
+    inputs = [
+        ArxivIdentifier(id="1706.03762v7"),
+        ArxivIdentifier(id="2404.11584"),
+    ]
+
+    result = await batch_run(
+        paper_flow, inputs, cfg=cfg, concurrency=2, prices=prices
+    )
+
+    print(f"success={result.success_count} failure={result.failure_count}")
+    print(f"tokens={result.tokens_total}")
+    print(f"cost_usd≈{result.cost_estimate_usd:.4f}")
+    skipped = [
+        i for i, e in result.errors if isinstance(e, BudgetExceededError)
+    ]
+    if skipped:
+        print(f"budget-skipped inputs: {skipped}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

quantmind/flows/__init__.py

@@ -12,15 +12,19 @@
 - ``paper_flow`` is a thin compatibility function for the semantic
   chunk/summary shape (``PaperFlowResult``).
 - ``batch_run`` runs any flow over a list of inputs with bounded
-  concurrency and aggregated results.
-- ``BatchResult`` is the shape returned by ``batch_run``.
+  concurrency and aggregated results, reporting aggregate token usage
+  (and optionally priced cost) and enforcing the ``cfg`` budget guardrails.
+- ``BatchResult`` is the shape returned by ``batch_run``; ``UsageSummary``
+  and ``PriceRate`` describe its usage/cost fields, and
+  ``BudgetExceededError`` marks inputs skipped after a budget tripped.
 - ``UnsupportedContentTypeError`` is raised when a paper pipeline does not
   resolve a page-aware PDF.
 - ``PaperStructureError`` is raised when structure building exceeds its
   runtime boundary.
 """
 
-from quantmind.flows.batch import BatchResult, batch_run
+from quantmind.flows._usage import PriceRate, UsageSummary
+from quantmind.flows.batch import BatchResult, BudgetExceededError, batch_run
 from quantmind.flows.news import collect_news
 from quantmind.flows.paper import (
     PaperFlow,
@@ -32,10 +36,13 @@
 
 __all__ = [
     "BatchResult",
+    "BudgetExceededError",
     "PaperCitationValidationError",
     "PaperFlow",
     "PaperStructureError",
+    "PriceRate",
     "UnsupportedContentTypeError",
+    "UsageSummary",
     "batch_run",
     "collect_news",
     "paper_flow",

quantmind/flows/_runner.py

@@ -12,6 +12,7 @@
 from agents import Agent, RunConfig, RunHooks, Runner
 
 from quantmind.configs import BaseFlowCfg
+from quantmind.flows._usage import record_usage
 
 
 async def run_with_observability(
@@ -55,6 +56,9 @@ async def run_with_observability(
         hooks=hooks,
         max_turns=cfg.max_turns,
     )
+    # No-op unless a `usage_scope` is active (e.g. inside `batch_run`);
+    # keeps every existing caller's behaviour unchanged.
+    record_usage(result.context_wrapper.usage)
     _archive_run_artifacts(cfg, memory, result)
     return result.final_output
 

quantmind/flows/_usage.py

@@ -0,0 +1,112 @@
+"""Per-run usage accounting layered on the flows observability seam.
+
+The Agents SDK computes token usage as ``RunResult.context_wrapper.usage``,
+but ``run_with_observability`` returns only ``final_output`` and drops the
+rest. Flows return domain objects, so ``batch_run`` never sees usage.
+
+This module surfaces usage across that seam without changing any flow's
+return type. ``run_with_observability`` calls ``record_usage`` after every
+run; a caller opens a ``usage_scope`` around the work it wants measured, and
+each run folds its usage into the active accumulator. ``asyncio`` copies the
+context (and with it the accumulator *reference*) into every child task, so
+usage from nested ``gather`` / ``wait_for`` fan-outs accumulates into the one
+scope the caller opened. Mutation has no ``await`` points, so no lock is
+needed on the single-threaded event loop.
+"""
+
+import contextlib
+from collections.abc import Iterator
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class PriceRate:
+    """Per-token USD pricing for one model (rates are per one million tokens)."""
+
+    input_usd_per_1m: float
+    output_usd_per_1m: float
+
+    def cost(self, input_tokens: int, output_tokens: int) -> float:
+        """USD cost for the given input / output token counts."""
+        return (
+            input_tokens / 1_000_000 * self.input_usd_per_1m
+            + output_tokens / 1_000_000 * self.output_usd_per_1m
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class UsageSummary:
+    """Immutable token-usage snapshot returned to callers.
+
+    ``cost_usd`` is ``None`` unless the caller supplied a price table; the
+    library reports tokens and leaves per-model pricing to the caller.
+    """
+
+    requests: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    cost_usd: float | None = None
+
+    def as_tokens_dict(self) -> dict[str, int]:
+        """Return the token counts as a plain dict (no cost)."""
+        return {
+            "requests": self.requests,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "total_tokens": self.total_tokens,
+        }
+
+
+@dataclass(slots=True)
+class _Accumulator:
+    """Mutable running total; duck-types the SDK ``Usage`` object."""
+
+    requests: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+
+    def add(self, usage: object) -> None:
+        """Fold one SDK ``Usage`` (or ``UsageSummary``) into the total."""
+        self.requests += getattr(usage, "requests", 0)
+        self.input_tokens += getattr(usage, "input_tokens", 0)
+        self.output_tokens += getattr(usage, "output_tokens", 0)
+        self.total_tokens += getattr(usage, "total_tokens", 0)
+
+    def summary(self) -> UsageSummary:
+        """Snapshot the running total as an immutable ``UsageSummary``."""
+        return UsageSummary(
+            requests=self.requests,
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+            total_tokens=self.total_tokens,
+        )
+
+
+_usage_var: ContextVar[_Accumulator | None] = ContextVar(
+    "quantmind_usage", default=None
+)
+
+
+@contextlib.contextmanager
+def usage_scope() -> Iterator[_Accumulator]:
+    """Accumulate usage from every run inside this context.
+
+    Yields the accumulator; read ``.summary()`` after the block. Nested
+    ``usage_scope`` calls each measure only their own runs.
+    """
+    accumulator = _Accumulator()
+    token = _usage_var.set(accumulator)
+    try:
+        yield accumulator
+    finally:
+        _usage_var.reset(token)
+
+
+def record_usage(usage: object) -> None:
+    """Fold one run's SDK usage into the active scope (no-op if none)."""
+    accumulator = _usage_var.get()
+    if accumulator is not None:
+        accumulator.add(usage)

quantmind/flows/batch.py

@@ -13,11 +13,16 @@
 from typing import Any, Generic, Literal, TypeVar
 
 from quantmind.configs import BaseFlowCfg, BaseInput
+from quantmind.flows._usage import PriceRate, usage_scope
 
 InputT = TypeVar("InputT", bound=BaseInput)
 OutputT = TypeVar("OutputT")
 
 
+class BudgetExceededError(Exception):
+    """Raised for inputs skipped after a batch token / cost budget tripped."""
+
+
 @dataclass(slots=True)
 class BatchResult(Generic[OutputT]):
     """Aggregate result of running a flow over many inputs.
@@ -56,6 +61,7 @@ async def batch_run(
     concurrency: int = 4,
     on_error: Literal["raise", "skip"] = "skip",
     on_progress: Callable[[int, int], None] | None = None,
+    prices: dict[str, PriceRate] | None = None,
     **flow_kwargs: Any,
 ) -> BatchResult[OutputT]:
     """Run ``flow_fn`` over ``inputs`` with bounded concurrency.
@@ -75,19 +81,28 @@ async def batch_run(
             completion (success or failure). Must be cheap and
             non-blocking — callbacks are invoked synchronously inside
             the worker loop.
+        prices: Optional ``{model: PriceRate}`` table. When it covers
+            ``cfg.model``, ``cost_estimate_usd`` is filled and the
+            ``max_total_cost_usd`` guardrail is enforced; otherwise cost
+            stays ``0.0`` and only the token guardrail applies. Pricing is
+            caller-supplied so the library never ships model prices.
         **flow_kwargs: Forwarded verbatim to ``flow_fn``. ``memory=`` is
             **forbidden** in MVP; passing it raises ``ValueError``.
 
     Returns:
         ``BatchResult`` with ``results`` parallel to ``inputs`` (None for
-        failures) and ``errors`` sorted by index.
+        failures) and ``errors`` sorted by index. ``tokens_total`` holds
+        aggregate SDK usage (empty when nothing was recorded) and
+        ``cost_estimate_usd`` the priced total.
 
     Raises:
         ValueError: If ``memory=`` is passed via ``flow_kwargs``, or if
             ``concurrency < 1``.
         Exception: Re-raised when ``on_error="raise"`` and any input
-            fails. The exception is the first one raised by a worker;
-            other workers may already be cancelled when this surfaces.
+            fails (including a ``BudgetExceededError`` for an input skipped
+            after ``cfg.max_total_input_tokens`` / ``max_total_cost_usd``
+            tripped). Other workers may already be cancelled when this
+            surfaces.
     """
     if "memory" in flow_kwargs:
         raise ValueError(
@@ -105,11 +120,52 @@ async def batch_run(
     started = time.monotonic()
     done_counter = 0
 
+    price = prices.get(cfg.model) if prices is not None and cfg else None
+    max_tokens = cfg.max_total_input_tokens if cfg else None
+    max_cost = cfg.max_total_cost_usd if cfg else None
+    # asyncio is single-threaded; `.add` and these reads have no `await`
+    # between them, so this shared running total needs no lock.
+    running = {
+        "requests": 0,
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "total_tokens": 0,
+    }
+    running_cost = 0.0
+    budget_tripped = False
+
     async def run_one(i: int, inp: InputT) -> None:
-        nonlocal done_counter
+        nonlocal done_counter, running_cost, budget_tripped
         async with sem:
             try:
-                results[i] = await flow_fn(inp, cfg=cfg, **flow_kwargs)
+                # ponytail: post-hoc gate — under concurrency we can't
+                # pre-check spend before it happens; we stop *launching*
+                # new work once the budget trips. Upgrade path: pre-flight
+                # token estimate per input if strict caps matter.
+                if budget_tripped:
+                    raise BudgetExceededError(
+                        f"input {i} skipped: batch budget already exceeded"
+                    )
+                with usage_scope() as acc:
+                    results[i] = await flow_fn(inp, cfg=cfg, **flow_kwargs)
+                summary = acc.summary()
+                running["requests"] += summary.requests
+                running["input_tokens"] += summary.input_tokens
+                running["output_tokens"] += summary.output_tokens
+                running["total_tokens"] += summary.total_tokens
+                if price is not None:
+                    running_cost = price.cost(
+                        running["input_tokens"], running["output_tokens"]
+                    )
+                if (
+                    max_tokens is not None
+                    and running["input_tokens"] > max_tokens
+                ) or (
+                    max_cost is not None
+                    and price is not None
+                    and running_cost > max_cost
+                ):
+                    budget_tripped = True
             except Exception as exc:
                 errors.append((i, exc))
                 if on_error == "raise":
@@ -133,4 +189,6 @@ async def run_one(i: int, inp: InputT) -> None:
         results=results,
         errors=sorted(errors, key=lambda t: t[0]),
         duration_seconds=time.monotonic() - started,
+        tokens_total=dict(running) if any(running.values()) else {},
+        cost_estimate_usd=running_cost,
     )