From 5cfec1fc19708a55e015954000e2c261ad7de8dc Mon Sep 17 00:00:00 2001 From: satoridev01 Date: Mon, 2 Nov 2026 21:25:00 -0300 Subject: [PATCH] perf(cost): keep running totals instead of recounting every response total_cost summed every recorded response on each read, and it is read on every save and budget check, so a scan's total counting work grew with the square of its call count. Keep the paired cost and the unmatched charges as running totals, updated under a lock when an estimate or a charge is recorded, so reading the total is O(1). --- strix/report/usage.py | 56 +++++++++++++++++++++-------- tests/test_cost_response_pairing.py | 26 ++++++++++++++ 2 files changed, 67 insertions(+), 15 deletions(-) diff --git a/strix/report/usage.py b/strix/report/usage.py index cccaba8da..19992b9c9 100644 --- a/strix/report/usage.py +++ b/strix/report/usage.py @@ -4,6 +4,7 @@ from __future__ import annotations import logging import math +import threading from typing import Any from agents.usage import Usage, deserialize_usage, serialize_usage @@ -76,6 +77,14 @@ class LLMUsageLedger: # contract, see ``total_cost``. self._estimated_by_response: dict[str, float] = {} self._observed_by_response: dict[str, float] = {} + # Running totals over the two maps, kept in step on every write so + # reading ``total_cost`` (each save, each budget check) is O(1): + # the cost of id-matched responses, and charges with no estimate. + self._paired_cost = 0.0 + self._unpaired_observed_by_id = 0.0 + # Estimates arrive from the event loop, charges from provider + # callbacks; one lock keeps a map and its running total in step. + self._response_lock = threading.Lock() self._persisted_cost = 0.0 # Keyed by upstream provider name, e.g. "Z.AI" or "DeepInfra". self._providers: dict[str, ProviderUsage] = {} @@ -111,7 +120,7 @@ class LLMUsageLedger: estimated = _estimate_litellm_cost(usage, model) if estimated and math.isfinite(estimated): if response_id: - self._estimated_by_response[response_id] = estimated + self._set_response_cost(response_id, estimated=estimated) else: self._estimated_cost += estimated self._agent_estimated_cost[normalized_agent_id] = ( @@ -125,10 +134,35 @@ class LLMUsageLedger: return if isinstance(cost, int | float) and math.isfinite(cost) and cost > 0: if response_id: - self._observed_by_response[response_id] = float(cost) + self._set_response_cost(response_id, observed=float(cost)) else: self._observed_cost += float(cost) + def _response_contribution(self, response_id: str) -> tuple[float, float]: + """This response's share of (paired cost, unpaired keyed charges).""" + observed = self._observed_by_response.get(response_id) + estimated = self._estimated_by_response.get(response_id) + if estimated is not None: + return (estimated if observed is None else observed), 0.0 + return 0.0, (observed or 0.0) + + def _set_response_cost( + self, + response_id: str, + *, + estimated: float | None = None, + observed: float | None = None, + ) -> None: + with self._response_lock: + old_paired, old_unpaired = self._response_contribution(response_id) + if estimated is not None: + self._estimated_by_response[response_id] = estimated + if observed is not None: + self._observed_by_response[response_id] = observed + new_paired, new_unpaired = self._response_contribution(response_id) + self._paired_cost += new_paired - old_paired + self._unpaired_observed_by_id += new_unpaired - old_unpaired + def record_provider( self, provider: str, @@ -165,26 +199,16 @@ class LLMUsageLedger: return _round_cost( sum(self._configured_agent_cost(aid, model, prices) for aid in self._agent_usage) ) - # Copies: tool threads read the total while the event loop records new - # responses, and iterating a dict that grows raises RuntimeError. - estimated_by_response = dict(self._estimated_by_response) - observed_by_response = dict(self._observed_by_response) # Matched by response id: the provider's charge replaces that response's # estimate, and a response without a charge keeps its own estimate. - paired = sum( - observed_by_response.get(response_id, estimate) - for response_id, estimate in estimated_by_response.items() - ) # Whatever id matching cannot settle keeps the upstream contract: a # provider charge with no estimate to match replaces the unkeyed # estimates instead of adding to them, because without a shared id they # may describe the same response. Adding both would overstate the spend # and stop a scan early at --max-budget. - unpaired_observed = self._observed_cost + sum( - cost - for response_id, cost in observed_by_response.items() - if response_id not in estimated_by_response - ) + with self._response_lock: + paired = self._paired_cost + unpaired_observed = self._observed_cost + self._unpaired_observed_by_id unpaired = unpaired_observed if unpaired_observed > 0 else self._estimated_cost return _round_cost(self._persisted_cost + paired + unpaired) @@ -236,6 +260,8 @@ class LLMUsageLedger: self._estimated_cost = 0.0 self._estimated_by_response.clear() self._observed_by_response.clear() + self._paired_cost = 0.0 + self._unpaired_observed_by_id = 0.0 self._persisted_cost = 0.0 self._providers = {} diff --git a/tests/test_cost_response_pairing.py b/tests/test_cost_response_pairing.py index dc2283abc..4e23f98aa 100644 --- a/tests/test_cost_response_pairing.py +++ b/tests/test_cost_response_pairing.py @@ -2,6 +2,7 @@ from __future__ import annotations +import random from types import SimpleNamespace from typing import TYPE_CHECKING @@ -159,3 +160,28 @@ def test_sdk_usage_correlates_request_ids_per_agent_and_keeps_upstream_contract_ state.record_observed_llm_cost(0.2, response_id="ambiguous-b") state.record_sdk_usage(agent_id="root", usage=_usage(500), model="model") assert state.get_total_llm_cost() == pytest.approx(0.97) + + +def test_running_totals_match_a_full_recount(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + usage_module, + "_estimate_litellm_cost", + lambda used, _model: used.input_tokens / 1000, + ) + rng = random.Random(1642) # noqa: S311 - a seeded, repeatable test sequence + ledger = LLMUsageLedger() + estimated: dict[str, float] = {} + observed: dict[str, float] = {} + for _ in range(500): + response_id = f"r{rng.randrange(60)}" + if rng.random() < 0.5: + tokens = rng.randrange(1, 900) + ledger.record(agent_id="a", usage=_usage(tokens), model="m", response_id=response_id) + estimated[response_id] = tokens / 1000 + else: + cost = rng.randrange(1, 900) / 1000 + ledger.record_observed_cost(cost, response_id=response_id) + observed[response_id] = cost + paired = sum(observed.get(rid, est) for rid, est in estimated.items()) + unpaired = sum(cost for rid, cost in observed.items() if rid not in estimated) + assert ledger.total_cost == pytest.approx(paired + unpaired, abs=1e-9)