diff --git a/docs/README.md b/docs/README.md index 43a3ab0..e59eb1c 100644 --- a/docs/README.md +++ b/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) | diff --git a/examples/flows/batch_usage.py b/examples/flows/batch_usage.py new file mode 100644 index 0000000..f7a94f1 --- /dev/null +++ b/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()) diff --git a/quantmind/flows/__init__.py b/quantmind/flows/__init__.py index faed970..0aa0be2 100644 --- a/quantmind/flows/__init__.py +++ b/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", diff --git a/quantmind/flows/_runner.py b/quantmind/flows/_runner.py index 1f3b507..e669cdb 100644 --- a/quantmind/flows/_runner.py +++ b/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 diff --git a/quantmind/flows/_usage.py b/quantmind/flows/_usage.py new file mode 100644 index 0000000..79e138e --- /dev/null +++ b/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) diff --git a/quantmind/flows/batch.py b/quantmind/flows/batch.py index 9e6c659..3da8287 100644 --- a/quantmind/flows/batch.py +++ b/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, ) diff --git a/tests/flows/test_batch.py b/tests/flows/test_batch.py index 6df8c5c..7fef1b4 100644 --- a/tests/flows/test_batch.py +++ b/tests/flows/test_batch.py @@ -2,11 +2,27 @@ import asyncio import unittest +from dataclasses import dataclass from typing import Any from quantmind.configs import PaperFlowCfg from quantmind.configs.paper import RawText -from quantmind.flows.batch import BatchResult, batch_run +from quantmind.flows._usage import PriceRate, record_usage +from quantmind.flows.batch import ( + BatchResult, + BudgetExceededError, + batch_run, +) + + +@dataclass +class _FakeUsage: + """Duck-types the SDK ``Usage`` object for tests.""" + + requests: int = 1 + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 class BatchResultPropertiesTests(unittest.TestCase): @@ -158,3 +174,85 @@ async def flow( await batch_run(flow, [RawText(text="x")], concurrency=1, marker="here") self.assertEqual(seen, ["here"]) + + +class BatchUsageTests(unittest.IsolatedAsyncioTestCase): + async def test_usage_aggregated_across_inputs(self) -> None: + async def flow(input: RawText, *, cfg: Any = None) -> str: + record_usage( + _FakeUsage(input_tokens=10, output_tokens=4, total_tokens=14) + ) + return f"ok:{input.text}" + + inputs = [RawText(text=str(i)) for i in range(3)] + result = await batch_run(flow, inputs, concurrency=2) + self.assertEqual( + result.tokens_total, + { + "requests": 3, + "input_tokens": 30, + "output_tokens": 12, + "total_tokens": 42, + }, + ) + self.assertEqual(result.cost_estimate_usd, 0.0) + + async def test_prices_populate_cost_estimate(self) -> None: + async def flow(input: RawText, *, cfg: Any = None) -> str: + record_usage(_FakeUsage(input_tokens=1_000_000, output_tokens=0)) + return "ok" + + cfg = PaperFlowCfg(model="gpt-test") + prices = { + "gpt-test": PriceRate(input_usd_per_1m=2.0, output_usd_per_1m=8.0) + } + result = await batch_run( + flow, [RawText(text="x")], cfg=cfg, concurrency=1, prices=prices + ) + # 1M input tokens @ $2/1M = $2.00. + self.assertAlmostEqual(result.cost_estimate_usd, 2.0) + + async def test_token_budget_trips_and_skips_later_inputs(self) -> None: + async def flow(input: RawText, *, cfg: Any = None) -> str: + record_usage(_FakeUsage(input_tokens=100)) + return f"ok:{input.text}" + + cfg = PaperFlowCfg(model="gpt-test", max_total_input_tokens=150) + inputs = [RawText(text=str(i)) for i in range(4)] + # concurrency=1 makes the trip point deterministic: after input 1 the + # running total (200) exceeds 150, so inputs 2 and 3 are skipped. + result = await batch_run(flow, inputs, cfg=cfg, concurrency=1) + self.assertEqual(result.success_count, 2) + self.assertEqual([i for i, _ in result.errors], [2, 3]) + self.assertTrue( + all(isinstance(e, BudgetExceededError) for _, e in result.errors) + ) + + async def test_cost_budget_trips_with_prices(self) -> None: + async def flow(input: RawText, *, cfg: Any = None) -> str: + record_usage(_FakeUsage(input_tokens=1_000_000)) + return "ok" + + cfg = PaperFlowCfg(model="gpt-test", max_total_cost_usd=1.5) + prices = { + "gpt-test": PriceRate(input_usd_per_1m=1.0, output_usd_per_1m=1.0) + } + inputs = [RawText(text=str(i)) for i in range(4)] + result = await batch_run( + flow, inputs, cfg=cfg, concurrency=1, prices=prices + ) + # $1/input; after input 1 cumulative $2.00 > $1.50 → skip 2 and 3. + self.assertEqual(result.success_count, 2) + self.assertEqual([i for i, _ in result.errors], [2, 3]) + + async def test_budget_raise_propagates(self) -> None: + async def flow(input: RawText, *, cfg: Any = None) -> str: + record_usage(_FakeUsage(input_tokens=100)) + return "ok" + + cfg = PaperFlowCfg(model="gpt-test", max_total_input_tokens=50) + inputs = [RawText(text=str(i)) for i in range(3)] + with self.assertRaises(BudgetExceededError): + await batch_run( + flow, inputs, cfg=cfg, concurrency=1, on_error="raise" + ) diff --git a/tests/flows/test_usage.py b/tests/flows/test_usage.py new file mode 100644 index 0000000..f6d1c6e --- /dev/null +++ b/tests/flows/test_usage.py @@ -0,0 +1,86 @@ +"""Tests for ``quantmind.flows._usage``.""" + +import asyncio +import unittest +from dataclasses import dataclass + +from quantmind.flows._usage import ( + PriceRate, + UsageSummary, + record_usage, + usage_scope, +) + + +@dataclass +class _FakeUsage: + """Duck-types the SDK ``Usage`` object for tests.""" + + requests: int = 1 + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + + +class UsageSummaryTests(unittest.TestCase): + def test_as_tokens_dict_excludes_cost(self) -> None: + summary = UsageSummary( + requests=2, input_tokens=10, output_tokens=3, total_tokens=13 + ) + self.assertEqual( + summary.as_tokens_dict(), + { + "requests": 2, + "input_tokens": 10, + "output_tokens": 3, + "total_tokens": 13, + }, + ) + + +class PriceRateTests(unittest.TestCase): + def test_cost_is_per_million(self) -> None: + rate = PriceRate(input_usd_per_1m=300.0, output_usd_per_1m=600.0) + # 1M input @ 300 + 0.5M output @ 600 = 300 + 300 = 600. + self.assertAlmostEqual(rate.cost(1_000_000, 500_000), 600.0) + + +class UsageScopeTests(unittest.IsolatedAsyncioTestCase): + def test_record_outside_scope_is_noop(self) -> None: + # Must not raise when no scope is active. + record_usage(_FakeUsage(input_tokens=99)) + + def test_accumulator_sums_multiple_records(self) -> None: + with usage_scope() as acc: + record_usage(_FakeUsage(input_tokens=5, output_tokens=2)) + record_usage(_FakeUsage(input_tokens=3, output_tokens=1)) + summary = acc.summary() + self.assertEqual(summary.requests, 2) + self.assertEqual(summary.input_tokens, 8) + self.assertEqual(summary.output_tokens, 3) + + async def test_scopes_isolate_across_concurrent_tasks(self) -> None: + async def worker(unit: int) -> int: + with usage_scope() as acc: + record_usage(_FakeUsage(input_tokens=unit)) + await asyncio.sleep(0) # force interleave with the sibling + record_usage(_FakeUsage(input_tokens=unit)) + return acc.summary().input_tokens + + one, ten = await asyncio.gather(worker(1), worker(10)) + self.assertEqual(one, 2) + self.assertEqual(ten, 20) + + async def test_nested_gather_accumulates_into_outer_scope(self) -> None: + # The batch case: children run as separate tasks yet fold into the + # one accumulator the outer scope set (shared mutable reference). + async def child(unit: int) -> None: + record_usage(_FakeUsage(input_tokens=unit)) + + with usage_scope() as acc: + await asyncio.gather(child(3), child(4)) + self.assertEqual(acc.summary().input_tokens, 7) + + +if __name__ == "__main__": + unittest.main()