Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/gemini_helpers.py: 13%

1"""Shared helpers for Google Gemini and Vertex AI clients. 

2 

3Both :class:`~lexigram.ai.llm.clients.gemini.GeminiClient` and 

4:class:`~lexigram.ai.llm.clients.vertex_ai.VertexAIClient` target the same 

5Gemini model contract. This module holds all shared message conversion, 

6response parsing, thinking injection, and tool-formatting utilities so that 

7neither client imports private symbols from the other. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import TYPE_CHECKING, Any 

13 

14from lexigram.ai.llm.clients._message_utils import serialize_content_for_gemini 

15from lexigram.ai.llm.clients._tools_utils import parse_json_arguments 

16from lexigram.ai.llm.types import ( 

17 AIError, 

18 Completion, 

19 FunctionCall, 

20 StreamChunk, 

21 ThinkingResult, 

22 TokenUsage, 

23 ToolCall, 

24) 

25from lexigram.logging import ( 

26 get_logger, 

27) 

28from lexigram.serialization import dumps_str 

29from lexigram.serialization import loads as _loads 

30 

31if TYPE_CHECKING: 

32 from collections.abc import Iterator 

33 

34 from lexigram.ai.llm.config import ClientConfig 

35 

36logger = get_logger(__name__) 

37 

38 

39__all__ = [ 

40 "inject_thinking_config", 

41 "messages_to_gemini", 

42 "parse_gemini_response", 

43 "parse_gemini_response_with_tools", 

44 "parse_gemini_sse_body", 

45 "tool_to_gemini_function", 

46] 

47 

48 

49def inject_thinking_config(gen_config: dict[str, Any], config: ClientConfig) -> None: 

50 """Inject ``thinkingConfig`` into a Gemini ``generationConfig`` dict. 

51 

52 When ``config.thinking.suppress`` is set, injects ``thinkingBudget: 0`` to 

53 disable thinking. Gemini 3 models use ``thinkingLevel``; Gemini 2.5 models 

54 use ``thinkingBudget``. 

55 

56 Args: 

57 gen_config: The ``generationConfig`` sub-dict to mutate in place. 

58 config: LLM configuration. 

59 """ 

60 if config.thinking is None: 

61 return 

62 if config.thinking.suppress: 

63 gen_config["thinkingConfig"] = {"thinkingBudget": 0} 

64 return 

65 if config.thinking.level: 

66 gen_config["thinkingConfig"] = {"thinkingLevel": config.thinking.level} 

67 else: 

68 gen_config["thinkingConfig"] = {"thinkingBudget": config.thinking.budget_tokens} 

69 

70 

71def messages_to_gemini( 

72 messages: list[dict[str, Any]], 

73) -> list[dict[str, Any]]: 

74 """Convert OpenAI-format messages to Gemini ``contents`` format. 

75 

76 System messages are prepended as a text part to the first user turn 

77 (Gemini has no top-level system role). Image URL content parts are 

78 converted to ``inline_data`` for data URIs, or kept as text references 

79 for external URLs. 

80 

81 Assistant turns that requested tools emit ``functionCall`` parts; tool 

82 results (``Role.TOOL``) emit ``functionResponse`` parts so multi-turn 

83 tool conversations round-trip correctly. 

84 

85 Args: 

86 messages: OpenAI-compatible message list or ChatMessage objects. 

87 

88 Returns: 

89 Gemini-format ``contents`` list. 

90 """ 

91 system_text: str | None = None 

92 contents: list[dict[str, Any]] = [] 

93 

94 for msg in messages: 

95 role: str = _role_str(msg) 

96 content: Any = ( 

97 msg.get("content", "") 

98 if isinstance(msg, dict) 

99 else getattr(msg, "content", "") 

100 ) 

101 tool_calls: Any = ( 

102 msg.get("tool_calls") 

103 if isinstance(msg, dict) 

104 else getattr(msg, "tool_calls", None) 

105 ) 

106 

107 if role == "system": 

108 system_text = ( 

109 content if isinstance(content, str) else _extract_text(content) 

110 ) 

111 continue 

112 

113 gemini_role = "user" if role in ("user", "tool", "function") else "model" 

114 parts: list[dict[str, Any]] = [] 

115 

116 # Prepend system text to the first user turn 

117 if system_text and gemini_role == "user" and not contents: 

118 parts.append({"text": system_text}) 

119 system_text = None 

120 

121 # Use the multimodal serializer for MessageContent 

122 serialized = serialize_content_for_gemini(content) 

123 if serialized and role != "tool": 

124 parts.extend(serialized) 

125 

126 if role == "tool": 

127 tool_call_id: str = ( 

128 msg.get("tool_call_id", "") 

129 if isinstance(msg, dict) 

130 else getattr(msg, "tool_call_id", "") 

131 ) 

132 if tool_call_id: 

133 parts.append( 

134 { 

135 "functionResponse": { 

136 "name": tool_call_id, 

137 "response": {"content": _extract_text(content)}, 

138 } 

139 } 

140 ) 

141 elif role == "assistant" and tool_calls: 

142 for call in tool_calls: 

143 fn = getattr(call, "function", None) 

144 if fn is None or not getattr(fn, "name", None): 

145 continue 

146 parts.append( 

147 { 

148 "functionCall": { 

149 "name": fn.name, 

150 "args": parse_json_arguments(fn.arguments), 

151 } 

152 } 

153 ) 

154 

155 # Drop empty text parts when the turn carries tool parts; Gemini 

156 # rejects empty text blocks alongside functionCall/functionResponse. 

157 if any("functionCall" in p or "functionResponse" in p for p in parts): 

158 parts = [p for p in parts if p != {"text": ""}] 

159 

160 contents.append({"role": gemini_role, "parts": parts}) 

161 

162 if system_text: 

163 logger.warning( 

164 "gemini_system_message_dropped", 

165 reason="first_non_system_message_was_not_user_role", 

166 ) 

167 

168 return contents 

169 

170 

171def parse_gemini_response(data: dict[str, Any], model: str) -> Completion: 

172 """Parse a Gemini ``generateContent`` response into a ``Completion``. 

