Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-relay-gateway/src/lexigram/ai/relay/gateway/buffered.py: 27%

110 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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"]