Coverage for src / lexigram / ai / relay / gateway / upstream.py: 100%

93 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-08-08 23:08 +0800

1"""HTTP upstream transport for the relay gateway. 

2 

3Sends resolved ``UpstreamRequest`` values to model providers through the 

4contracts-level ``HTTPClientProtocol`` and maps transport outcomes into 

5typed ``RelayGatewayError`` or ``Ok(UpstreamResponse)`` results. 

6 

7The transport protocol keeps ``lexigram-http`` behind the DI boundary: 

8this module never imports it and classifies failures from stdlib and 

9contracts-level exceptions only. 

10""" 

11 

12from __future__ import annotations 

13 

14import asyncio 

15from collections.abc import AsyncIterator, Iterator 

16 

17from lexigram.contracts.ai.relay import ( 

18 RelayGatewayError, 

19 UpstreamChunk, 

20 UpstreamRequest, 

21 UpstreamResponse, 

22) 

23from lexigram.contracts.core.result import Err, Ok, Result 

24from lexigram.contracts.exceptions import InfrastructureError 

25from lexigram.contracts.web import HTTPClientProtocol, HttpResponse 

26from lexigram.serialization import dumps_str, loads 

27 

28__all__ = ["HTTPUpstreamAdapter"] 

29 

30 

31class HTTPUpstreamAdapter: 

32 """Sends upstream requests through an injected HTTP client. 

33 

34 The adapter classifies failures from stdlib and contracts-level 

35 exceptions only (by design); implementations of 

36 ``HTTPClientProtocol`` keep transport libraries like 

37 ``lexigram-http`` behind the DI boundary. 

38 

39 Attributes: 

40 _http: The injected HTTP client resolved from DI. 

41 _cancelled: Request identifiers whose streaming cancel was 

42 observed (after-the-fact only; outbound frames are not 

43 interrupted). 

44 """ 

45 

46 def __init__(self, http: HTTPClientProtocol) -> None: 

47 """Bind the adapter to an HTTP client. 

48 

49 Args: 

50 http: Any ``HTTPClientProtocol`` implementation driving the 

51 outbound request. 

52 """ 

53 self._http = http 

54 self._cancelled: set[str] = set() 

55 

56 async def request( 

57 self, request: UpstreamRequest 

58 ) -> Result[UpstreamResponse, RelayGatewayError]: 

59 """Send *request* upstream and classify the outcome. 

60 

61 Method, URL, headers, JSON payload, and timeout come straight 

62 from the ``UpstreamRequest``. 2xx responses are decoded into a 

63 typed ``UpstreamResponse`` (empty bodies yield a ``None`` 

64 payload); non-2xx responses and transport failures map to 

65 ``RelayGatewayError`` values that never carry raw bodies, 

66 headers, or credentials. 

67 

68 Args: 

69 request: Fully-resolved upstream request. 

70 

71 Returns: 

72 ``Ok(UpstreamResponse)`` for 2xx responses (the response 

73 headers are preserved verbatim), or ``Err`` classifying 

74 transport cancellation (``UPSTREAM_CANCELLED``, 499), 

75 timeouts (``UPSTREAM_TIMEOUT``, 504), generic transport 

76 failures (``UPSTREAM_FAILED``, 502), malformed 2xx bodies 

77 (``UPSTREAM_MALFORMED``, 502), and non-2xx responses 

78 (``UPSTREAM_ERROR`` with a safe public message). 

79 """ 

80 try: 

81 response = await self._http.request( 

82 method=request.method, 

83 url=request.url, 

84 headers=dict(request.headers), 

85 json=request.payload, 

86 timeout=request.timeout_seconds, 

87 channel_name=request.channel_name, 

88 ) 

89 except asyncio.CancelledError: 

90 return Err( 

91 RelayGatewayError( 

92 code="UPSTREAM_CANCELLED", 

93 message="upstream request cancelled", 

94 status_code=499, 

95 request_id=request.request_id, 

96 ) 

97 ) 

98 except TimeoutError: 