173 

174 Separates ``thought`` parts (thinking) from answer parts. Populates 

175 :attr:`~lexigram.ai.llm.types.Completion.thinking` with a 

176 :class:`~lexigram.contracts.ai.thinking.ThinkingResult` when the model 

177 produced thinking output. 

178 

179 Args: 

180 data: Parsed JSON response dict from the Gemini API. 

181 model: Model identifier used for the ``Completion`` metadata. 

182 

183 Returns: 

184 Normalised :class:`~lexigram.ai.llm.types.Completion`. 

185 

186 Raises: 

187 AIError: When the response has no candidates (e.g. blocked by safety filters). 

188 """ 

189 candidates = data.get("candidates") 

190 if not candidates: 

191 block_reason = data.get("promptFeedback", {}).get("blockReason", "unknown") 

192 raise AIError(f"Gemini returned no candidates (blockReason={block_reason!r})") 

193 

194 candidate = candidates[0] 

195 parts = candidate.get("content", {}).get("parts", []) 

196 

197 thinking_parts: list[str] = [] 

198 answer_parts: list[str] = [] 

199 for p in parts: 

200 if not isinstance(p, dict): 

201 continue 

202 text = p.get("text", "") 

203 if p.get("thought"): 

204 thinking_parts.append(text) 

205 else: 

206 answer_parts.append(text) 

207 

208 usage_meta = data.get("usageMetadata", {}) 

209 reasoning_tokens: int | None = usage_meta.get("thoughtsTokenCount") or None 

210 usage = TokenUsage( 

211 prompt_tokens=usage_meta.get("promptTokenCount", 0), 

212 completion_tokens=usage_meta.get("candidatesTokenCount", 0), 

213 total_tokens=usage_meta.get("totalTokenCount", 0), 

214 ) 

215 

216 thinking: ThinkingResult | None = None 

217 if thinking_parts: 

218 thinking = ThinkingResult( 

219 content="".join(thinking_parts), 

220 tokens=reasoning_tokens, 

221 ) 

222 

223 return Completion( 

224 content="".join(answer_parts), 

225 model=model, 

226 thinking=thinking, 

227 usage=usage, 

228 ) 

229 

230 

231def parse_gemini_response_with_tools(data: dict[str, Any], model: str) -> Completion: 

232 """Parse a Gemini response that may contain function-call parts. 

233 

234 Handles both plain-text and ``functionCall`` parts in the response 

235 candidates, returning a :class:`Completion` with optional 

236 :attr:`~Completion.tool_calls` populated. 

237 

238 Args: 

239 data: Parsed JSON from the Gemini ``generateContent`` endpoint. 

240 model: Model identifier used for the ``Completion`` metadata. 

241 

242 Returns: 

243 Normalised :class:`~lexigram.ai.llm.types.Completion`. 

244 

245 Raises: 

246 AIError: When the response has no candidates. 

247 """ 

248 candidates = data.get("candidates") 

249 if not candidates: 

250 block_reason = data.get("promptFeedback", {}).get("blockReason", "unknown") 

251 raise AIError(f"Gemini returned no candidates (blockReason={block_reason!r})") 

252 

253 candidate = candidates[0] 

254 parts = candidate.get("content", {}).get("parts", []) 

255 

256 text_parts: list[str] = [] 

257 tool_calls: list[ToolCall] = [] 

258 

259 for part in parts: 

260 if not isinstance(part, dict): 

261 continue 

262 if "text" in part: 

263 text_parts.append(part["text"]) 

264 elif "functionCall" in part: 

265 fc = part["functionCall"] 

266 name: str = fc.get("name", "") 

267 args: dict[str, Any] = fc.get("args", {}) 

268 tool_calls.append( 

269 ToolCall( 

270 id=name, 

271 type="function", 

272 function=FunctionCall( 

273 name=name, 

274 arguments=dumps_str(args), 

275 ), 

276 ) 

277 ) 

278 

279 usage_meta = data.get("usageMetadata", {}) 

280 usage = TokenUsage( 

281 prompt_tokens=usage_meta.get("promptTokenCount", 0), 

282 completion_tokens=usage_meta.get("candidatesTokenCount", 0), 

283 total_tokens=usage_meta.get("totalTokenCount", 0), 

284 ) 

285 

286 return Completion( 

287 content="".join(text_parts), 

288 model=model, 

289 usage=usage, 

290 tool_calls=tool_calls or None, 

291 ) 

292 

293 

294def parse_gemini_sse_body(body: str, model: str) -> Iterator[StreamChunk]: 

295 """Yield :class:`StreamChunk` objects from a Gemini SSE response body. 

