Coverage for src / lexigram / ai / relay / gateway / web / image_endpoints.py: 97%

91 statements  

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

1"""Image passthrough routes for the relay gateway web layer. 

2 

3``POST /v1/images/generations`` and ``POST /v1/images/edits`` serve 

4image endpoint kinds through the same ``PassthroughService`` as the 

5embeddings route: same request resolution, header filtering, and OpenAI 

6error envelope machinery, with no wire-format conversion. Generations 

7bodies are JSON and forward encoded unchanged; edits bodies are 

8``multipart/form-data`` and forward byte-for-byte with the ``model`` 

9form field lifted into the gateway request so channel selection and 

10model-suffix substitution work without ever parsing the body as JSON. 

11Responses pass through verbatim: decoded JSON when the upstream returned 

12JSON, raw bytes otherwise, ``204`` when the result carries neither. 

13""" 

14 

15from __future__ import annotations 

16 

17from collections.abc import Mapping 

18from functools import partial 

19 

20from starlette.requests import Request 

21from starlette.responses import JSONResponse, Response 

22from starlette.routing import Route 

23 

24from lexigram.ai.relay.gateway.passthrough import RelayPassthroughBody 

25from lexigram.ai.relay.gateway.web.shared import ( 

26 ResolvePassthrough, 

27 _error_response, 

28 _parse_body, 

29 _safe_headers, 

30) 

31from lexigram.contracts.ai.relay import ( 

32 JsonValue, 

33 RelayFormat, 

34 RelayGatewayError, 

35 RelayGatewayRequest, 

36) 

37from lexigram.contracts.ai.relay.gateway import RelayGatewayErrorCode 

38from lexigram.identity.ambient import new_uuid 

39 

40__all__ = [ 

41 "IMAGE_ROUTE_PATHS", 

42 "IMAGE_ROUTE_TABLE", 

43 "build_image_routes", 

44 "image_endpoint", 

45] 

46 

47IMAGE_ROUTE_TABLE: tuple[tuple[str, str], ...] = ( 

48 ("/v1/images/generations", "image_generation"), 

49 ("/v1/images/edits", "image_edit"), 

50) 

51"""Inbound path to endpoint kind for image passthrough routes.""" 

52 

53IMAGE_ROUTE_PATHS: tuple[str, ...] = tuple(path for path, _ in IMAGE_ROUTE_TABLE) 

54"""Inbound image paths registered by ``build_image_routes``.""" 

55 

56_FORM_FIELD_HEADER_MARKER = b'name="' 

57_FORM_FIELD_HEADER_SUFFIX = b'"' 

58"""Multipart ``Content-Disposition`` attribute delimiters used by the model read.""" 

59 

60_MULTIPART_MEDIA_TYPE = "multipart/form-data" 

61 

62 

63async def image_endpoint( 

64 kind: str, 

65 resolve_passthrough: ResolvePassthrough, 

66 request: Request, 

67) -> Response: 

68 """Serve one inbound image request in its own wire format. 

69 

70 The body is read exactly once. Generations bodies are JSON and 

71 parse exactly once, like ``passthrough_endpoint``; edits bodies are 

72 ``multipart/form-data`` and travel verbatim with the ``model`` form 

73 field lifted into the gateway request, so a raw multipart body can 

74 never be mistaken for JSON. The service is resolved per request, 

75 never cached. Image endpoints are OpenAI-shaped by convention, so 

76 failures render in the OpenAI error envelope. 

77 

78 Args: 

79 kind: The image endpoint kind owned by this route. 

80 resolve_passthrough: Resolver of the passthrough service. 

81 request: The Starlette request being served. 

82 

83 Returns: 

84 The upstream payload JSON verbatim, the upstream binary body 

85 verbatim, ``204`` when the result carries neither, or the OpenAI 

86 error envelope for gateway failures. 

87 """ 

88 raw = await request.body() 

89 request_id = getattr(request.state, "request_id", None) or new_uuid() 

90 trace_id = request.headers.get("x-trace-id", "") or "" 

91 content_type = request.headers.get("content-type", "application/json") 

92 if _is_multipart(content_type): 

93 boundary = _multipart_boundary(content_type) 

94 model_value = _multipart_model_value(raw, boundary) if boundary else None 

95 if not model_value: 

96 return _model_required_error(request_id) 

97 payload: Mapping[str, JsonValue] | RelayPassthroughBody = ( 

98 RelayPassthroughBody.raw(raw, content_type) 

99 ) 

100 else: 

101 body = _parse_body(raw, RelayFormat.OPENAI_CHAT, request_id) 

102 if isinstance(body, Response): 

103 return body 

104 model_value = body.get("model") 

105 if not isinstance(model_value, str) or not model_value: 

106 return _model_required_error(request_id) 

107 payload = body 

108 user = getattr(request.state, "user", None) 

109 tenant_id = "" 

110 if isinstance(user, dict): 

111 tenant = user.get("tenant_id") or user.get("tenant") 

112 if isinstance(tenant, str): 

113 tenant_id = tenant 

