diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index b466258..1c995e4 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 118, + "modules": 120, "counts": { - "load_bearing": 109, + "load_bearing": 111, "standalone_capability": 9 }, "orphaned": [], @@ -400,6 +400,7 @@ "model.scan", "model.specs", "model.taint", + "plane.budget", "preprocess", "push.assistance", "push.cache", @@ -409,7 +410,8 @@ "push.reuse", "push.runner", "push.snapshot", - "push.vetoes" + "push.vetoes", + "tool_hunter" ] }, { @@ -479,7 +481,9 @@ "reason": "model-grounded-detection Req 8: the judge client, wired by model/judge.py in group 3", "importers": [ "cli", - "push.assistance" + "plane.budget", + "push.assistance", + "tool_hunter" ] }, { @@ -645,6 +649,18 @@ "pair_gate" ] }, + { + "module": "plane", + "classification": "load_bearing", + "reason": "ai-service-plane Req 1-4: the service plane package; runner and reconciler integration in tasks 3-4", + "importers": [] + }, + { + "module": "plane.budget", + "classification": "load_bearing", + "reason": "ai-service-plane Req 2.3/2.4: metered chat client with usd and calls ceilings; task integration in tasks 3, 6", + "importers": [] + }, { "module": "policy", "classification": "load_bearing", @@ -1041,6 +1057,7 @@ "mapping", "model.candidates", "model.layout", + "model.scan", "model.specs", "pairs", "semantic.cst", @@ -1048,6 +1065,7 @@ "semantic.functions", "semantic.ir", "semantic.obligations.check", + "semantic.obligations.facts", "semantic.obligations.operations", "semantic.obligations.shapes", "semantic.overlay", @@ -1093,10 +1111,12 @@ "mapping", "model.candidates", "model.layout", + "model.scan", "model.specs", "pairs", "semantic", "semantic.obligations.check", + "semantic.obligations.facts", "semantic.obligations.operations", "semantic.obligations.shapes", "semantic.overlay", @@ -1293,7 +1313,8 @@ "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ - "model.endpoint" + "model.endpoint", + "plane.budget" ] }, { diff --git a/src/openultrasast/model/audit.py b/src/openultrasast/model/audit.py index 7544f0e..99132c1 100644 --- a/src/openultrasast/model/audit.py +++ b/src/openultrasast/model/audit.py @@ -54,6 +54,8 @@ SPEC_OWNED = { "cpg": "model-grounded-detection Req 4: the CPG seam package", "cpg.backend": "model-grounded-detection Req 4.1/4.5: the single Joern subprocess boundary", "cpg.capability": "model-grounded-detection Req 4.2/4.3: the engine probe that degrades to suspicion", + "plane": "ai-service-plane Req 1-4: the service plane package; runner and reconciler integration in tasks 3-4", + "plane.budget": "ai-service-plane Req 2.3/2.4: metered chat client with usd and calls ceilings; task integration in tasks 3, 6", } STANDALONE = { diff --git a/src/openultrasast/plane/__init__.py b/src/openultrasast/plane/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/openultrasast/plane/budget.py b/src/openultrasast/plane/budget.py new file mode 100644 index 0000000..7db2e68 --- /dev/null +++ b/src/openultrasast/plane/budget.py @@ -0,0 +1,163 @@ +"""Budget metering for a task's model calls (ai-service-plane, Requirement 2.3 / 2.4). + +`MeteredClient` wraps any `ChatClient`, sums the provider's usage fields per call, prices them from a Model +manifest's parameters, and stops the task at its ceiling. Two failure shapes are distinct on purpose: a ceiling +is `BudgetExhausted` (the task is `unfinished` and resumes under a larger budget), an account or authentication +failure is `AccountError` (the task is `failed` and the run starts nothing further). A model without prices is +`unpriced`, which is not zero dollars: such a client refuses a usd budget rather than reporting a clean 0. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from ..model.endpoint import Prices +from ..tool_hunter import ChatClient, ChatResponse + +_ACCOUNT_MARKERS = ("HTTP Error 401", "HTTP Error 402", "Insufficient Balance") +_USAGE_FIELDS = ("prompt_tokens", "prompt_cache_hit_tokens", "completion_tokens") + + +class BudgetExhausted(RuntimeError): + """Raised before a call would start once the usd or calls ceiling is reached.""" + + def __init__(self, message: str, *, usd: float | None, calls: int) -> None: + super().__init__(message) + self.usd = usd + self.calls = calls + + +class AccountError(RuntimeError): + """Raised when the provider refuses the account: 401, 402 or an insufficient-balance body.""" + + def __init__(self, provider_message: str) -> None: + super().__init__(provider_message) + self.provider_message = provider_message + + +def is_account_error(exc: BaseException) -> bool: + text = str(exc).lower() + return any(marker.lower() in text for marker in _ACCOUNT_MARKERS) + + +def prices_from(parameters: Mapping[str, object] | Prices | None) -> Prices | None: + """Prices from a Model manifest's `spec.parameters` (or an existing `Prices`); None when any rate is missing.""" + if parameters is None or isinstance(parameters, Prices): + return parameters + values: list[float] = [] + for key in ("cache_hit_per_m", "input_per_m", "output_per_m"): + value = parameters.get(key) + if not isinstance(value, int | float) or isinstance(value, bool): + return None + values.append(float(value)) + return Prices(cache_hit_per_m=values[0], input_per_m=values[1], output_per_m=values[2]) + + +def cost_of(usage: Mapping[str, object], prices: Prices) -> float: + """Dollars for one call from the provider's usage fields, the arithmetic of `ChatEndpoint.cost`.""" + cache_hit = _tokens(usage.get("prompt_cache_hit_tokens")) + miss = usage.get("prompt_cache_miss_tokens") + input_tokens = _tokens(miss) if miss is not None else max(_tokens(usage.get("prompt_tokens")) - cache_hit, 0) + output_tokens = _tokens(usage.get("completion_tokens")) + return (cache_hit * prices.cache_hit_per_m + input_tokens * prices.input_per_m + output_tokens * prices.output_per_m) / 1_000_000 + + +class MeteredClient: + """A `ChatClient` that meters spend from the wrapped client's provider usage rows. + + The wrapped client is expected to record one usage dict per call in a `usage` list (as `DeepSeekChatClient` + does); rows appended during a call are attributed to it. A client without that list still has its calls + counted, with zero tokens. + """ + + def __init__( + self, + inner: ChatClient, + *, + prices: Mapping[str, object] | Prices | None = None, + budget_usd: float | None = None, + budget_calls: int | None = None, + ) -> None: + self.inner = inner + self.prices = prices_from(prices) + if budget_usd is not None and self.prices is None: + raise ValueError("unpriced model cannot run under a usd budget") + if budget_usd is not None and budget_usd < 0: + raise ValueError("budget_usd must not be negative") + if budget_calls is not None and budget_calls < 0: + raise ValueError("budget_calls must not be negative") + self.budget_usd = budget_usd + self.budget_calls = budget_calls + self.calls = 0 + self.usage: dict[str, int] = dict.fromkeys(_USAGE_FIELDS, 0) + self._usd = 0.0 + self._seen_rows = len(self._inner_rows()) + + @property + def priced(self) -> bool: + return self.prices is not None + + @property + def usd(self) -> float | None: + return self._usd if self.prices is not None else None + + def check(self) -> None: + """Raise `BudgetExhausted` when the next call must not start.""" + if self.budget_usd is not None and self._usd >= self.budget_usd: + raise BudgetExhausted( + f"usd budget exhausted: spent {self._usd:.6f} of {self.budget_usd:.6f} after {self.calls} calls", + usd=self._usd, + calls=self.calls, + ) + if self.budget_calls is not None and self.calls >= self.budget_calls: + raise BudgetExhausted(f"calls budget exhausted: {self.calls} of {self.budget_calls}", usd=self.usd, calls=self.calls) + + def complete( + self, + *, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]], + timeout_seconds: int = 60, + json_object: bool = False, + ) -> ChatResponse: + self.check() + try: + response = self.inner.complete( + model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object + ) + except Exception as exc: + self._absorb_rows() + if is_account_error(exc): + raise AccountError(str(exc)) from exc + raise + self.calls += 1 + self._absorb_rows() + return response + + def summary(self) -> dict[str, Any]: + """The spend block of a task's `summary.json`: `usd` is None, not 0, for an unpriced model.""" + return {"usd": self.usd, "calls": self.calls, "usage": dict(self.usage), "priced": self.priced} + + def _inner_rows(self) -> list[Any]: + rows = getattr(self.inner, "usage", None) + return rows if isinstance(rows, list) else [] + + def _absorb_rows(self) -> None: + rows = self._inner_rows() + for row in rows[self._seen_rows :]: + if not isinstance(row, Mapping): + continue + for field in _USAGE_FIELDS: + self.usage[field] += _tokens(row.get(field)) + if self.prices is not None: + self._usd += cost_of(row, self.prices) + self._seen_rows = len(rows) + + +def _tokens(value: object) -> int: + return int(value) if isinstance(value, int | float) and not isinstance(value, bool) else 0 + + +__all__ = ["AccountError", "BudgetExhausted", "MeteredClient", "cost_of", "is_account_error", "prices_from"] diff --git a/tests/test_plane_budget.py b/tests/test_plane_budget.py new file mode 100644 index 0000000..266e600 --- /dev/null +++ b/tests/test_plane_budget.py @@ -0,0 +1,193 @@ +"""Budget metering (ai-service-plane task 2): ceilings, account errors, unpriced refusal, cache-hit pricing.""" + +from __future__ import annotations + +import pytest + +from openultrasast.model.endpoint import Prices +from openultrasast.plane.budget import AccountError, BudgetExhausted, MeteredClient, cost_of, prices_from +from openultrasast.tool_hunter import ChatResponse + +FLASH = {"cache_hit_per_m": 0.014, "input_per_m": 0.44, "output_per_m": 1.32} +MESSAGES: list[dict[str, object]] = [{"role": "user", "content": "hi"}] + + +class UsageClient: + """Records provider usage rows like DeepSeekChatClient; a scripted step may raise instead.""" + + def __init__(self, steps: list[dict[str, object] | Exception]) -> None: + self.steps = list(steps) + self.usage: list[dict[str, object]] = [] + self.calls = 0 + + def complete( + self, + *, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]], + timeout_seconds: int = 60, + json_object: bool = False, + ) -> ChatResponse: + self.calls += 1 + step = self.steps.pop(0) + if isinstance(step, Exception): + raise step + self.usage.append(dict(step)) + return ChatResponse(content="ok") + + +def _row(prompt: int, cache_hit: int, completion: int) -> dict[str, object]: + return {"prompt_tokens": prompt, "prompt_cache_hit_tokens": cache_hit, "completion_tokens": completion} + + +def _call(client: MeteredClient) -> ChatResponse: + return client.complete(model="m", messages=MESSAGES, tools=[]) + + +def test_cache_hit_pricing_is_exact_dollars() -> None: + # 1_000_000 prompt tokens of which 250_000 hit the cache, 100_000 completion tokens: + # 250_000 * 0.014/M + 750_000 * 0.44/M + 100_000 * 1.32/M = 0.0035 + 0.33 + 0.132 + inner = UsageClient([_row(1_000_000, 250_000, 100_000)]) + metered = MeteredClient(inner, prices=FLASH) + _call(metered) + assert metered.usd == pytest.approx(0.4655, abs=1e-12) + assert metered.calls == 1 + assert metered.usage == {"prompt_tokens": 1_000_000, "prompt_cache_hit_tokens": 250_000, "completion_tokens": 100_000} + + +def test_cost_of_honours_explicit_cache_miss_field() -> None: + prices = Prices(cache_hit_per_m=1.0, input_per_m=10.0, output_per_m=100.0) + row = {"prompt_tokens": 30, "prompt_cache_hit_tokens": 10, "prompt_cache_miss_tokens": 20, "completion_tokens": 5} + assert cost_of(row, prices) == pytest.approx((10 * 1.0 + 20 * 10.0 + 5 * 100.0) / 1_000_000) + + +def test_usage_sums_across_calls() -> None: + inner = UsageClient([_row(100, 40, 10), _row(200, 0, 20)]) + metered = MeteredClient(inner, prices=FLASH) + _call(metered) + _call(metered) + assert metered.usage == {"prompt_tokens": 300, "prompt_cache_hit_tokens": 40, "completion_tokens": 30} + expected = (40 * 0.014 + 60 * 0.44 + 10 * 1.32 + 200 * 0.44 + 20 * 1.32) / 1_000_000 + assert metered.usd == pytest.approx(expected) + + +def test_usd_ceiling_stops_before_the_next_call() -> None: + # each call costs exactly 1.32 usd (1M completion tokens); the budget of 2.0 admits two calls, not three + inner = UsageClient([_row(0, 0, 1_000_000)] * 3) + metered = MeteredClient(inner, prices=FLASH, budget_usd=2.0) + _call(metered) + _call(metered) + with pytest.raises(BudgetExhausted) as info: + _call(metered) + assert inner.calls == 2, "the third call must not reach the provider" + assert info.value.calls == 2 + assert info.value.usd == pytest.approx(2.64) + assert isinstance(info.value, RuntimeError) + + +def test_calls_ceiling_stops_before_the_next_call() -> None: + inner = UsageClient([_row(1, 0, 1)] * 3) + metered = MeteredClient(inner, prices=FLASH, budget_calls=2) + _call(metered) + _call(metered) + with pytest.raises(BudgetExhausted) as info: + _call(metered) + assert inner.calls == 2 + assert info.value.calls == 2 + + +def test_zero_calls_budget_refuses_the_first_call() -> None: + inner = UsageClient([_row(1, 0, 1)]) + metered = MeteredClient(inner, budget_calls=0) + with pytest.raises(BudgetExhausted): + _call(metered) + assert inner.calls == 0 + + +@pytest.mark.parametrize( + "message", + [ + 'OpenRouter chat request failed: HTTPError: HTTP Error 402: Payment Required: {"error":{"message":"Insufficient credits"}}', + "OpenRouter chat request failed: HTTPError: HTTP Error 401: Unauthorized", + 'DeepSeek chat request failed: {"error":{"message":"Insufficient Balance","type":"unknown_error"}}', + ], +) +def test_account_errors_are_distinct_and_carry_the_provider_message(message: str) -> None: + inner = UsageClient([RuntimeError(message)]) + metered = MeteredClient(inner, prices=FLASH, budget_usd=5.0) + with pytest.raises(AccountError) as info: + _call(metered) + assert info.value.provider_message == message + assert isinstance(info.value.__cause__, RuntimeError) + assert metered.calls == 0 + + +def test_other_provider_errors_pass_through_unchanged() -> None: + inner = UsageClient([RuntimeError("OpenRouter chat request failed: HTTPError: HTTP Error 503: Service Unavailable")]) + metered = MeteredClient(inner, prices=FLASH) + with pytest.raises(RuntimeError, match="503") as info: + _call(metered) + assert not isinstance(info.value, AccountError) + + +def test_unpriced_model_refuses_a_usd_budget() -> None: + with pytest.raises(ValueError, match="unpriced model cannot run under a usd budget"): + MeteredClient(UsageClient([]), budget_usd=1.0) + with pytest.raises(ValueError, match="unpriced"): + MeteredClient(UsageClient([]), prices={"input_per_m": 0.44}, budget_usd=1.0) + + +def test_unpriced_model_reports_none_not_zero_under_a_calls_budget() -> None: + inner = UsageClient([_row(500, 100, 50)]) + metered = MeteredClient(inner, budget_calls=3) + _call(metered) + assert metered.priced is False + assert metered.usd is None + assert metered.usage["prompt_tokens"] == 500 + assert metered.summary() == { + "usd": None, + "calls": 1, + "usage": {"prompt_tokens": 500, "prompt_cache_hit_tokens": 100, "completion_tokens": 50}, + "priced": False, + } + + +def test_summary_shape_for_a_priced_model() -> None: + inner = UsageClient([_row(1_000_000, 0, 0)]) + metered = MeteredClient(inner, prices=Prices(**FLASH), budget_usd=10.0) + _call(metered) + summary = metered.summary() + assert set(summary) == {"usd", "calls", "usage", "priced"} + assert summary["priced"] is True + assert summary["calls"] == 1 + assert summary["usd"] == pytest.approx(0.44) + assert summary["usage"] == {"prompt_tokens": 1_000_000, "prompt_cache_hit_tokens": 0, "completion_tokens": 0} + + +def test_prices_from_accepts_manifest_parameters_and_rejects_partial_ones() -> None: + assert prices_from(FLASH) == Prices(cache_hit_per_m=0.014, input_per_m=0.44, output_per_m=1.32) + assert prices_from({"cache_hit_per_m": 0.014, "input_per_m": 0.44}) is None + assert prices_from({"cache_hit_per_m": "0.014", "input_per_m": 0.44, "output_per_m": 1.32}) is None + assert prices_from(None) is None + + +def test_client_without_usage_rows_still_counts_calls() -> None: + class Bare: + def complete( + self, + *, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]], + timeout_seconds: int = 60, + json_object: bool = False, + ) -> ChatResponse: + return ChatResponse(content="ok") + + metered = MeteredClient(Bare(), prices=FLASH, budget_calls=1) + _call(metered) + assert metered.calls == 1 + assert metered.usd == 0.0 + with pytest.raises(BudgetExhausted): + _call(metered)