1"""Buffered dispatch pipeline for :class:`RelayGatewayService`.
2
3Runs the ordered dependency pipeline for one non-streaming request:
4authorization, channel selection with retry/failover, billing admission,
5request conversion, the protected upstream call, response decoding and
6conversion, billing settlement, and result metadata assembly.
7"""
8
9from __future__ import annotations
10
11from typing import cast
12
13from lexigram.ai.relay.gateway.channels import RelayChannelRegistry
14from lexigram.ai.relay.gateway.codec import RelayPayloadCodec
15from lexigram.ai.relay.gateway.config import RelayGatewayConfig
16from lexigram.ai.relay.gateway.errors import (
17 auth_denied,
18 conversion_error_to_gateway,
19 with_request_id,
20)
21from lexigram.ai.relay.gateway.operations import billing as billing_ops
22from lexigram.ai.relay.gateway.operations import telemetry
23from lexigram.ai.relay.gateway.operations import upstream as upstream_ops
24from lexigram.ai.relay.gateway.operations.failover import RelayFailoverTracker
25from lexigram.ai.relay.gateway.upstream import HTTPUpstreamAdapter
26from lexigram.contracts.ai.governance import RelayBillingProtocol, RelayUsageReservation
27from lexigram.contracts.ai.relay import (
28 MediaResolverProtocol,
29 RelayChannel,
30 RelayConversionContext,
31 RelayConverterProtocol,
32 RelayGatewayError,
33 RelayGatewayMetadata,
34 RelayGatewayRequest,
35 RelayGatewayResult,
36 RelayOptions,
37 RelayRequestPayload,
38)
39from lexigram.contracts.ai.relay.gateway import RelayGatewayErrorCode
40from lexigram.contracts.auth.guard import AuthorizerProtocol
41from lexigram.contracts.core.result import Err, Ok, Result
42from lexigram.logging import get_logger
43
44logger = get_logger(__name__)
45
46
47class BufferedDispatchMixin:
48 """Buffered request pipeline shared into ``RelayGatewayService``.
49
50 Requires the host service to provide the gateway dependencies as
51 private attributes (``_converter``, ``_codec``, ``_registry``,
52 ``_upstream``, ``_config``, ``_authorizer``, ``_billing``,
53 ``_media_resolver``, ``_failover``).
54 """
55
56 _converter: RelayConverterProtocol
57 _codec: RelayPayloadCodec
58 _registry: RelayChannelRegistry
59 _upstream: HTTPUpstreamAdapter
60 _config: RelayGatewayConfig
61 _authorizer: AuthorizerProtocol | None
62 _billing: RelayBillingProtocol | None
63 _media_resolver: MediaResolverProtocol | None
64 _failover: RelayFailoverTracker | None
65
66 async def _dispatch(
67 self, request: RelayGatewayRequest
68 ) -> tuple[Result[RelayGatewayResult, RelayGatewayError], str]:
69 """Run the ordered dependency pipeline for one request.
70
71 Returns:
72 ``tuple`` of the pipeline result and the selected channel
73 name. The channel name is ``""`` when selection failed
74 before a channel was chosen.
75 """
76 if self._authorizer is not None:
77 allowed = await self._authorizer.authorize(
78 user=request.tenant_id,
79 action="relay.invoke",
80 resource=request.model,
81 )
82 if not allowed:
83 return Err(auth_denied(request.request_id)), ""
84 max_attempts = self._config.max_upstream_retries + 1
85 tried: set[str] = set()
86 last_upstream_error: RelayGatewayError | None = None
87 last_channel_name = ""
88 last_channel: RelayChannel | None = None
89 held_reservation: RelayUsageReservation | None = None
90 for attempt in range(1, max_attempts + 1):
91 selected = self._registry.select(
92 source=request.source,
93 model=request.model,
94 stream=request.stream,
95 preferred=request.channel.name if request.channel else None,
96 exclude=frozenset(tried),
97 )
98 if selected.is_err():
99 if last_upstream_error is not None:
100 billing = self._billing
101 if (
102 held_reservation is not None
103 and billing is not None
104 and last_channel is not None
105 ):
106 await billing_ops.settle(
107 billing,
108 held_reservation,
109 billing_ops.empty_settle_result(request, last_channel),
110 status="failed",
111 )
112 return (
113 Err(with_request_id(last_upstream_error, request.request_id)),
114 last_channel_name,
115 )
116 return (
117 Err(with_request_id(selected.unwrap_err(), request.request_id)),
118 "",
119 )
120 channel = selected.unwrap()
121 last_channel = channel
122 last_channel_name = channel.name
123 logger.info(
124 "relay_gateway_channel_selected",
125 request_id=request.request_id,
126 channel=channel.name,
127 target_format=channel.target_format,
128 model=request.model,
129 )
130 billing = self._billing
131 reservation: RelayUsageReservation | None = None
132 if billing is not None:
133 admitted = await billing_ops.pre_consume(
134 self._codec, request, billing, channel
135 )
136 if admitted.is_err():
137 return Err(admitted.unwrap_err()), channel.name
138 reservation = admitted.unwrap()
139 if held_reservation is not None:
140 await billing.release(held_reservation)
141 held_reservation = None
142 outbound_model = upstream_ops.outbound_model(
143 self._config, channel, request.model
144 )
145 context = RelayConversionContext(
146 request_id=request.request_id,
147 channel_name=channel.name,
148 upstream_model=outbound_model,
149 options=RelayOptions(),
150 media_resolver=self._media_resolver,
151 )
152 conv = self._converter.convert_request(
153 payload=cast("RelayRequestPayload", request.payload),
154 source=request.source,
155 target=channel.target_format,
156 context=context,
157 )
158 if conv.is_err():
159 if reservation is not None and billing is not None:
160 await billing.release(reservation)
161 return (
162 Err(
163 conversion_error_to_gateway(
164 conv.unwrap_err(), request.request_id
165 )
166 ),
167 channel.name,
168 )
169 converted_request = conv.unwrap()
170 telemetry.log_conversion_loss(
171 request.request_id,
172 converted_request.converter_id,
173 converted_request.losses,
174 )
175 upstream_response = await upstream_ops.call_upstream(
176 self._upstream,
177 channel,
178 outbound_model,
179 converted_request.value.to_dict(),
180 request,
181 )
182 if upstream_response.is_err():
183 upstream_error = upstream_response.unwrap_err()
184 if upstream_ops.should_track_upstream_failure(upstream_error.code):
185 upstream_ops.note_failure(self._failover, channel.name)
186 if upstream_error.retryable and attempt < max_attempts:
187 tried.add(channel.name)
188 last_upstream_error = upstream_error
189 held_reservation = reservation
190 logger.info(
191 "relay_gateway_upstream_retry",
192 request_id=request.request_id,
193 channel=channel.name,
194 error_code=upstream_error.code,
195 attempt=attempt,
196 )
197 continue
198 if reservation is not None and billing is not None:
199 await billing_ops.settle(
200 billing,
201 reservation,
202 billing_ops.empty_settle_result(request, channel),
203 status="failed",
204 )
205 return (
206 Err(with_request_id(upstream_error, request.request_id)),
207 channel.name,
208 )
209 resp = upstream_response.unwrap()
210 if resp.payload is None:
211 if reservation is not None and billing is not None:
212 await billing_ops.settle(
213 billing,
214 reservation,
215 billing_ops.empty_settle_result(request, channel),
216 status="failed",
217 )
218 return (
219 Err(
220 RelayGatewayError(
221 code=RelayGatewayErrorCode.UPSTREAM_MALFORMED,
222 message="malformed upstream response",
223 status_code=502,
224 request_id=request.request_id,
225 retryable=False,
226 )
227 ),
228 channel.name,
229 )
230 decoded = self._codec.decode_response_payload(
231 target=channel.target_format,
232 data=dict(resp.payload),
233 request_id=request.request_id,
234 )
235 if decoded.is_err():
236 if reservation is not None and billing is not None:
237 await billing_ops.settle(
238 billing,
239 reservation,
240 billing_ops.empty_settle_result(request, channel),
241 status="failed",
242 )
243 return Err(decoded.unwrap_err()), channel.name
244 back = self._converter.convert_response(
245 payload=decoded.unwrap(),
246 source=channel.target_format,
247 target=request.source,
248 context=context,
249 )
250 if back.is_err():
251 if reservation is not None and billing is not None:
252 await billing_ops.settle(
253 billing,
254 reservation,
255 billing_ops.empty_settle_result(request, channel),
256 status="failed",
257 )
258 return (
259 Err(
260 conversion_error_to_gateway(
261 back.unwrap_err(), request.request_id
262 )
263 ),
264 channel.name,
265 )
266 converted = back.unwrap()
267 telemetry.log_conversion_loss(
268 request.request_id, converted.converter_id, converted.losses
269 )
270 if reservation is not None and billing is not None:
271 await billing_ops.settle(
272 billing, reservation, converted, status="completed"
273 )
274 upstream_ops.note_success(self._failover, channel.name)
275 metadata = RelayGatewayMetadata(
276 converter_id=converted.converter_id,
277 source=request.source,
278 target=channel.target_format,
279 quality=converted.quality,
280 loss_codes=tuple(loss.reason for loss in converted.losses),
281 warnings=converted.warnings,
282 )
283 return (
284 Ok(
285 RelayGatewayResult(
286 status_code=resp.status_code,
287 headers={**resp.headers, "x-request-id": request.request_id},
288 payload=converted.value.to_dict(),
289 stream=None,
290 metadata=metadata,
291 )
292 ),
293 channel.name,
294 )
295 error = last_upstream_error or RelayGatewayError(
296 code=RelayGatewayErrorCode.CHANNEL_DISABLED,
297 message="no channels available",
298 status_code=404,
299 request_id=request.request_id,
300 )
301 return Err(with_request_id(error, request.request_id)), last_channel_name
302
303
304__all__ = ["BufferedDispatchMixin"]