99 return Err( 

100 RelayGatewayError( 

101 code="UPSTREAM_TIMEOUT", 

102 message="upstream request timed out", 

103 status_code=504, 

104 request_id=request.request_id, 

105 retryable=True, 

106 ) 

107 ) 

108 except InfrastructureError: 

109 return Err( 

110 RelayGatewayError( 

111 code="UPSTREAM_FAILED", 

112 message="upstream transport failure", 

113 status_code=502, 

114 request_id=request.request_id, 

115 retryable=True, 

116 ) 

117 ) 

118 if 200 <= response.status < 300: 

119 return self._decode_success(request, response) 

120 return self._decode_error(request, response) 

121 

122 async def stream(self, request: UpstreamRequest) -> AsyncIterator[UpstreamChunk]: 

123 """Consume one upstream SSE response as a stream of chunks. 

124 

125 The whole stream arrives through a single ``request`` call whose 

126 body is parsed into ``data:`` frames; frames are emitted one by 

127 one as the consumer iterates. 2xx responses are parsed into 

128 chunk frames; non-2xx responses and transport failures surface 

129 as one terminal ``UpstreamChunk`` carrying a safe public 

130 ``{"code", "message"}`` payload. 

131 

132 Args: 

133 request: Fully-resolved upstream request. 

134 

135 Yields: 

136 One ``UpstreamChunk`` per SSE ``data:`` line, with the 

137 OpenAI ``[DONE]`` marker flagged terminal. 

138 """ 

139 try: 

140 response = await self._http.request( 

141 method=request.method, 

142 url=request.url, 

143 headers=dict(request.headers), 

144 json=request.payload, 

145 timeout=request.timeout_seconds, 

146 channel_name=request.channel_name, 

147 ) 

148 except asyncio.CancelledError: 

149 yield HTTPUpstreamAdapter._error_chunk( 

150 "UPSTREAM_CANCELLED", "upstream stream cancelled" 

151 ) 

152 return 

153 except TimeoutError: 

154 yield HTTPUpstreamAdapter._error_chunk( 

155 "UPSTREAM_TIMEOUT", "upstream stream timed out" 

156 ) 

157 return 

158 except InfrastructureError: 

159 yield HTTPUpstreamAdapter._error_chunk( 

160 "UPSTREAM_FAILED", "upstream transport failure" 

161 ) 

162 return 

163 if not 200 <= response.status < 300: 

164 message = HTTPUpstreamAdapter._safe_error_message(response) 

165 yield HTTPUpstreamAdapter._error_chunk( 

166 "UPSTREAM_ERROR", 

167 message or "upstream request failed", 

168 ) 

169 return 

170 for chunk in HTTPUpstreamAdapter._iter_sse_frames(response.body): 

171 yield chunk 

172 

173 async def cancel(self, request_id: str) -> None: 

174 """Record a streaming cancellation request (always succeeds). 

175 

176 The fake-safe transport cannot interrupt an in-flight response 

177 body, so cancellation is observed after the fact; the stream 

178 loop stops consulting this adapter once its cancel is recorded. 

179 

180 Args: 

181 request_id: Identifier of the stream being cancelled. 

182 """ 

183 self._cancelled.add(request_id) 

184 

185 @staticmethod 

186 def _error_chunk(code: str, message: str) -> UpstreamChunk: 

187 """Build one terminal error chunk from a safe public message. 

188 

189 Args: 

190 code: Stable machine-readable error code. 

191 message: Public, non-credential error message. 

192 

193 Returns: 

194 A terminal ``UpstreamChunk`` whose JSON payload carries 

195 ``code`` and ``message`` keys. 

196 """ 

197 return UpstreamChunk( 

198 event="error", 

199 data=( 

200 '{"code": ' 

201 + dumps_str(code) 

202 + ', "message": ' 

203 + dumps_str(message) 

204 + "}" 

205 ), 

206 terminal=True, 

207 ) 

208 

209 @staticmethod 

210 def _iter_sse_frames(body: bytes) -> Iterator[UpstreamChunk]: 

