« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Billing lifecycle helpers for the relay gateway service.
3Admission control reserves capacity before conversion and upstream I/O;
4settlement runs exactly once per attempt after the terminal outcome is
5known. Settlement failures are logged and never fail the response.
6"""
8from __future__ import annotations
10from collections.abc import Mapping
11from typing import Any, Literal
13from lexigram.ai.relay.gateway.codec import RelayPayloadCodec
14from lexigram.ai.relay.gateway.errors import billing_error_to_gateway
15from lexigram.contracts.ai.governance import (
16 RelayBillingProtocol,
17 RelayUsageReservation,
18 RelayUsageScope,
19)
20from lexigram.contracts.ai.relay import (
21 ConversionQuality,
22 RelayChannel,
23 RelayConvertResult,
24 RelayGatewayError,
25 RelayGatewayRequest,
26 RelayStreamSessionProtocol,
27 RelayUsage,
28)
29from lexigram.contracts.core.result import Err, Ok, Result
30from lexigram.logging import get_logger
31from lexigram.serialization import dumps
33__all__ = [
34 "StreamStatus",
35 "empty_settle_result",
36 "pre_consume",
37 "settle",
38 "stream_settle_result",
39 "usage_from_snapshot",
40]
42logger = get_logger(__name__)
44StreamStatus = Literal["completed", "failed", "cancelled", "truncated"]
47def usage_from_snapshot(snapshot: object) -> RelayUsage | None:
48 """Extract ``RelayUsage`` from a session snapshot when present.
50 Snapshots are opaque; only a ``Mapping`` carrying a ``usage``
51 sub-mapping is inspected, accepting either the OpenAI-style
52 ``prompt_tokens``/``completion_tokens`` keys or the Claude-style
53 ``input_tokens``/``output_tokens`` keys.
55 Args:
56 snapshot: The session snapshot returned by
57 ``RelayStreamSessionProtocol.snapshot``.
59 Returns:
60 A normalized ``RelayUsage`` when the snapshot exposes one,
61 else ``None`` (settlement then records no usage).
62 """
63 if not isinstance(snapshot, Mapping):
64 return None
65 usage = snapshot.get("usage")
66 if not isinstance(usage, Mapping):
67 return None
68 prompt = usage.get("prompt_tokens", usage.get("input_tokens"))
69 completion = usage.get("completion_tokens", usage.get("output_tokens"))
70 if not isinstance(prompt, int) or not isinstance(completion, int):
71 return None
72 return RelayUsage(
73 prompt_tokens=prompt,
74 completion_tokens=completion,
75 )
78def empty_settle_result(
79 request: RelayGatewayRequest, channel: RelayChannel
80) -> RelayConvertResult[Any]:
81 """Build the usage-free result settled for failed attempts.
83 Args:
84 request: The gateway request being dispatched.
85 channel: The selected channel.
87 Returns:
88 A ``RelayConvertResult`` carrying no usage so the billing
89 pipeline records an attempted-but-unbilled attempt.
90 """
91 return RelayConvertResult(
92 value=None,
93 source=request.source,
94 target=channel.target_format,
95 converter_id="",
96 quality=ConversionQuality.GOOD,
97 )
100def stream_settle_result(
101 request: RelayGatewayRequest,
102 channel: RelayChannel,
103 *,
104 converter_id: str,
105 session: RelayStreamSessionProtocol,
106) -> RelayConvertResult[Any]:
107 """Build the settled result from the stream session snapshot.
109 Args:
110 request: The gateway request being settled.
111 channel: The selected channel.
112 converter_id: Converter that produced the stream session.
113 session: The stream session whose snapshot carries the settled
114 usage.
116 Returns:
117 A ``RelayConvertResult`` carrying normalized usage extracted
118 from the session snapshot (or no usage when the snapshot
119 exposes none).
120 """
121 usage = usage_from_snapshot(session.snapshot())
122 return RelayConvertResult(
123 value=None,
124 source=request.source,
125 target=channel.target_format,
126 converter_id=converter_id,
127 quality=ConversionQuality.GOOD,
128 usage=usage,
129 )
132async def pre_consume(
133 codec: RelayPayloadCodec,
134 request: RelayGatewayRequest,
135 billing: RelayBillingProtocol,
136 channel: RelayChannel,
137) -> Result[RelayUsageReservation, RelayGatewayError]:
138 """Reserve billing capacity before conversion and upstream I/O.
140 The inbound payload is re-decoded into a typed request DTO so the
141 billing pipeline can estimate prompt and output budgets; a payload
142 that rejects decoding fails the request here, before any upstream
143 I/O. Billing denials short-circuit the pipeline and are classified
144 through :func:`billing_error_to_gateway`.
146 Args:
147 codec: Wire DTO decoder used to re-decode the inbound payload.
148 request: The gateway request being dispatched.
149 billing: The billing lifecycle to reserve through.
150 channel: The selected channel.
152 Returns:
153 ``Ok(reservation)`` when admission is proven, or
154 ``Err(RelayGatewayError)`` carrying the classified failure.
155 """
156 scope = RelayUsageScope(
157 tenant_id=request.tenant_id,
158 model=request.model,
159 channel=channel.name,
160 )
161 dto = codec.decode_request(
162 source=request.source,
163 raw=dumps(dict(request.payload)),
164 request_id=request.request_id,
165 )
166 if dto.is_err():
167 return Err(dto.unwrap_err())
168 admitted = await billing.pre_consume(request.request_id, scope, dto.unwrap())
169 if admitted.is_err():
170 error = admitted.unwrap_err()
171 logger.warning(
172 "relay_gateway_billing_denied",
173 request_id=request.request_id,
174 channel=channel.name,
175 code=error.code,
176 error=error.message,
177 )
178 return Err(billing_error_to_gateway(error, request.request_id))
179 return Ok(admitted.unwrap())
182async def settle(
183 billing: RelayBillingProtocol,
184 reservation: RelayUsageReservation,
185 result: RelayConvertResult[Any],
186 *,
187 status: StreamStatus,
188) -> None:
189 """Settle the reservation exactly once without failing the response.
191 Settlement failures are logged and never propagate: the response
192 path has already completed by the time accounting runs.
194 Args:
195 billing: The billing lifecycle to settle through.
196 reservation: The reservation granted by ``pre_consume``.
197 result: The conversion result carrying settled usage, or an
198 empty result when the attempt produced no billable usage.
199 status: Terminal lifecycle status of the attempt.
200 """
201 settled = await billing.settle(reservation, result, status=status)
202 if settled.is_err():
203 error = settled.unwrap_err()
204 logger.warning(
205 "relay_gateway_settle_failed",
206 request_id=reservation.request_id,
207 status=status,
208 code=error.code,
209 error=error.message,
210 )