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

1"""Google Gemini REST API client for the Lexigram LLM routing system. 

2 

3Implements the :class:`~lexigram.contracts.ai.protocols.LLMClientProtocol` protocol 

4against the ``generativelanguage.googleapis.com`` v1beta endpoint. 

5 

6Key responsibilities: 

7* Translate OpenAI-compatible ``role``/``content`` messages into Gemini's 

8 ``contents[].parts`` format (including image ``inline_data`` parts). 

9* Map Gemini errors to the framework's typed exception hierarchy so the 

10 router can classify them without provider-specific logic. 

11* Return a normalised :class:`~lexigram.ai.llm.types.Completion` object. 

12 

13Notes: 

14 Gemini does *not* expose an OpenAI-compatible endpoint — this client 

15 uses the native REST API directly via :class:`~lexigram.http.BaseURLHTTPClient`. 

16""" 

17 

18from __future__ import annotations 

19 

20from typing import TYPE_CHECKING, Any 

21 

22from lexigram.ai.llm.exceptions import ( 

23 LLMAuthenticationError, 

24 LLMContentFilterError, 

25 LLMError, 

26 LLMModelNotFoundError, 

27 LLMRateLimitError, 

28) 

29from lexigram.ai.llm.http.client import ResilientHTTPClient 

30from lexigram.ai.llm.types import ( 

31 AIError, 

32 Completion, 

33 StreamChunk, 

34) 

35from lexigram.contracts.core.health import HealthCheckResult, HealthStatus 

36from lexigram.contracts.web.http_models import HttpStatusError 

37from lexigram.logging import ( 

38 get_logger, 

39) 

40from lexigram.result import Err, Ok, Result 

41 

42logger = get_logger(__name__) 

43 

44_BASE_URL = "https://generativelanguage.googleapis.com" 

45 

46from lexigram.ai.llm.clients.base import AbstractLLMClient 

47from lexigram.ai.llm.clients.gemini_helpers import ( 

48 inject_thinking_config, 

49 messages_to_gemini, 

50 parse_gemini_response_with_tools, 

51 parse_gemini_sse_body, 

52 tool_to_gemini_function, 

53) 

54 

55if TYPE_CHECKING: 

56 from collections.abc import AsyncIterator 

57 

58 from lexigram.ai.llm.config import ClientConfig 

59 

60__all__ = ["GeminiClient", "_messages_to_gemini"] 

61 

62_messages_to_gemini = messages_to_gemini 

63 

64 

65class GeminiClient(AbstractLLMClient): 

66 """Client for the Google Gemini REST API. 

67 

68 Wraps ``generativelanguage.googleapis.com/v1beta/models/{model}:generateContent`` 

69 and normalises the response to a standard :class:`~lexigram.ai.llm.types.Completion`. 

70 

71 The Gemini API uses API-key query-parameter authentication, not Bearer 

72 tokens. The key is read from ``ClientConfig.api_key`` and is never logged. 

73 

74 Example: 

75 >>> config = ClientConfig( 

76 ... provider="gemini", 

77 ... model="gemini-2.5-flash", 

78 ... api_key="AIza…", 

79 ... timeout=60.0, 

80 ... ) 

81 >>> client = GeminiClient(config) 

82 >>> completion = await client.complete( 

83 ... messages=[{"role": "user", "content": "Describe quantum entanglement."}], 

84 ... model="gemini-2.5-flash", 

85 ... max_tokens=512, 

86 ... ) 

87 >>> print(completion.content) 

88 """ 

89 

90 def __init__(self, config: ClientConfig) -> None: 

91 """Initialise the Gemini client. 

92 

93 Args: 

94 config: LLM configuration. ``config.api_key`` must be set; 

95 ``config.api_base`` overrides the default base URL when set. 

96 """ 

97 super().__init__(config=config) 

98 self._http: ResilientHTTPClient | None = None 

99 

100 @property 

101 def _api_key(self) -> str: 

102 """Return the raw API key string. 

103 

104 Raises: 

105 LLMAuthenticationError: When ``config.api_key`` is not set. 

106 """ 

107 key = self.config.api_key 

108 if key is None: 

109 msg = "GeminiClient requires api_key in ClientConfig" 

110 raise LLMAuthenticationError(msg) 

111 return key.get_secret_value() 

112 

113 @property 

114 def _base_url(self) -> str: 