211 """Slice an SSE body into ``data:`` chunks, calling out ``[DONE]``. 

212 

213 Blocks separated by blank lines are scanned for ``data:`` lines; 

214 every block yields at most one chunk (multi-line data payloads 

215 are joined with ``\\n``). The OpenAI ``[DONE]`` marker is 

216 flagged terminal. 

217 

218 Args: 

219 body: The raw SSE response body. 

220 

221 Returns: 

222 An iterator of ``UpstreamChunk`` values covering every 

223 ``data:`` block in order. 

224 """ 

225 for block in body.replace(b"\r\n", b"\n").split(b"\n\n"): 

226 data_lines = [ 

227 line[5:].strip() 

228 for line in block.split(b"\n") 

229 if line.startswith(b"data:") 

230 ] 

231 if not data_lines: 

232 continue 

233 raw = b"\n".join(data_lines).decode("utf-8", errors="replace") 

234 yield UpstreamChunk(event=None, data=raw, terminal=raw == "[DONE]") 

235 

236 @staticmethod 

237 def _decode_success( 

238 request: UpstreamRequest, response: HttpResponse 

239 ) -> Result[UpstreamResponse, RelayGatewayError]: 

240 """Map a 2xx response to an ``UpstreamResponse``. 

241 

242 Empty bodies yield a ``None`` payload; non-empty bodies must be a 

243 JSON object or the response is classified as malformed. 

244 """ 

245 if not response.body: 

246 return Ok( 

247 UpstreamResponse( 

248 status_code=response.status, 

249 headers=dict(response.headers), 

250 payload=None, 

251 ) 

252 ) 

253 try: 

254 parsed = loads(response.body) 

255 except ValueError: 

256 return Err( 

257 RelayGatewayError( 

258 code="UPSTREAM_MALFORMED", 

259 message="malformed upstream response", 

260 status_code=502, 

261 request_id=request.request_id, 

262 ) 

263 ) 

264 if not isinstance(parsed, dict): 

265 return Err( 

266 RelayGatewayError( 

267 code="UPSTREAM_MALFORMED", 

268 message="malformed upstream response", 

269 status_code=502, 

270 request_id=request.request_id, 

271 ) 

272 ) 

273 return Ok( 

274 UpstreamResponse( 

275 status_code=response.status, 

276 headers=dict(response.headers), 

277 payload=parsed, 

278 ) 

279 ) 

280 

281 @staticmethod 

282 def _decode_error( 

283 request: UpstreamRequest, response: HttpResponse 

284 ) -> Result[UpstreamResponse, RelayGatewayError]: 

285 """Map a non-2xx response to an ``UPSTREAM_ERROR``. 

286 

287 The error message is a safe public string extracted from the 

288 body (``error.message``, ``error``, or ``message`` — first 

289 present wins); raw bodies, headers, and their keys are never 

290 propagated. 

291 """ 

292 message = HTTPUpstreamAdapter._safe_error_message(response) 

293 status = response.status if 400 <= response.status <= 599 else 502 

294 return Err( 

295 RelayGatewayError( 

296 code="UPSTREAM_ERROR", 

297 message=message or "upstream request failed", 

298 status_code=status, 

299 request_id=request.request_id, 

300 retryable=response.status >= 500, 

301 ) 

302 ) 

303 

304 @staticmethod 

305 def _safe_error_message(response: HttpResponse) -> str | None: 

306 """Extract a public error message from a non-2xx body. 

307 

308 Returns: 

309 The first string found under ``error.message``, ``error``, 

310 or ``message`` (in that order), or ``None`` when the body is 

311 empty, not a JSON object, or holds no public message key. 

312 """ 

313 if not response.body: 

314 return None 

315 try: 

316 parsed = loads(response.body) 

317 except ValueError: 

318 return None 

319 if not isinstance(parsed, dict): 

320 return None 

321 error = parsed.get("error") 

322 if isinstance(error, dict): 

323 nested = error.get("message") 

324 if isinstance(nested, str): 

325 return nested 

326 elif isinstance(error, str): 

327 return error 

328 message = parsed.get("message") 

329 if isinstance(message, str): 

330 return message 

331 return None