1"""Google Gemini ``generateContent`` response conversion.
2
3Parses :class:`GeminiResponse` wire DTOs into canonical
4:class:`RelayResponse` (:func:`response_to_ir`) and rebuilds them from
5canonical IR (:func:`ir_to_response`). Stream conversion is handled by
6the shared stream lifecycle task and reports ``unsupported_feature``
7until then.
8"""
9
10from __future__ import annotations
11
12from typing import Any
13
14from lexigram.ai.relay.context import ConversionContext
15from lexigram.ai.relay.errors import translate, unsupported_format
16from lexigram.ai.relay.finish_reasons import (
17 FINISH_REASON_TO_WIRE,
18 finish_reason_to_wire,
19)
20from lexigram.ai.relay.mappers.base import record_loss
21from lexigram.ai.relay.mappers.gemini_request import (
22 _TARGET,
23 _tool_call_from_part,
24 _tool_call_to_part,
25)
26from lexigram.contracts.ai.exceptions import RelayError
27from lexigram.contracts.ai.llm import ToolCall
28from lexigram.contracts.ai.relay.dto import (
29 GeminiCandidate,
30 GeminiContent,
31 GeminiGroundingMetadata,
32 GeminiPart,
33 GeminiPromptFeedback,
34 GeminiResponse,
35 GeminiSafetyRating,
36 GeminiUsageMetadata,
37)
38from lexigram.contracts.ai.relay.ir import RelayResponse, normalize_finish_reason
39from lexigram.contracts.ai.relay.types import RelayUsage
40from lexigram.contracts.ai.thinking import ThinkingResult
41from lexigram.contracts.core.result import Err, Ok, Result
42
43__all__ = ["ir_to_response", "response_to_ir"]
44
45
46def response_to_ir(
47 payload: Any, *, context: ConversionContext
48) -> Result[RelayResponse, RelayError]:
49 """Convert a ``GeminiResponse`` into canonical ``RelayResponse``.
50
51 Args:
52 payload: A wire response DTO.
53 context: Per-conversion context with loss sink.
54
55 Returns:
56 Ok(response) on success, Err(relay_error) on malformed payload.
57 """
58 if not isinstance(payload, GeminiResponse):
59 return Err(
60 unsupported_format(f"expected GeminiResponse, got {type(payload).__name__}")
61 )
62 try:
63 candidates = payload.candidates or []
64 if len(candidates) > 1:
65 record_loss(
66 context,
67 field="candidates",
68 target=_TARGET,
69 reason="multiple_candidates_collapsed",
70 )
71 candidate = candidates[0] if candidates else None
72 passthrough: dict[str, Any] = dict(payload.passthrough)
73 if payload.model_version is not None:
74 passthrough["model_version"] = payload.model_version
75 if payload.create_time is not None:
76 passthrough["create_time"] = payload.create_time
77 if payload.prompt_feedback is not None:
78 passthrough["prompt_feedback"] = payload.prompt_feedback.to_dict()
79 content = ""
80 thinking: ThinkingResult | None = None
81 tool_calls: list[ToolCall] = []
82 if candidate is not None and candidate.content is not None:
83 text_parts: list[str] = []
84 think_parts: list[str] = []
85 think_signature: str | None = None
86 for part in candidate.content.parts:
87 if part.thought:
88 think_parts.append(part.text or "")
89 think_signature = part.thought_signature
90 elif part.text is not None:
91 text_parts.append(part.text)
92 elif part.function_call is not None:
93 tool_calls.append(_tool_call_from_part(part))
94 elif part.inline_data is not None or part.function_response is not None:
95 record_loss(
96 context,
97 field="content.part",
98 target=_TARGET,
99 reason="unrepresentable_part_dropped",
100 )
101 content = "".join(text_parts)
102 if think_parts:
103 thinking = ThinkingResult(
104 content="".join(think_parts),
105 signature=think_signature,
106 tokens=_thought_tokens(payload),
107 )
108 _preserve_candidate_metadata(candidate, passthrough)
109 return Ok(
110 RelayResponse(
111 model=payload.model_version or "",
112 id=payload.response_id,
113 content=content,
114 thinking=thinking,
115 tool_calls=tool_calls,
116 finish_reason=normalize_finish_reason(
117 candidate.finish_reason if candidate else None
118 ),
119 usage=_usage_from_wire(payload.usage_metadata),
120 passthrough=passthrough,
121 )
122 )
123 except (RelayError, ValueError, TypeError, KeyError) as exc:
124 return Err(translate(exc, detail="response_to_ir"))
125
126
127def ir_to_response(
128 response: RelayResponse, *, context: ConversionContext
129) -> Result[Any, RelayError]:
130 """Convert canonical ``RelayResponse`` into a ``GeminiResponse``.
131
132 Args:
133 response: Canonical response IR.
134 context: Per-conversion context with loss sink.
135
136 Returns:
137 Ok(response) on success, Err(relay_error) on failure.
138 """
139 try:
140 passthrough = dict(response.passthrough)
141 model_version = passthrough.pop("model_version", None)
142 prompt_feedback = passthrough.pop("prompt_feedback", None)
143 create_time = passthrough.pop("create_time", None)
144 safety_ratings = passthrough.pop("safety_ratings", None)
145 grounding_metadata = passthrough.pop("grounding_metadata", None)
146 citation_metadata = passthrough.pop("citation_metadata", None)
147 token_count = passthrough.pop("token_count", None)
148 avg_logprobs = passthrough.pop("avg_logprobs", None)
149 parts: list[GeminiPart] = []
150 if response.content:
151 parts.append(GeminiPart(text=response.content))
152 for tool_call in response.tool_calls:
153 parts.append(_tool_call_to_part(tool_call))
154 candidate = GeminiCandidate(
155 content=GeminiContent(role="model", parts=parts),
156 finish_reason=_finish_reason_from_ir(response.finish_reason, context),
157 index=0,
158 safety_ratings=_safety_ratings_from_passthrough(safety_ratings) or [],
159 grounding_metadata=_grounding_from_passthrough(grounding_metadata),
160 citation_metadata=(
161 citation_metadata if isinstance(citation_metadata, dict) else None
162 ),
163 token_count=token_count if isinstance(token_count, int) else None,
164 avg_logprobs=(
165 avg_logprobs if isinstance(avg_logprobs, (int, float)) else None
166 ),
167 passthrough=dict(passthrough),
168 )
169 return Ok(
170 GeminiResponse(
171 candidates=[candidate],
172 prompt_feedback=_prompt_feedback_from_passthrough(prompt_feedback),
173 usage_metadata=_usage_to_wire(response.usage),
174 model_version=model_version if isinstance(model_version, str) else None,
175 create_time=create_time if isinstance(create_time, str) else None,
176 passthrough=passthrough,
177 )
178 )
179 except (RelayError, ValueError, TypeError, KeyError) as exc:
180 return Err(translate(exc, detail="ir_to_response"))
181
182
183def _usage_from_wire(usage: GeminiUsageMetadata | None) -> RelayUsage | None:
184 """Map a wire ``GeminiUsageMetadata`` into canonical ``RelayUsage``.
185
186 Mirrors relaykit's ``UsageFromGeminiMetadata``: completion counts
187 thinking tokens, the prompt adds tool-use tokens, and the explicit
188 total is preserved because Gemini counts thoughts within both.
189 """
190 if usage is None:
191 return None
192 return RelayUsage(
193 prompt_tokens=(usage.prompt_token_count + usage.tool_use_prompt_token_count),
194 completion_tokens=(
195 usage.candidates_token_count + (usage.thoughts_token_count or 0)
196 ),
197 cache_read_tokens=usage.cached_content_token_count or 0,
198 reasoning_tokens=usage.thoughts_token_count or 0,
199 total_tokens_override=usage.total_token_count or None,
200 )
201
202
203def _usage_to_wire(usage: RelayUsage | None) -> GeminiUsageMetadata | None:
204 """Serialize canonical ``RelayUsage`` into a ``GeminiUsageMetadata``.
205
206 Gemini reports thinking tokens as a subset of the candidate
207 tokens and does not surface cache or reasoning fields in the
208 generated payload, so those counters are emitted as zeros.
209 """
210 if usage is None:
211 return None
212 return GeminiUsageMetadata(
213 prompt_token_count=usage.prompt_tokens,
214 candidates_token_count=usage.completion_tokens,
215 total_token_count=usage.total_tokens,
216 cached_content_token_count=0,
217 thoughts_token_count=0,
218 tool_use_prompt_token_count=0,
219 )
220
221
222def _thought_tokens(payload: GeminiResponse) -> int | None:
223 """Read thinking tokens from the usage metadata."""
224 if (
225 payload.usage_metadata is None
226 or not payload.usage_metadata.thoughts_token_count
227 ):
228 return None
229 return payload.usage_metadata.thoughts_token_count
230
231
232def _preserve_candidate_metadata(
233 candidate: GeminiCandidate, passthrough: dict[str, Any]
234) -> None:
235 """Preserve candidate-level provider metadata as passthrough."""
236 if candidate.safety_ratings:
237 passthrough["safety_ratings"] = [
238 rating.to_dict() for rating in candidate.safety_ratings
239 ]
240 if candidate.grounding_metadata is not None:
241 passthrough["grounding_metadata"] = candidate.grounding_metadata.to_dict()
242 if candidate.citation_metadata is not None:
243 passthrough["citation_metadata"] = candidate.citation_metadata
244 if candidate.token_count is not None:
245 passthrough["token_count"] = candidate.token_count
246 if candidate.avg_logprobs is not None:
247 passthrough["avg_logprobs"] = candidate.avg_logprobs
248 passthrough.update(candidate.passthrough)
249
250
251def _finish_reason_from_ir(
252 finish_reason: str | None, context: ConversionContext
253) -> str | None:
254 """Map a canonical finish reason back to a Gemini value."""
255 if finish_reason is None:
256 return None
257 if finish_reason == "function_call":
258 record_loss(
259 context,
260 field="finish_reason",
261 target=_TARGET,
262 reason="function_call_adapted",
263 )
264 elif finish_reason not in FINISH_REASON_TO_WIRE:
265 record_loss(
266 context,
267 field="finish_reason",
268 target=_TARGET,
269 reason="finish_reason_adapted",
270 )
271 return finish_reason_to_wire(finish_reason, _TARGET)
272
273
274def _safety_ratings_from_passthrough(
275 raw: Any,
276) -> list[GeminiSafetyRating] | None:
277 """Rebuild safety ratings from passthrough dicts."""
278 if not isinstance(raw, list):
279 return None
280 ratings = [
281 GeminiSafetyRating.from_dict(item) for item in raw if isinstance(item, dict)
282 ]
283 return ratings or None
284
285
286def _grounding_from_passthrough(raw: Any) -> GeminiGroundingMetadata | None:
287 """Rebuild grounding metadata from a passthrough dict."""
288 if not isinstance(raw, dict):
289 return None
290 return GeminiGroundingMetadata.from_dict(raw)
291
292
293def _prompt_feedback_from_passthrough(
294 raw: Any,
295) -> GeminiPromptFeedback | None:
296 """Rebuild prompt feedback from a passthrough dict."""
297 if not isinstance(raw, dict):
298 return None
299 return GeminiPromptFeedback.from_dict(raw)