Coverage for src / lexigram / contracts / ai / governance / relay_billing.py: 32%
114 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-19 05:41 +0800
1"""Relay usage, reservation, and billing contracts.
3Shared accounting protocols and value types for relay gateway billing:
4usage scopes, reservations, settlement records, charge breakdowns, and
5the store/estimator/billing protocols implemented by the governance
6package. All monetary fields use non-negative ``Decimal`` values; the
7contracts never carry raw request payloads.
8"""
10from __future__ import annotations
12from collections.abc import Mapping, Sequence
13from dataclasses import dataclass
14from datetime import UTC, datetime
15from decimal import Decimal
16from typing import Literal, Protocol, runtime_checkable
18from lexigram.contracts.ai.relay import (
19 JsonValue,
20 RelayConvertResult,
21 RelayRequestPayload,
22 RelayUsage,
23)
24from lexigram.contracts.core.result import Result
26__all__ = [
27 "RelayBillingError",
28 "RelayBillingProtocol",
29 "RelayChargeBreakdown",
30 "RelayPriceEstimatorProtocol",
31 "RelayUsageRecord",
32 "RelayUsageReservation",
33 "RelayUsageScope",
34 "RelayUsageStoreProtocol",
35 "billing_store_unavailable",
36 "charge_overflow",
37 "duplicate_settlement",
38 "invalid_usage",
39 "quota_exhausted",
40 "reservation_expired",
41 "unknown_price",
42]
44_RELAY_USAGE_TOKEN_FIELDS = (
45 "prompt_tokens",
46 "completion_tokens",
47 "cache_read_tokens",
48 "cache_creation_tokens",
49 "reasoning_tokens",
50 "audio_input_tokens",
51 "audio_output_tokens",
52 "image_tokens",
53 "input_tokens",
54 "output_tokens",
55)
58@dataclass(frozen=True, slots=True)
59class RelayUsageScope:
60 """Accounting scope for one relay request.
62 Attributes:
63 tenant_id: Tenant the charge is attributed to (required).
64 account_id: Account the charge is attributed to, when applicable.
65 user_id: User the charge is attributed to, when applicable.
66 model: Model alias used for the request.
67 provider: Provider name for the request.
68 channel: Channel name for the request.
69 """
71 tenant_id: str
72 account_id: str | None = None
73 user_id: str | None = None
74 model: str = ""
75 provider: str = ""
76 channel: str = ""
78 def __post_init__(self) -> None:
79 """Reject scopes without a tenant."""
80 if not self.tenant_id:
81 raise ValueError("tenant_id must be non-empty")
84@dataclass(frozen=True, slots=True)
85class RelayUsageRecord:
86 """The settled accounting record for one request attempt.
88 Attributes:
89 request_id: Gateway request identifier.
90 attempt_id: Attempt key; ``(request_id, attempt_id)`` is unique.
91 scope: Accounting scope of the request.
92 usage: Actual normalized usage that was billed.
93 charge: Final charge in ``currency`` units (never negative).
94 currency: ISO currency code of the charge.
95 status: Terminal lifecycle status of the attempt.
96 converter_id: Converter that produced the usage, when known.
97 loss_codes: Conversion loss codes recorded alongside usage.
98 """
100 request_id: str
101 attempt_id: str
102 scope: RelayUsageScope
103 usage: RelayUsage
104 charge: Decimal
105 currency: str
106 status: Literal["completed", "failed", "cancelled", "truncated"]
107 converter_id: str | None = None
108 loss_codes: tuple[str, ...] = ()
110 def __post_init__(self) -> None:
111 """Reject invalid identities, usage, charges, and statuses."""
112 if not self.request_id:
113 raise ValueError("request_id must be non-empty")
114 if not self.attempt_id:
115 raise ValueError("attempt_id must be non-empty")
116 for field_name in _RELAY_USAGE_TOKEN_FIELDS:
117 if getattr(self.usage, field_name) < 0:
118 raise ValueError(f"usage.{field_name} must be non-negative")
119 if (
120 self.usage.total_tokens_override is not None
121 and self.usage.total_tokens_override < 0
122 ):
123 raise ValueError("usage.total_tokens_override must be non-negative")
124 if self.charge < 0:
125 raise ValueError("charge must be non-negative")
126 if not self.currency:
127 raise ValueError("currency must be non-empty")
128 if self.status not in ("completed", "failed", "cancelled", "truncated"):
129 raise ValueError(f"unknown billing status: {self.status}")
132@dataclass(frozen=True, slots=True)
133class RelayUsageReservation:
134 """Capacity reserved for a request before upstream admission.
136 Attributes:
137 reservation_id: Unique reservation identifier.
138 request_id: Gateway request identifier the reservation belongs to.
139 estimated_tokens: Prompt estimate used for admission only; never
140 billed as final usage.
141 estimated_charge: Maximum charge covered by the reservation.
142 expires_at: Expiry instant (timezone-aware); a non-positive TTL
143 is rejected at construction.
144 """
146 reservation_id: str
147 request_id: str
148 estimated_tokens: int
149 estimated_charge: Decimal
150 expires_at: datetime
152 def __post_init__(self) -> None:
153 """Reject invalid identities, estimates, and expiries."""
154 if not self.reservation_id:
155 raise ValueError("reservation_id must be non-empty")
156 if not self.request_id:
157 raise ValueError("request_id must be non-empty")
158 if self.estimated_tokens < 0:
159 raise ValueError("estimated_tokens must be non-negative")
160 if self.estimated_charge < 0:
161 raise ValueError("estimated_charge must be non-negative")
162 if self.expires_at <= datetime.now(UTC):
163 raise ValueError("reservation TTL must be positive")
166@dataclass(frozen=True, slots=True)
167class RelayChargeBreakdown:
168 """Per-dimension charge breakdown of one estimate or settlement.
170 Attributes:
171 prompt: Charge for input tokens.
172 cached_prompt: Charge for cached input tokens.
173 completion: Charge for output tokens.
174 reasoning: Charge for reasoning tokens.
175 audio_input: Charge for audio input tokens.
176 audio_output: Charge for audio output tokens.
177 image: Charge for image input tokens.
178 total: Total charge; must equal the configured sum after rounding.
179 """
181 prompt: Decimal
182 cached_prompt: Decimal
183 completion: Decimal
184 reasoning: Decimal
185 audio_input: Decimal
186 audio_output: Decimal
187 image: Decimal
188 total: Decimal
190 def __post_init__(self) -> None:
191 """Reject negative prices in any dimension."""
192 for field_name in (
193 "prompt",
194 "cached_prompt",
195 "completion",
196 "reasoning",
197 "audio_input",
198 "audio_output",
199 "image",
200 "total",
201 ):
202 if getattr(self, field_name) < 0:
203 raise ValueError(f"breakdown.{field_name} must be non-negative")
206@dataclass(frozen=True, slots=True)
207class RelayBillingError:
208 """A domain error returned from billing operations.
210 Attributes:
211 code: Machine-readable error code (e.g. ``unknown_price``).
212 message: Public, redaction-safe error message. Never contains
213 raw request payloads, prompts, or credentials.
214 request_id: Request identifier the error applies to, when known.
215 tenant_id: Tenant the error applies to, when known.
216 """
218 code: str
219 message: str
220 request_id: str | None = None
221 tenant_id: str | None = None
224@runtime_checkable
225class RelayUsageStoreProtocol(Protocol):
226 """Persistence boundary for reservations and settled usage records."""
228 async def save_reservation(self, reservation: RelayUsageReservation) -> None:
229 """Persist a reservation atomically."""
231 async def settle_once(self, record: RelayUsageRecord) -> RelayUsageRecord:
232 """Settle a record exactly once, returning the stored record."""
234 async def release(self, reservation_id: str) -> None:
235 """Release a reservation, marking it as released."""
237 async def query(
238 self, filters: Mapping[str, JsonValue]
239 ) -> Sequence[RelayUsageRecord]:
240 """Query settled records by scope filters."""
243@runtime_checkable
244class RelayPriceEstimatorProtocol(Protocol):
245 """Computes charges from normalized usage and configured prices.
247 Pricing failures (unknown model, malformed price expressions,
248 overflow) are returned as ``Err(RelayBillingError)`` — estimators
249 never clamp silently.
250 """
252 def estimate_charge(
253 self,
254 model: str,
255 usage: RelayUsage,
256 *,
257 provider: str = "",
258 channel: str = "",
259 ) -> Result[RelayChargeBreakdown, RelayBillingError]:
260 """Return the charge breakdown for *usage* on *model*."""
263@runtime_checkable
264class RelayBillingProtocol(Protocol):
265 """Billing lifecycle for one relay request.
267 Callers invoke ``pre_consume`` after channel selection and before
268 upstream I/O, then ``settle`` on completion, stream finalization, or
269 failure, or ``release`` when the request never reached upstream.
270 """
272 async def pre_consume(
273 self,
274 request_id: str,
275 scope: RelayUsageScope,
276 payload: RelayRequestPayload,
277 ) -> Result[RelayUsageReservation, RelayBillingError]:
278 """Reserve capacity for a request before upstream admission."""
280 async def settle(
281 self,
282 reservation: RelayUsageReservation,
283 result: RelayConvertResult[RelayRequestPayload],
284 *,
285 status: Literal["completed", "failed", "cancelled", "truncated"],
286 ) -> Result[RelayUsageRecord, RelayBillingError]:
287 """Settle actual usage for a request attempt exactly once."""
289 async def release(self, reservation: RelayUsageReservation) -> None:
290 """Release a reservation when the request never reached upstream."""
293def unknown_price(
294 *,
295 message: str = "unknown model price",
296 request_id: str | None = None,
297 tenant_id: str | None = None,
298) -> RelayBillingError:
299 """Return an ``unknown_price`` billing error."""
300 return RelayBillingError(
301 code="unknown_price",
302 message=message,
303 request_id=request_id,
304 tenant_id=tenant_id,
305 )
308def quota_exhausted(
309 *,
310 message: str = "quota exhausted",
311 request_id: str | None = None,
312 tenant_id: str | None = None,
313) -> RelayBillingError:
314 """Return a ``quota_exhausted`` billing error."""
315 return RelayBillingError(
316 code="quota_exhausted",
317 message=message,
318 request_id=request_id,
319 tenant_id=tenant_id,
320 )
323def reservation_expired(
324 *,
325 message: str = "reservation expired",
326 request_id: str | None = None,
327 tenant_id: str | None = None,
328) -> RelayBillingError:
329 """Return a ``reservation_expired`` billing error."""
330 return RelayBillingError(
331 code="reservation_expired",
332 message=message,
333 request_id=request_id,
334 tenant_id=tenant_id,
335 )
338def duplicate_settlement(
339 *,
340 message: str = "duplicate settlement",
341 request_id: str | None = None,
342 tenant_id: str | None = None,
343) -> RelayBillingError:
344 """Return a ``duplicate_settlement`` billing error."""
345 return RelayBillingError(
346 code="duplicate_settlement",
347 message=message,
348 request_id=request_id,
349 tenant_id=tenant_id,
350 )
353def invalid_usage(
354 *,
355 message: str = "invalid usage",
356 request_id: str | None = None,
357 tenant_id: str | None = None,
358) -> RelayBillingError:
359 """Return an ``invalid_usage`` billing error."""
360 return RelayBillingError(
361 code="invalid_usage",
362 message=message,
363 request_id=request_id,
364 tenant_id=tenant_id,
365 )
368def billing_store_unavailable(
369 *,
370 message: str = "billing store unavailable",
371 request_id: str | None = None,
372 tenant_id: str | None = None,
373) -> RelayBillingError:
374 """Return a ``billing_store_unavailable`` billing error."""
375 return RelayBillingError(
376 code="billing_store_unavailable",
377 message=message,
378 request_id=request_id,
379 tenant_id=tenant_id,
380 )
383def charge_overflow(
384 *,
385 message: str = "charge overflow",
386 request_id: str | None = None,
387 tenant_id: str | None = None,
388) -> RelayBillingError:
389 """Return a ``charge_overflow`` billing error."""
390 return RelayBillingError(
391 code="charge_overflow",
392 message=message,
393 request_id=request_id,
394 tenant_id=tenant_id,
395 )