Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-relay/src/lexigram/ai/relay/mappers/gemini_response.py: 21%

119 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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)