296 

297 Gemini streams a JSON array over SSE. Each ``data:`` line carries a 

298 ``generateContentResponse`` object whose ``candidates[0].content.parts`` 

299 contain incremental text deltas. Parts with ``thought=True`` are yielded 

300 as thinking chunks. 

301 

302 Args: 

303 body: Full SSE response body text. 

304 model: Model label embedded in each :class:`StreamChunk`. 

305 

306 Yields: 

307 :class:`StreamChunk` with incremental text or thinking deltas. 

308 """ 

309 for line in body.splitlines(): 

310 stripped = line.strip() 

311 if not stripped.startswith("data:"): 

312 continue 

313 raw = stripped[len("data:") :].strip() 

314 if not raw or raw == "[DONE]": 

315 continue 

316 try: 

317 chunk_data: dict[str, Any] = _loads(raw) 

318 except (ValueError, TypeError): 

319 continue 

320 candidates = chunk_data.get("candidates", []) 

321 for i, cand in enumerate(candidates): 

322 parts = cand.get("content", {}).get("parts", []) 

323 finish = cand.get("finishReason") 

324 for part in parts: 

325 if not isinstance(part, dict) or "text" not in part: 

326 continue 

327 if part.get("thought"): 

328 yield StreamChunk( 

329 thinking_delta=part["text"], 

330 is_thinking=True, 

331 model=model, 

332 finish_reason=finish, 

333 index=i, 

334 ) 

335 else: 

336 yield StreamChunk( 

337 delta=part["text"], 

338 model=model, 

339 finish_reason=finish, 

340 index=i, 

341 ) 

342 

343 

344def tool_to_gemini_function(tool: Any) -> dict[str, Any]: 

345 """Convert a tool descriptor to a Gemini ``FunctionDeclaration``. 

346 

347 Supports objects that expose a ``__tool_schema__`` class attribute 

348 (following the Lexigram tool registration convention) as well as plain 

349 dictionaries in OpenAI tool format. 

350 

351 Args: 

352 tool: A :class:`ToolCall`, class with ``__tool_schema__``, or dict 

353 with ``function`` key in OpenAI tool format. 

354 

355 Returns: 

356 Gemini ``FunctionDeclaration`` dict with ``name``, ``description``, 

357 and ``parameters`` keys. 

358 """ 

359 if hasattr(tool, "__tool_schema__"): 

360 schema: dict[str, Any] = tool.__tool_schema__ 

361 return { 

362 "name": schema["name"], 

363 "description": schema.get("description", ""), 

364 "parameters": schema.get("parameters", {}), 

365 } 

366 if isinstance(tool, dict): 

367 func = tool.get("function", tool) 

368 return { 

369 "name": func.get("name", ""), 

370 "description": func.get("description", ""), 

371 "parameters": func.get("parameters", {}), 

372 } 

373 return { 

374 "name": getattr(tool, "name", str(tool)), 

375 "description": getattr(tool, "description", ""), 

376 "parameters": getattr(tool, "parameters", None) or {}, 

377 } 

378 

379 

380# ────────────────────────────────────────────────────────────────────── 

381# Internal helpers 

382# ────────────────────────────────────────────────────────────────────── 

383 

384 

385def _role_str(msg: Any) -> str: 

386 """Convert a Lexigram role string to a Gemini role string. 

387 

388 Args: 

389 msg: A dict or object with a ``role`` attribute (e.g. 'user', 'assistant', 'tool'). 

390 

391 Returns: 

392 'user' for user/tool/function roles, 'model' for all others. 

393 """ 

394 role = ( 

395 msg.get("role", "user") 

396 if isinstance(msg, dict) 

397 else getattr(msg, "role", "user") 

398 ) 

399 return role.value if hasattr(role, "value") else str(role) 

400 

401 

402def _extract_text(content: Any) -> str: 

403 """Extract plain text from a mixed OpenAI content list. 

404 

405 Args: 

406 content: Content field from an OpenAI message; may be a list of 

407 typed parts or any other value. 

408 

409 Returns: 

410 Concatenated text string. 

411 """ 

412 if isinstance(content, list): 

413 return " ".join( 

414 item.get("text", "") 

415 for item in content 

416 if isinstance(item, dict) and item.get("type") == "text" 

417 ) 

418 return str(content)