115 """Return the API base URL, using config override when provided.""" 

116 return self.config.api_base or _BASE_URL 

117 

118 def _get_http(self) -> ResilientHTTPClient: 

119 """Return a lazily-created HTTP client instance.""" 

120 if self._http is None: 

121 self._http = ResilientHTTPClient( 

122 base_url=self._base_url, 

123 headers={"Content-Type": "application/json"}, 

124 timeout=self.config.timeout, 

125 name="gemini-client", 

126 ) 

127 return self._http 

128 

129 # ────────────────────────────────────────────────────────────────── 

130 # LLMClientProtocol protocol implementation 

131 # ────────────────────────────────────────────────────────────────── 

132 

133 async def _do_complete( # type: ignore[override] 

134 self, 

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

136 *, 

137 model: str | None = None, 

138 temperature: float = 0.2, 

139 max_tokens: int | None = None, 

140 **kwargs: Any, 

141 ) -> Result[Completion, LLMError]: 

142 """Generate a completion from Gemini. 

143 

144 Args: 

145 messages: OpenAI-compatible message list. 

146 model: Model identifier override. When ``None``, uses 

147 ``config.model``. 

148 temperature: Sampling temperature (0.0–2.0). 

149 max_tokens: Maximum output tokens. 

150 **kwargs: Ignored for protocol compatibility. 

151 

152 Returns: 

153 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable 

154 failures (rate limit, quota, content filter, model not found). 

155 

156 Raises: 

157 LLMAuthenticationError: When the API key is invalid (HTTP 401/403). 

158 AIError: For unexpected infrastructure failures. 

159 """ 

160 active_model = model or self.config.model 

161 contents = messages_to_gemini(messages) 

162 payload: dict[str, Any] = { 

163 "contents": contents, 

164 "generationConfig": { 

165 "temperature": temperature, 

166 }, 

167 } 

168 if max_tokens is not None: 

169 payload["generationConfig"]["maxOutputTokens"] = max_tokens 

170 inject_thinking_config(payload["generationConfig"], self.config) 

171 tools = kwargs.pop("tools", None) 

172 if tools: 

173 payload["tools"] = [ 

174 {"functionDeclarations": [tool_to_gemini_function(t) for t in tools]} 

175 ] 

176 

177 path = f"/v1beta/models/{active_model}:generateContent?key={self._api_key}" 

178 

179 try: 

180 http = self._get_http() 

181 response = await http.post(path, json=payload) 

182 response.raise_for_status() 

183 except ( 

184 HttpStatusError, 

185 OSError, 

186 ConnectionError, 

187 TimeoutError, 

188 RuntimeError, 

189 ) as exc: 

190 return self._handle_error_as_result(exc) 

191 

192 data: dict[str, Any] = response.json 

193 try: 

194 return Ok(parse_gemini_response_with_tools(data, active_model)) 

195 except AIError as exc: 

196 return Err(LLMContentFilterError(str(exc))) 

197 

198 async def _do_stream_chat( 

199 self, 

200 messages: list[Any], 

201 *, 

202 model: str | None = None, 

203 temperature: float = 0.2, 

204 max_tokens: int | None = None, 

205 **kwargs: Any, 

206 ) -> Result[AsyncIterator[StreamChunk], LLMError]: 

207 """Stream completion tokens from Gemini. 

208 

209 Uses ``streamGenerateContent`` with SSE (``alt=sse``) to yield 

210 incremental text deltas as :class:`StreamChunk` objects. 

211 

212 Args: 

213 messages: OpenAI-compatible message list. 

214 model: Model override. 

215 temperature: Sampling temperature. 

216 max_tokens: Maximum output tokens. 

217 **kwargs: Ignored for protocol compatibility. 

218 

219 Returns: 

220 ``Ok(AsyncIterator[StreamChunk])`` on successful connection. 

221 ``Err(LLMError)`` for recoverable failures. 

222 

223 Raises: 

224 LLMAuthenticationError: When credentials are invalid. 

225 AIError: For unexpected infrastructure failures. 

226 """ 

227 active_model = model or self.config.model 

228 contents = messages_to_gemini(messages) 

229 payload: dict[str, Any] = { 

230 "contents": contents, 

231 "generationConfig": {"temperature": temperature}, 

232 } 

233 if max_tokens is not None: 

234 payload["generationConfig"]["maxOutputTokens"] = max_tokens 

