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 )