114 gateway_request = RelayGatewayRequest( 

115 request_id=request_id, 

116 tenant_id=tenant_id, 

117 source=RelayFormat.OPENAI_CHAT, 

118 model=model_value, 

119 stream=False, 

120 payload=payload, 

121 headers=dict(request.headers.items()), 

122 channel=None, 

123 ) 

124 service = await resolve_passthrough(request) 

125 result = await service.handle(kind, gateway_request) 

126 if result.is_err(): 

127 return _error_response(RelayFormat.OPENAI_CHAT, result.unwrap_err()) 

128 ok_result = result.unwrap() 

129 headers = _safe_headers(ok_result.headers, request_id, trace_id) 

130 if ok_result.payload is not None: 

131 return JSONResponse( 

132 content=ok_result.payload, 

133 status_code=ok_result.status_code, 

134 headers=headers, 

135 ) 

136 if ok_result.body: 

137 headers = {**headers, "content-type": ok_result.content_type} 

138 return Response( 

139 content=ok_result.body, 

140 status_code=ok_result.status_code, 

141 headers=headers, 

142 ) 

143 return Response(status_code=204, headers=headers) 

144 

145 

146def build_image_routes( 

147 resolve_passthrough: ResolvePassthrough, 

148) -> list[Route]: 

149 """Build the image passthrough POST routes bound to a service resolver. 

150 

151 Args: 

152 resolve_passthrough: Async callable resolving a 

153 ``PassthroughService`` from the request; wired to 

154 request-time DI by the contributor. 

155 

156 Returns: 

157 One ``Route`` per image path, in ``IMAGE_ROUTE_PATHS`` order. 

158 """ 

159 return [ 

160 Route( 

161 path, 

162 partial(image_endpoint, kind, resolve_passthrough), 

163 methods=["POST"], 

164 ) 

165 for path, kind in IMAGE_ROUTE_TABLE 

166 ] 

167 

168 

169def _is_multipart(content_type: str) -> bool: 

170 """Tell whether a content-type header denotes multipart form data. 

171 

172 Args: 

173 content_type: A content-type header value. 

174 

175 Returns: 

176 ``True`` for a ``multipart/form-data`` media type with or 

177 without parameters; ``False`` otherwise. 

178 """ 

179 media_type = content_type.partition(";")[0].strip().lower() 

180 return media_type == _MULTIPART_MEDIA_TYPE 

181 

182 

183def _multipart_boundary(content_type: str) -> str | None: 

184 """Extract the ``boundary`` parameter from a content-type header. 

185 

186 Args: 

187 content_type: The raw content-type header value. 

188 

189 Returns: 

190 The boundary token without surrounding quotes, or ``None`` when 

191 the header carries no boundary parameter. 

192 """ 

193 for parameter in content_type.split(";"): 

194 key, separator, raw_value = parameter.strip().partition("=") 

195 if separator and key.lower() == "boundary": 

196 return raw_value.strip().strip('"') 

197 return None 

198 

199 

200def _multipart_model_value(body: bytes, boundary: str) -> str | None: 

201 """Read the ``model`` form field's value from a multipart body. 

202 

203 A narrow boundary-aware read (not a general multipart parser): the 

204 body is split on the ``--<boundary>`` framing marker and the first 

205 part whose ``Content-Disposition`` header carries ``name="model"`` 

206 has its value character returned; every check mirrors 

207 ``rewrite_multipart_form_field`` so both functions agree on the 

208 same parts. 

209 

210 Args: 

211 body: The raw ``multipart/form-data`` body bytes. 

212 boundary: The boundary token from the content-type header. 

213 

214 Returns: 

215 The ``model`` field value, or ``None`` when the body has no 

216 model part (or no boundary marker at all). 

217 """ 

218 marker = b"--" + boundary.encode("utf-8") 

219 target = _FORM_FIELD_HEADER_MARKER + b"model" + _FORM_FIELD_HEADER_SUFFIX 

220 segments = body.split(marker) 

221 if len(segments) < 2: 

222 return None 

223 for index in range(1, len(segments) - 1): 

224 part = segments[index] 

225 separator = part.find(b"\r\n\r\n") 

226 if separator < 0: 

227 continue 

228 headers = part[2:separator].lower() 

229 if target not in headers: 

230 continue 

231 value_start = separator + 4 

232 value_end = len(part) - 2 if part.endswith(b"\r\n") else len(part) 

233 value = part[value_start:value_end] 

234 if not value: 

235 return None 

236 return value.decode("utf-8", errors="replace") 

237 return None 

238 

239 

240def _model_required_error(request_id: str) -> Response: 

241 """Build the 400 OpenAI envelope for a missing image model field. 

242 

243 Args: 

244 request_id: Request id stamped on the error. 

245 

246 Returns: 

247 The ``INVALID_REQUEST`` envelope with status 400. 

248 """ 

249 return _error_response( 

250 RelayFormat.OPENAI_CHAT, 

251 RelayGatewayError( 

252 code=RelayGatewayErrorCode.INVALID_REQUEST, 

253 message="model is required", 

254 status_code=400, 

255 request_id=request_id, 

256 ), 

257 )