235 inject_thinking_config(payload["generationConfig"], self.config) 

236 

237 path = ( 

238 f"/v1beta/models/{active_model}:streamGenerateContent" 

239 f"?key={self._api_key}&alt=sse" 

240 ) 

241 try: 

242 http = self._get_http() 

243 response = await http.post(path, json=payload) 

244 response.raise_for_status() 

245 except ( 

246 HttpStatusError, 

247 OSError, 

248 ConnectionError, 

249 TimeoutError, 

250 RuntimeError, 

251 ) as exc: 

252 return self._handle_error_as_result(exc) 

253 

254 # The non-streaming post returns the full SSE body as text; 

255 # parse newline-delimited JSON objects from the array. 

256 result = parse_gemini_sse_body(response.text or "", active_model) 

257 

258 async def _to_async() -> AsyncIterator[StreamChunk]: 

259 for chunk in result: 

260 yield chunk 

261 

262 return Ok(_to_async()) 

263 

264 async def _do_chat( # type: ignore[override] 

265 self, 

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

267 tools: list[Any] | None = None, 

268 *, 

269 model: str | None = None, 

270 temperature: float = 0.2, 

271 max_tokens: int | None = None, 

272 **kwargs: Any, 

273 ) -> Result[Completion, LLMError]: 

274 """Generate completion with optional Gemini function calling. 

275 

276 Tool calling is handled by :meth:`_do_complete` via 

277 ``complete(..., tools=...)``; this method forwards the tool 

278 descriptors to keep the ``chat`` code path consistent. 

279 

280 Args: 

281 messages: OpenAI-compatible message list. 

282 tools: Optional tool/function descriptors. 

283 model: Model override. 

284 temperature: Sampling temperature. 

285 max_tokens: Maximum output tokens. 

286 **kwargs: Ignored for protocol compatibility. 

287 

288 Returns: 

289 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable 

290 failures. 

291 """ 

292 return await self._do_complete( 

293 messages, 

294 model=model, 

295 temperature=temperature, 

296 max_tokens=max_tokens, 

297 tools=tools, 

298 **kwargs, 

299 ) 

300 

301 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

302 """Perform a lightweight health check against the Gemini API. 

303 

304 Attempts to list models — a zero-cost metadata call that verifies 

305 the API key is valid and the service is reachable. 

306 

307 Args: 

308 timeout: Seconds to wait for the response (informational; the 

309 client-level timeout in ``config.timeout`` is used). 

310 

311 Returns: 

312 :class:`~lexigram.contracts.core.health.HealthCheckResult`. 

313 """ 

314 try: 

315 http = self._get_http() 

316 response = await http.get( 

317 f"/v1beta/models?key={self._api_key}&pageSize=1", 

318 ) 

319 response.raise_for_status() 

320 return HealthCheckResult(component="gemini", status=HealthStatus.HEALTHY) 

321 except (OSError, ConnectionError, TimeoutError, RuntimeError) as exc: 

322 return HealthCheckResult( 

323 component="gemini", 

324 status=HealthStatus.UNHEALTHY, 

325 error=str(exc), 

326 ) 

327 

328 def _handle_error_as_result(self, error: Exception) -> Result[Any, LLMError]: 

329 """Map a caught exception to ``Err`` (recoverable) or re-raise (infra).""" 

330 status: int | None = None 

331 if isinstance(error, HttpStatusError): 

332 status = error.status 

333 

334 if status in (401, 403): 

335 raise LLMAuthenticationError( 

336 f"gemini: authentication failed ({status}): {error}" 

337 ) from error 

338 if status == 429: 

339 return Err(LLMRateLimitError(f"gemini: rate limit exceeded: {error}")) 

340 if status == 400: 

341 err_str = str(error).lower() 

342 if "not found" in err_str or "does not exist" in err_str: 

343 return Err(LLMModelNotFoundError(f"gemini: model not found: {error}")) 

344 return Err(LLMError(f"gemini: invalid request: {error}")) 

345 if status == 404: 

346 return Err(LLMModelNotFoundError(f"gemini: model not found: {error}")) 

347 raise AIError(f"gemini: infrastructure error: {error}") from error 

348 

349 async def close(self) -> None: 

350 """Close the underlying HTTP client and release connections.""" 

351 if self._http is not None: 

352 await self._http.close() 

353 self._http = None 

354 await super().close()