Coverage for src / lexigram / ai / relay / gateway / codec.py: 93%

45 statements  

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

1"""Wire JSON decode/encode with protocol-specific DTO construction. 

2 

3The relay gateway speaks four wire formats (OpenAI Chat, OpenAI 

4Responses, Claude, Gemini). This module owns the boundary between raw 

5wire JSON bytes and the typed request DTOs from ``lexigram-contracts``: 

6decode validates the JSON object root and required fields, and encode 

7serializes a converted DTO back to bytes while preserving the DTOs' 

8None-omission semantics. 

9""" 

10 

11from __future__ import annotations 

12 

13from typing import Any, TypeAlias 

14 

15from lexigram.contracts.ai.exceptions import RelayError 

16from lexigram.contracts.ai.relay import ( 

17 ClaudeRequest, 

18 GeminiRequest, 

19 OpenAIChatRequest, 

20 RelayFormat, 

21 RelayGatewayError, 

22 RelayResponsePayload, 

23 ResponsesRequest, 

24) 

25from lexigram.contracts.ai.relay.dto import ( 

26 ClaudeResponse, 

27 GeminiResponse, 

28 OpenAIChatResponse, 

29 ResponsesResponse, 

30) 

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

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

33from lexigram.serialization import dumps, loads 

34 

35__all__ = ["RelayPayloadCodec"] 

36 

37WireRequest: TypeAlias = ( 

38 OpenAIChatRequest | ResponsesRequest | ClaudeRequest | GeminiRequest 

39) 

40"""Any typed request DTO accepted by the relay gateway.""" 

41 

42_REQUEST_TYPES: dict[RelayFormat, type[WireRequest]] = { 

43 RelayFormat.OPENAI_CHAT: OpenAIChatRequest, 

44 RelayFormat.OPENAI_RESPONSES: ResponsesRequest, 

45 RelayFormat.CLAUDE: ClaudeRequest, 

46 RelayFormat.GEMINI: GeminiRequest, 

47} 

48"""Map each wire format to its request DTO class.""" 

49 

50_RESPONSE_TYPES: dict[RelayFormat, type[RelayResponsePayload]] = { 

51 RelayFormat.OPENAI_CHAT: OpenAIChatResponse, 

52 RelayFormat.OPENAI_RESPONSES: ResponsesResponse, 

53 RelayFormat.CLAUDE: ClaudeResponse, 

54 RelayFormat.GEMINI: GeminiResponse, 

55} 

56"""Map each wire format to its response DTO class.""" 

57 

58 

59class RelayPayloadCodec: 

60 """Decode and encode relay wire payloads as typed DTOs. 

61 

62 The codec is stateless; a single instance can be shared. Decoding 

63 rejects malformed JSON, non-object roots, and DTOs missing required 

64 fields; unknown wire fields are preserved verbatim in the DTO 

65 ``passthrough`` dict and re-emitted on encode. 

66 """ 

67 

68 def decode_request( 

69 self, 

70 source: RelayFormat, 

71 raw: bytes, 

72 request_id: str, 

73 ) -> Result[WireRequest, RelayGatewayError]: 

74 """Decode wire JSON bytes into the request DTO for *source*. 

75 

76 Args: 

77 source: Wire format the payload claims to be. 

78 raw: Raw request body bytes. 

79 request_id: Caller-supplied request id stamped on errors. 

80 

81 Returns: 

82 ``Ok(dto)`` with unknown fields preserved in the DTO's 

83 ``passthrough``, or ``Err(RelayGatewayError)`` classifying 

84 malformed JSON (``INVALID_REQUEST``), non-object roots 

85 (``INVALID_REQUEST``), unknown formats 

86 (``UNSUPPORTED_FORMAT``), and missing required fields 

87 (``INVALID_REQUEST`` carrying the field path). 

88 """ 

89 try: 

90 decoded = loads(raw) 

91 except ValueError: 

92 return Err( 

93 RelayGatewayError( 

94 code="INVALID_REQUEST", 

95 message="malformed JSON", 

96 status_code=400, 

97 request_id=request_id, 

98 ) 

99 ) 

100 if not isinstance(decoded, dict): 

101 return Err( 

102 RelayGatewayError( 

103 code="INVALID_REQUEST", 

104 message="payload must be a JSON object", 

105 status_code=400, 

106 request_id=request_id, 

107 ) 

108 ) 

109 dto_type = _REQUEST_TYPES.get(source) 

110 if dto_type is None: 

111 return Err( 

112 RelayGatewayError( 

113 code="UNSUPPORTED_FORMAT", 

114 message=f"unsupported relay format: {source}", 

115 status_code=400, 

116 request_id=request_id, 

117 ) 

118 ) 

119 try: 

120 dto = dto_type.from_dict(decoded) 

121 except RelayError as relay_error: 

122 return Err( 

123 RelayGatewayError( 

124 code="INVALID_REQUEST", 

125 message=str(relay_error), 

126 status_code=400, 

127 request_id=request_id, 

128 ) 

129 ) 

130 return Ok(dto) 

131 

132 def decode_response_payload( 

133 self, 

134 target: RelayFormat, 

135 data: dict[str, Any], 

136 request_id: str, 

137 ) -> Result[RelayResponsePayload, RelayGatewayError]: 

138 """Decode an upstream wire dict into the response DTO for *target*. 

139 

140 Args: 

141 target: Wire format the upstream claims to speak. 

142 data: Decoded upstream response body. 

143 request_id: Caller-supplied request id stamped on errors. 

144 

145 Returns: 

146 ``Ok(dto)`` with unknown fields preserved in the DTO's 

147 ``passthrough``, or ``Err(RelayGatewayError)`` classifying 

148 unknown formats (``UNSUPPORTED_FORMAT``) and DTOs missing 

149 required fields (``UPSTREAM_MALFORMED`` — a malformed 

150 upstream response is a 502, not a client 400). 

151 """ 

152 dto_type = _RESPONSE_TYPES.get(target) 

153 if dto_type is None: 

154 return Err( 

155 RelayGatewayError( 

156 code=RelayGatewayErrorCode.UNSUPPORTED_FORMAT, 

157 message=f"unsupported relay format: {target}", 

158 status_code=400, 

159 request_id=request_id, 

160 ) 

161 ) 

162 try: 

163 dto = dto_type.from_dict(data) 

164 except RelayError as relay_error: 

165 return Err( 

166 RelayGatewayError( 

167 code=RelayGatewayErrorCode.UPSTREAM_MALFORMED, 

168 message=relay_error.message, 

169 status_code=502, 

170 request_id=request_id, 

171 retryable=False, 

172 ) 

173 ) 

174 return Ok(dto) 

175 

176 def encode(self, dto: WireRequest) -> Result[bytes, RelayGatewayError]: 

177 """Serialize a request DTO to wire JSON bytes. 

178 

179 Args: 

180 dto: Typed request DTO to serialize. 

181 

182 Returns: 

183 ``Ok(bytes)`` with ``None`` fields omitted and falsey values 

184 preserved, or ``Err(RelayGatewayError)`` with code 

185 ``ENCODE_FAILED`` when the payload cannot be serialized. 

186 """ 

187 try: 

188 return Ok(dumps(dto.to_dict())) 

189 except (TypeError, ValueError): 

190 return Err( 

191 RelayGatewayError( 

192 code="ENCODE_FAILED", 

193 message="failed to serialize payload", 

194 status_code=500, 

195 request_id="", 

196 ) 

197 )