« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""DI registrations and gateway hooks for relay billing.
3The governance provider root registers the billing hierarchy by contract
4during ``register()`` and swaps the concrete lifecycle in ``boot()``:
6- :func:`register_relay_billing` binds ``RelayBillingProtocol`` to a
7 no-op admission policy (:class:`NoopRelayBilling`) when billing is
8 disabled, and to the placeholder instances the boot phase replaces
9 when billing is enabled.
10- :func:`boot_relay_billing` resolves the token counter, price
11 estimator, event bus, audit store, database, and reservation manager
12 through their contracts, builds ``RelayBillingService``, and rebinds
13 it behind :class:`RelayBillingHooks` so audit/budget emission and
14 gateway calls share one protocol binding.
16Nothing in this module imports gateway or converter implementations.
17"""
19from __future__ import annotations
21import asyncio
22from datetime import UTC, datetime, timedelta
23from decimal import Decimal
24import time
25from typing import TYPE_CHECKING, Literal
27from lexigram.ai.governance import GovernanceConfig
28from lexigram.ai.governance.budget import BudgetTracker
29from lexigram.ai.governance.relay_billing.models import RelayBillingConfig
30from lexigram.ai.governance.relay_billing.persistence import DatabaseRelayUsageStore
31from lexigram.ai.governance.relay_billing.pricing import SimpleCostEstimator
32from lexigram.ai.governance.relay_billing.reservations import RelayReservationManager
33from lexigram.ai.governance.relay_billing.service import RelayBillingService
34from lexigram.contracts.ai.governance import (
35 AIAuditEvent,
36 AIAuditStoreProtocol,
37 AuditEventType,
38 RelayBillingError,
39 RelayBillingProtocol,
40 RelayPriceEstimatorProtocol,
41 RelayUsageRecord,
42 RelayUsageReservation,
43 RelayUsageScope,
44 RelayUsageStoreProtocol,
45)
46from lexigram.contracts.ai.llm import (
47 CostEstimatorProtocol,
48 TokenCounterProtocol,
49)
50from lexigram.contracts.ai.relay import (
51 RelayConvertResult,
52 RelayRequestPayload,
53 RelayUsage,
54)
55from lexigram.contracts.core.result import Ok, Result
56from lexigram.contracts.data import DatabaseProviderProtocol
57from lexigram.contracts.events import EventBusProtocol
58from lexigram.logging import get_logger
60if TYPE_CHECKING:
61 from lexigram.contracts.core.di import (
62 BootContainerProtocol,
63 ContainerRegistrarProtocol,
64 )
66logger = get_logger(__name__)
68__all__ = [
69 "NoopRelayBilling",
70 "RelayBillingHooks",
71 "boot_relay_billing",
72 "register_relay_billing",
73]
76class NoopRelayBilling(RelayBillingProtocol):
77 """Admit-everything billing policy used when billing is disabled.
79 ``pre_consume`` always grants a zero-cost reservation, ``settle``
80 returns a zero-charge record built from the converter result, and
81 ``release`` does nothing. The gateway can therefore drive the same
82 protocol-shaped code path whether or not billing is enabled.
84 Attributes:
85 _scopes: Reservation identifier to scope map captured at
86 admission time (mirrors the real service for settle).
87 """
89 def __init__(self) -> None:
90 """Bind an empty scope map."""
91 self._scopes: dict[str, RelayUsageScope] = {}
93 async def pre_consume(
94 self,
95 request_id: str,
96 scope: RelayUsageScope,
97 payload: RelayRequestPayload,
98 ) -> Result[RelayUsageReservation, RelayBillingError]:
99 """Always admit with a zero-cost no-op reservation.
101 Args:
102 request_id: Gateway request identifier.
103 scope: Accounting scope of the request.
104 payload: Relay request payload (unused by the no-op policy).
106 Returns:
107 ``Ok`` with a non-expiring zero-cost reservation.
108 """
109 del payload
110 reservation_id = f"noop:{request_id}"
111 self._scopes[reservation_id] = scope
112 return Ok(
113 RelayUsageReservation(
114 reservation_id=reservation_id,
115 request_id=request_id,
116 estimated_tokens=0,
117 estimated_charge=Decimal(0),
118 expires_at=datetime.now(UTC) + timedelta(days=1),
119 )
120 )
122 async def settle(
123 self,
124 reservation: RelayUsageReservation,
125 result: RelayConvertResult,
126 *,
127 status: Literal["completed", "failed", "cancelled", "truncated"],
128 ) -> Result[RelayUsageRecord, RelayBillingError]:
129 """Return a zero-charge record for the attempt.
131 Args:
132 reservation: The no-op reservation issued by
133 :meth:`pre_consume`.
134 result: Converter result carrying the normalized usage.
135 status: Terminal status of the attempt.
137 Returns:
138 ``Ok`` with a zero-charge record (billing disabled).
139 """
140 scope = self._scopes.get(reservation.reservation_id) or RelayUsageScope(
141 tenant_id=reservation.request_id,
142 )
143 return Ok(
144 RelayUsageRecord(
145 request_id=reservation.request_id,
146 attempt_id=reservation.reservation_id,
147 scope=scope,
148 usage=result.usage or RelayUsage(),
149 charge=Decimal(0),
150 currency="USD",
151 status=status,
152 converter_id=result.converter_id,
153 loss_codes=tuple(loss.reason for loss in result.losses),
154 )
155 )
157 async def release(self, reservation: RelayUsageReservation) -> None:
158 """Do nothing; the no-op policy holds no capacity."""
159 self._scopes.pop(reservation.reservation_id, None)
162_TERMINAL_AUDIT_STATUS = {
163 "completed": "success",
164 "failed": "error",
165 "truncated": "truncated",
166 "cancelled": "cancelled",
167}
170class RelayBillingHooks(RelayBillingProtocol):
171 """Billing protocol wrapper that emits audit and budget events.
173 Delegates the lifecycle to the inner ``RelayBillingProtocol`` and,
174 after a successful settlement, records a redacted
175 :class:`~lexigram.contracts.ai.governance.AIAuditEvent` through the
176 optional audit store and feeds the settled usage into the optional
177 budget tracker (whose threshold alerts flow out over the event bus).
178 Both emissions happen on a background task so the caller never waits
179 on the persistence store.
181 Args:
182 delegate: The inner billing lifecycle (or a no-op policy).
183 audit_store: Audit event store; ``None`` disables audit events.
184 budget_tracker: Budget tracker emitting threshold alerts over
185 its event bus; ``None`` disables budget alerts.
186 """
188 def __init__(
189 self,
190 delegate: RelayBillingProtocol,
191 *,
192 audit_store: AIAuditStoreProtocol | None = None,
193 budget_tracker: BudgetTracker | None = None,
194 ) -> None:
195 """Bind the wrapper to the inner lifecycle and observers."""
196 self._delegate = delegate
197 self._audit_store = audit_store
198 self._budget_tracker = budget_tracker
199 self._started: dict[str, float] = {}
200 self._background_tasks: set[asyncio.Task[object]] = set()
202 async def pre_consume(
203 self,
204 request_id: str,
205 scope: RelayUsageScope,
206 payload: RelayRequestPayload,
207 ) -> Result[RelayUsageReservation, RelayBillingError]:
208 """Start a latency window and delegate admission to the service.
210 Args:
211 request_id: Gateway request identifier.
212 scope: Accounting scope of the request.
213 payload: Relay request payload used for prompt estimation.
215 Returns:
216 The inner admission result.
217 """
218 result = await self._delegate.pre_consume(request_id, scope, payload)
219 if result.is_ok():
220 reservation = result.unwrap()
221 self._started[reservation.reservation_id] = time.monotonic()
222 return result
224 async def settle(
225 self,
226 reservation: RelayUsageReservation,
227 result: RelayConvertResult,
228 *,
229 status: Literal["completed", "failed", "cancelled", "truncated"],
230 ) -> Result[RelayUsageRecord, RelayBillingError]:
231 """Settle through the inner service and emit observer events.
233 Args:
234 reservation: The pre-consume reservation for this attempt.
235 result: Converter result carrying the normalized usage.
236 status: Terminal status of the attempt.
238 Returns:
239 The settled record, or the inner error when settlement fails.
240 Observer events are scheduled only after a successful settle.
241 """
242 outcome = await self._delegate.settle(
243 reservation,
244 result,
245 status=status,
246 )
247 if outcome.is_err():
248 return outcome
249 record = outcome.unwrap()
250 latency_ms = self._started.pop(reservation.reservation_id, None)
251 if latency_ms is not None:
252 duration = max(0.0, time.monotonic() - latency_ms) * 1000.0
253 else:
254 duration = 0.0
255 self._schedule(self._emit(record, duration))
256 return outcome
258 async def release(self, reservation: RelayUsageReservation) -> None:
259 """Release through the inner service and drop latency state."""
260 self._started.pop(reservation.reservation_id, None)
261 await self._delegate.release(reservation)
263 def _schedule(self, coro: object) -> None:
264 """Start *coro* as a tracked background task (fire and forget)."""
265 loop = asyncio.get_running_loop()
266 task: asyncio.Task[object] = loop.create_task(coro) # type: ignore[arg-type]
267 self._background_tasks.add(task)
268 task.add_done_callback(self._background_tasks.discard)
270 async def _emit(
271 self,
272 record: RelayUsageRecord,
273 latency_ms: float,
274 ) -> None:
275 """Record the audit event and budget usage for one settlement.
277 Args:
278 record: The settled usage record.
279 latency_ms: Measured request latency, or ``0.0`` when unknown.
280 """
281 scope = record.scope
282 if self._audit_store is not None:
283 await self._audit_store.record(
284 AIAuditEvent(
285 event_type=AuditEventType.LLM_CALL,
286 model=scope.model,
287 provider=scope.provider,
288 user_id=scope.user_id,
289 status=_TERMINAL_AUDIT_STATUS.get(record.status, "success"),
290 tokens=record.usage.total_tokens,
291 cost=float(record.charge),
292 latency_ms=round(latency_ms, 2),
293 metadata={
294 "request_id": record.request_id,
295 "attempt_id": record.attempt_id,
296 "tenant_id": scope.tenant_id,
297 "account_id": scope.account_id,
298 "channel": scope.channel,
299 "converter_id": record.converter_id,
300 "loss_codes": list(record.loss_codes),
301 "currency": record.currency,
302 },
303 )
304 )
305 if self._budget_tracker is not None:
306 await self._budget_tracker.record_usage(
307 scope.model,
308 record.usage.total_tokens,
309 float(record.charge),
310 tenant_id=scope.tenant_id or None,
311 )
314def register_relay_billing(
315 container: ContainerRegistrarProtocol,
316 config: object,
317) -> None:
318 """Register the relay billing hierarchy by contract.
320 The root always exposes ``RelayBillingProtocol`` so the gateway can
321 resolve an admission policy even when billing is disabled. When
322 billing is enabled, the concrete lifecycle is built from resolved
323 contracts during :func:`boot_relay_billing`; until then the no-op
324 instance is the registered singleton the boot phase rebinds.
326 Args:
327 container: The container registrar to bind into.
328 config: Governance configuration; ``enabled`` gates billing.
329 """
330 relay_config = RelayBillingConfig()
331 container.singleton(RelayBillingConfig, relay_config)
332 if not isinstance(config, GovernanceConfig) or not config.enabled:
333 container.singleton(RelayBillingProtocol, NoopRelayBilling())
334 logger.info("relay_billing_disabled", reason="governance disabled")
335 return
337 container.singleton(RelayReservationManager, RelayReservationManager())
338 container.singleton(RelayBillingProtocol, NoopRelayBilling())
339 logger.info("relay_billing_registered")
342async def boot_relay_billing(
343 container: BootContainerProtocol,
344 config: object,
345) -> None:
346 """Build the live relay billing service and rebind it into the container.
348 Resolution is contract-scoped (database, estimator, token counter,
349 audit store, reservation manager). When a required contract is
350 missing (for example no database backend), the no-op admission
351 policy from :func:`register_relay_billing` remains bound and a
352 startup diagnostic is logged so the missing dependency is
353 discoverable.
355 Args:
356 container: The boot container used to resolve contracts.
357 config: Relay configuration driving the bootstrap.
358 """
359 if not isinstance(config, GovernanceConfig) or not config.enabled:
360 logger.info("relay_billing_boot_skipped", reason="governance disabled")
361 return
363 database = await container.resolve_optional(DatabaseProviderProtocol)
364 estimator = await container.resolve_optional(RelayPriceEstimatorProtocol)
365 token_counter = await container.resolve_optional(TokenCounterProtocol)
366 event_bus = await container.resolve_optional(EventBusProtocol)
368 if database is None:
369 logger.warning(
370 "relay_billing_missing_dependency",
371 missing="DatabaseProviderProtocol",
372 )
373 return
374 if estimator is None:
375 cost_estimator = await container.resolve_optional(CostEstimatorProtocol)
376 if cost_estimator is not None:
377 estimator = SimpleCostEstimator(cost_estimator)
378 if estimator is None:
379 logger.warning(
380 "relay_billing_missing_dependency",
381 missing="RelayPriceEstimatorProtocol",
382 )
383 return
385 store: RelayUsageStoreProtocol = DatabaseRelayUsageStore(database)
386 manager = await container.resolve(RelayReservationManager)
387 relay_config = await container.resolve(RelayBillingConfig)
388 budget_tracker = (
389 BudgetTracker(
390 tpm_limit=config.tpm_limit,
391 cost_limit_hourly=config.monthly_budget,
392 event_bus=event_bus,
393 )
394 if event_bus is not None
395 else None
396 )
397 audit_store = await container.resolve_optional(AIAuditStoreProtocol)
399 billing = RelayBillingService(
400 reservation_manager=manager,
401 estimator=estimator,
402 store=store,
403 token_counter=token_counter,
404 currency=relay_config.currency,
405 )
406 hooks = RelayBillingHooks(
407 billing,
408 audit_store=audit_store,
409 budget_tracker=budget_tracker,
410 )
411 container.bind(RelayBillingProtocol, hooks) # type: ignore[type-abstract]
412 logger.info("relay_billing_booted")