mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-17 20:19:29 +00:00
Marius CiepluchandClaude Fable 5.1
c33b053b15
budget_spent carries end_detail (spend, task_spend, reason_rounds, ...); a task whose allowance cannot cover the next call finishes with its facts; per-task metrics in pilot records; input reservation bytes/2 + 1024 per message (measured 3-4x over-reservation before). Co-Authored-By: Claude Fable 5.1 <[email protected]>
323 lines
13 KiB
Python
323 lines
13 KiB
Python
"""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 decimal import Decimal
|
|
from typing import Any, Protocol
|
|
|
|
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, end_detail: str = "spend") -> None:
|
|
super().__init__(f"{end_detail}: {message}")
|
|
self.end_detail = end_detail
|
|
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,
|
|
temperature: float | None = None,
|
|
logprobs: bool = False,
|
|
) -> ChatResponse:
|
|
"""One metered call. ``temperature`` and ``logprobs`` pass through only when set, so a wrapped client that
|
|
does not take them (a scripted one) is called as before."""
|
|
self.check()
|
|
extra: dict[str, Any] = {}
|
|
if temperature is not None:
|
|
extra["temperature"] = temperature
|
|
if logprobs:
|
|
extra["logprobs"] = True
|
|
try:
|
|
response = self.inner.complete(
|
|
model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds, json_object=json_object, **extra
|
|
)
|
|
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
|
|
|
|
|
|
class EmbeddingClient(Protocol):
|
|
def embed(self, *, model: str, inputs: list[str], timeout_seconds: int = 60) -> list[list[float]]: ...
|
|
|
|
|
|
class UnmeteredUsage(RuntimeError):
|
|
"""A priced embedding call whose provider reported no tokens: a working-looking zero, never a free call."""
|
|
|
|
|
|
class MeteredEmbeddingClient:
|
|
"""An embedding client metered like :class:`MeteredClient`: the wrapped client records one usage dict per call
|
|
in a ``usage`` list (``OpenRouterEmbeddingClient`` does), ``prompt_tokens`` are priced at the Model manifest's
|
|
``input_per_m``, and the usd/calls ceilings stop the next call. A priced call that reports no tokens raises
|
|
:class:`UnmeteredUsage` (the 2026-09 lesson: a silent zero is an instrument failure)."""
|
|
|
|
def __init__(
|
|
self,
|
|
inner: EmbeddingClient,
|
|
*,
|
|
prices: Mapping[str, object] | Prices | None = None,
|
|
budget_usd: float | None = None,
|
|
budget_calls: int | None = None,
|
|
require_usage: bool = True,
|
|
) -> 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")
|
|
self.budget_usd, self.budget_calls, self.require_usage = budget_usd, budget_calls, require_usage
|
|
self.calls = 0
|
|
self.tokens = 0
|
|
self._usd = 0.0
|
|
self._seen_rows = len(self._inner_rows())
|
|
|
|
@property
|
|
def usd(self) -> float | None:
|
|
return self._usd if self.prices is not None else None
|
|
|
|
def check(self) -> None:
|
|
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}", 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 embed(self, *, model: str, inputs: list[str], timeout_seconds: int = 60) -> list[list[float]]:
|
|
self.check()
|
|
try:
|
|
vectors = self.inner.embed(model=model, inputs=inputs, timeout_seconds=timeout_seconds)
|
|
except Exception as exc:
|
|
self._absorb_rows()
|
|
if is_account_error(exc):
|
|
raise AccountError(str(exc)) from exc
|
|
raise
|
|
self.calls += 1
|
|
tokens = self._absorb_rows()
|
|
if self.require_usage and self.prices is not None and tokens == 0 and inputs:
|
|
raise UnmeteredUsage(
|
|
f"embedding call {self.calls} reported no prompt tokens for {len(inputs)} inputs: the meter cannot price it"
|
|
)
|
|
return vectors
|
|
|
|
def summary(self) -> dict[str, Any]:
|
|
return {"usd": self.usd, "calls": self.calls, "usage": {"prompt_tokens": self.tokens}, "priced": self.prices is not None}
|
|
|
|
def _inner_rows(self) -> list[Any]:
|
|
rows = getattr(self.inner, "usage", None)
|
|
return rows if isinstance(rows, list) else []
|
|
|
|
def _absorb_rows(self) -> int:
|
|
rows = self._inner_rows()
|
|
tokens = 0
|
|
for row in rows[self._seen_rows :]:
|
|
if isinstance(row, Mapping):
|
|
tokens += _tokens(row.get("prompt_tokens") or row.get("total_tokens"))
|
|
self._seen_rows = len(rows)
|
|
self.tokens += tokens
|
|
if self.prices is not None:
|
|
self._usd += tokens * self.prices.input_per_m / 1_000_000
|
|
return tokens
|
|
|
|
|
|
__all__ = [
|
|
"AccountError", "BudgetExhausted", "EmbeddingClient", "MeteredClient", "MeteredEmbeddingClient", "UnmeteredUsage", "cost_of",
|
|
"is_account_error", "prices_from", "SpendBudget", "SpendReservation",
|
|
] # fmt: skip
|
|
|
|
|
|
class SpendBudget:
|
|
"""Thread-safe pre-call admission. Share one instance across concurrent callers.
|
|
|
|
A caller must supply a genuine upper bound, including all tool-loop calls.
|
|
Reservations remain charged until settled or explicitly released on failure.
|
|
Existing MeteredClient callers retain their historical post-call metering.
|
|
"""
|
|
|
|
def __init__(self, cap: float, *, spent: float = 0.0) -> None:
|
|
from threading import Lock
|
|
|
|
self._lock = Lock()
|
|
self._cap = self._amount(cap)
|
|
self._spent = self._amount(spent)
|
|
if self._spent > self._cap:
|
|
raise ValueError("spend exceeds cap")
|
|
self._reserved = Decimal(0)
|
|
|
|
@staticmethod
|
|
def _amount(value: float) -> Decimal:
|
|
|
|
amount = Decimal(str(value))
|
|
if not amount.is_finite() or amount < 0:
|
|
raise ValueError("cost must be finite and nonnegative")
|
|
return amount
|
|
|
|
@property
|
|
def spent(self) -> float:
|
|
with self._lock:
|
|
return float(self._spent)
|
|
|
|
@property
|
|
def reserved(self) -> float:
|
|
with self._lock:
|
|
return float(self._reserved)
|
|
|
|
def reserve(self, max_cost: float) -> SpendReservation:
|
|
amount = self._amount(max_cost)
|
|
with self._lock:
|
|
if self._spent + self._reserved + amount > self._cap:
|
|
raise BudgetExhausted("spend reservation exceeds cap", usd=float(self._spent), calls=0)
|
|
self._reserved += amount
|
|
return SpendReservation(self, amount)
|
|
|
|
|
|
class SpendReservation:
|
|
def __init__(self, budget: SpendBudget, amount: Decimal) -> None:
|
|
self._budget, self._amount, self._active = budget, amount, True
|
|
|
|
def settle(self, actual: float) -> None:
|
|
amount = self._budget._amount(actual)
|
|
with self._budget._lock:
|
|
if not self._active:
|
|
raise ValueError("reservation already closed")
|
|
if amount > self._amount:
|
|
# Retain the reservation: a broken upper bound must fail closed.
|
|
raise ValueError("actual cost exceeds reservation")
|
|
self._budget._reserved -= self._amount
|
|
self._budget._spent += amount
|
|
self._active = False
|
|
|
|
def release(self) -> None:
|
|
self.settle(0)
|