mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
feat(plane): budget metering (task 2)
MeteredClient wraps a ChatClient, sums the provider's usage rows per call, prices them from a Model manifest's parameters (cache-hit, input, output per million), and raises BudgetExhausted before a call once the usd or calls ceiling is reached. A 401/402 or "Insufficient Balance" failure becomes AccountError with the provider's message. An unpriced model reports usd None, never zero, and refuses a usd budget; summary() gives the spend block of a task's summary.json. The two plane modules are registered as spec-owned in the module audit (their consumers land in tasks 3-4) and the audit manifest is regenerated. Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
152728c052
commit
f460735a87
@@ -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"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user