« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Relay billing lifecycle service.
3Implements :class:`~lexigram.contracts.ai.governance.RelayBillingProtocol`
4for one relay request: pre-consume admission, settle-once settlement, and
5release. Prompt estimation, scope reservations, and per-dimension pricing
6are delegated to the reservation manager and the injected price estimator;
7actual usage always comes from ``RelayConvertResult``, never from an
8estimate.
9"""
11from __future__ import annotations
13from collections.abc import Callable
14from decimal import Decimal
15from typing import TYPE_CHECKING, Literal
17from lexigram.ai.governance.relay_billing.models import DEFAULT_CURRENCY, USAGE_MISSING
18from lexigram.ai.governance.relay_billing.reservations import (
19 RelayReservationManager,
20 estimate_prompt_tokens,
21 requested_max_output_tokens,
22)
23from lexigram.contracts.ai.governance import (
24 RelayBillingError,
25 RelayBillingProtocol,
26 RelayUsageRecord,
27 RelayUsageScope,
28 RelayUsageStoreProtocol,
29 billing_store_unavailable,
30)
31from lexigram.contracts.ai.relay import RelayUsage
32from lexigram.contracts.core.result import Err, Ok, Result
33from lexigram.identity import ambient as identity
34from lexigram.logging import get_logger
36if TYPE_CHECKING:
37 from lexigram.contracts.ai.governance import (
38 RelayPriceEstimatorProtocol,
39 RelayUsageReservation,
40 )
41 from lexigram.contracts.ai.llm import TokenCounterProtocol
42 from lexigram.contracts.ai.relay import RelayConvertResult, RelayRequestPayload
44logger = get_logger(__name__)
46__all__ = [
47 "RelayBillingService",
48 "RelayCostAdapter",
49]
52class RelayBillingService(RelayBillingProtocol):
53 """Billing lifecycle for one relay gateway request.
55 Args:
56 reservation_manager: Manager holding per-scope admission quotas.
57 estimator: Price estimator computing reservation and settled
58 charges from normalized usage.
59 store: Persistence boundary for reservations and settlement
60 records.
61 token_counter: Optional model-aware token counter for prompt
62 estimation.
63 currency: ISO currency code for charges.
64 """
66 def __init__(
67 self,
68 *,
69 reservation_manager: RelayReservationManager,
70 estimator: RelayPriceEstimatorProtocol,
71 store: RelayUsageStoreProtocol,
72 token_counter: TokenCounterProtocol | None = None,
73 currency: str = DEFAULT_CURRENCY,
74 ) -> None:
75 self._manager = reservation_manager
76 self._estimator = estimator
77 self._store = store
78 self._token_counter = token_counter
79 self._currency = currency
80 # reservation_id -> scope captured at admission time
81 self._scopes: dict[str, RelayUsageScope] = {}
83 async def pre_consume(
84 self,
85 request_id: str,
86 scope: RelayUsageScope,
87 payload: RelayRequestPayload,
88 ) -> Result[RelayUsageReservation, RelayBillingError]:
89 """Reserve admission capacity for a request before upstream I/O.
91 Args:
92 request_id: Gateway request identifier.
93 scope: Accounting scope of the request.
94 payload: Relay request payload used for prompt estimation.
96 Returns:
97 Ok(reservation) after the reservation is persisted, or
98 Err when admission cannot be proven (quota, pricing, or
99 persistence failure). On persistence failure the in-memory
100 reservation is released, leaving none behind.
101 """
102 prompt_tokens = estimate_prompt_tokens(payload, self._token_counter)
103 max_output = requested_max_output_tokens(payload)
104 estimate_tokens = prompt_tokens + max_output
106 estimated_charge = await self._max_charge(scope, prompt_tokens, max_output)
107 if estimated_charge.is_err():
108 return Err(
109 RelayBillingError(
110 code=estimated_charge.unwrap_err().code,
111 message=estimated_charge.unwrap_err().message,
112 request_id=request_id,
113 tenant_id=scope.tenant_id,
114 )
115 )
117 reservation_result = await self._manager.reserve(
118 request_id,
119 scope,
120 estimate_tokens,
121 estimated_charge.unwrap(),
122 )
123 if reservation_result.is_err():
124 return Err(reservation_result.unwrap_err())
125 reservation = reservation_result.unwrap()
127 try:
128 await self._store.save_reservation(reservation)
129 except Exception as exc: # noqa: BLE001 - infrastructure boundary
130 await self._manager.release(reservation.reservation_id)
131 logger.warning(
132 "relay_billing_persist_reservation_failed",
133 reservation_id=reservation.reservation_id,
134 request_id=request_id,
135 tenant_id=scope.tenant_id,
136 error=str(exc),
137 )
138 return Err(
139 billing_store_unavailable(
140 message="cannot persist reservation; admission rejected",
141 request_id=request_id,
142 tenant_id=scope.tenant_id,
143 )
144 )
146 self._scopes[reservation.reservation_id] = scope
147 return Ok(reservation)
149 async def _max_charge(
150 self,
151 scope: RelayUsageScope,
152 prompt_tokens: int,
153 max_output: int,
154 ) -> Result[Decimal, RelayBillingError]:
155 """Return the maximum reservation charge for the estimate.
157 Args:
158 scope: Accounting scope of the request.
159 prompt_tokens: Estimated prompt tokens.
160 max_output: Requested max output tokens.
162 Returns:
163 Ok(worst-case charge) for the full completion budget, or Err
164 when pricing is unavailable for the model.
165 """
166 usage = RelayUsage(
167 prompt_tokens=prompt_tokens,
168 completion_tokens=max_output,
169 )
170 result = self._estimator.estimate_charge(
171 scope.model,
172 usage,
173 provider=scope.provider,
174 channel=scope.channel,
175 )
176 if result.is_err():
177 return Err(
178 RelayBillingError(
179 code=result.unwrap_err().code,
180 message=result.unwrap_err().message,
181 tenant_id=scope.tenant_id,
182 )
183 )
184 return Ok(result.unwrap().total)
186 async def settle(
187 self,
188 reservation: RelayUsageReservation,
189 result: RelayConvertResult,
190 *,
191 status: Literal["completed", "failed", "cancelled", "truncated"],
192 ) -> Result[RelayUsageRecord, RelayBillingError]:
193 """Settle actual usage for a request attempt exactly once.
195 The final usage comes from ``result.usage``. The attempt key is
196 the reservation identifier, so retries with the same reservation
197 map to the same ``(request_id, attempt_id)`` record and never
198 charge twice.
200 Args:
201 reservation: The pre-consume reservation for this attempt.
202 result: Converter result carrying the normalized final usage.
203 status: Terminal status of the attempt.
205 Returns:
206 Ok(record) with the settled usage and charge, or Err when
207 pricing or persistence fails. A missing final usage object
208 settles as zero usage with an explicit ``usage_missing``
209 metadata code; the prompt estimate is never billed.
210 """
211 scope = self._scopes.get(reservation.reservation_id)
212 if scope is None:
213 return Err(
214 RelayBillingError(
215 code="invalid_usage",
216 message="settle without pre_consume admission",
217 request_id=reservation.request_id,
218 )
219 )
220 usage, loss_codes = self._usage_for_result(result)
222 charge_result = self._estimator.estimate_charge(
223 scope.model,
224 usage,
225 provider=scope.provider,
226 channel=scope.channel,
227 )
228 if charge_result.is_err():
229 return Err(
230 RelayBillingError(
231 code=charge_result.unwrap_err().code,
232 message=charge_result.unwrap_err().message,
233 request_id=reservation.request_id,
234 tenant_id=scope.tenant_id,
235 )
236 )
237 charge = charge_result.unwrap().total
239 record = RelayUsageRecord(
240 request_id=reservation.request_id,
241 attempt_id=reservation.reservation_id,
242 scope=scope,
243 usage=usage,
244 charge=charge,
245 currency=self._currency,
246 status=status,
247 converter_id=result.converter_id,
248 loss_codes=loss_codes,
249 )
251 try:
252 stored = await self._store.settle_once(record)
253 except Exception as exc: # noqa: BLE001 - infrastructure boundary
254 logger.warning(
255 "relay_billing_settle_store_failed",
256 request_id=reservation.request_id,
257 error=str(exc),
258 )
259 return Err(
260 billing_store_unavailable(
261 message="could not persist settlement",
262 request_id=reservation.request_id,
263 tenant_id=scope.tenant_id,
264 )
265 )
267 await self._manager.settle(reservation.reservation_id)
268 logger.info(
269 "relay_billing_settled",
270 request_id=reservation.request_id,
271 attempt_id=stored.attempt_id,
272 charge=str(stored.charge),
273 tokens=usage.total_tokens,
274 )
275 return Ok(stored)
277 def _usage_for_result(
278 self,
279 result: RelayConvertResult,
280 ) -> tuple[RelayUsage, tuple[str, ...]]:
281 """Extract the usage and loss codes to bill for an attempt.
283 ``completed``, ``cancelled``, and ``truncated`` bill observed
284 usage; ``failed`` bills only provider-reported usage. A missing
285 usage object always settles as zero usage with an explicit
286 ``usage_missing`` loss code.
288 Args:
289 result: Converter result carrying normalized usage.
291 Returns:
292 Billed usage and the loss codes to attach to the record.
293 """
294 usage = result.usage or RelayUsage()
295 loss_codes = tuple(item.reason for item in result.losses)
296 if result.usage is None:
297 loss_codes = (*loss_codes, USAGE_MISSING)
298 return usage, loss_codes
300 async def release(self, reservation: RelayUsageReservation) -> None:
301 """Release a reservation when the request never reached upstream.
303 Args:
304 reservation: The reservation to release.
305 """
306 await self._manager.release(reservation.reservation_id)
307 await self._store.release(reservation.reservation_id)
308 self._scopes.pop(reservation.reservation_id, None)
311class RelayCostAdapter:
312 """Cost-tracking adapter over the relay billing store.
314 Preserves compatibility with :class:`CostTrackingProtocol` callers:
315 :meth:`track_cost` records a completed, generic usage record and
316 :meth:`get_budget` delegates to a configured tenant/account budget
317 resolver. Existing governance callers keep working without importing
318 relay types.
320 Args:
321 store: Persistence boundary for settled records.
322 currency: Default currency code.
323 budget: Optional callable returning the remaining budget for a
324 user; ``None`` means an unbounded budget of 0.
325 """
327 def __init__(
328 self,
329 *,
330 store: RelayUsageStoreProtocol,
331 currency: str = DEFAULT_CURRENCY,
332 budget: Callable[[str | None], Decimal] | None = None,
333 ) -> None:
334 self._store = store
335 self._currency = currency
336 self._budget = budget
338 async def track_cost(
339 self,
340 cost: float,
341 model: str,
342 user_id: str | None = None,
343 *,
344 tenant_id: str = "global",
345 ) -> None:
346 """Track a cost by recording a completed generic usage record.
348 Args:
349 cost: USD cost to record.
350 model: Model the cost belongs to.
351 user_id: User the cost belongs to, when applicable.
352 tenant_id: Tenant the charge is attributed to.
353 """
354 request_id = f"adapter:{identity.new_uuid()}"
355 record = RelayUsageRecord(
356 request_id=request_id,
357 attempt_id=request_id,
358 scope=RelayUsageScope(
359 tenant_id=tenant_id,
360 model=model,
361 user_id=user_id or None,
362 ),
363 usage=RelayUsage(),
364 charge=Decimal(str(cost)),
365 currency=self._currency,
366 status="completed",
367 )
368 await self._store.settle_once(record)
370 async def get_budget(self, user_id: str | None = None) -> Decimal:
371 """Return the configured remaining budget.
373 Args:
374 user_id: User the budget applies to.
376 Returns:
377 The resolver's remaining budget, or ``0`` when unconfigured.
378 """
379 if self._budget is None:
380 return Decimal(0)
381 return self._budget(user_id)