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

1"""Google Vertex AI LLM client for the Lexigram LLM routing system. 

2 

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

4protocol against the Vertex AI ``predict`` REST endpoint using the 

5``google-auth`` library for service-account OAuth2 authentication. 

6 

7The Vertex AI endpoint for Gemini models follows this pattern:: 

8 

9 https://{region}-aiplatform.googleapis.com/v1/projects/{project}/ 

10 locations/{location}/publishers/google/models/{model}:generateContent 

11 

12Configuration is sourced from ``ClientConfig.extra``: 

13 

14* ``vertex_project`` — GCP project ID (required) 

15* ``vertex_location`` — Vertex location, e.g. ``us-central1`` (required) 

16* ``vertex_region`` — Region prefix for the hostname; defaults to 

17 ``vertex_location`` when omitted. 

18* ``vertex_credentials_file`` — Path to a service account JSON credentials 

19 file. When absent, Application Default Credentials (ADC) are used. 

20 

21Notes: 

22 ``google-auth`` and ``google-auth-httplib2`` are optional dependencies. 

23 An :class:`ImportError` is raised at construction time if they are absent. 

24""" 

25 

26from __future__ import annotations 

27 

28from typing import TYPE_CHECKING, Any, cast 

29 

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

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

32 inject_thinking_config, 

33 messages_to_gemini, 

34 parse_gemini_response_with_tools, 

35 parse_gemini_sse_body, 

36 tool_to_gemini_function, 

37) 

38from lexigram.ai.llm.exceptions import ( 

39 LLMAuthenticationError, 

40 LLMError, 

41 LLMModelNotFoundError, 

42 LLMRateLimitError, 

43) 

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

45from lexigram.ai.llm.types import ( 

46 AIError, 

47 Completion, 

48 StreamChunk, 

49) 

50from lexigram.contracts.core import HealthCheckResult, HealthStatus 

51from lexigram.contracts.web.http_models import HttpStatusError 

52from lexigram.logging import ( 

53 get_logger, 

54) 

55from lexigram.result import Err, Ok, Result 

56 

57if TYPE_CHECKING: 

58 from collections.abc import AsyncIterator 

59 

60 from lexigram.ai.llm.config import ClientConfig 

61 

62logger = get_logger(__name__) 

63 

64__all__ = ["VertexAIClient"] 

65 

66_SCOPE = "https://www.googleapis.com/auth/cloud-platform" 

67 

68 

69class VertexAIClient(AbstractLLMClient): 

70 """Google Vertex AI client using the native REST API. 

71 

72 Authenticates via ``google-auth`` service-account credentials and routes 

73 requests to the Vertex AI ``generateContent`` endpoint. All message and 

74 tool conversion reuses Gemini-native helpers because Vertex AI exposes the 

75 same Gemini model contract. 

76 

77 Args: 

78 config: LLM configuration. ``config.extra`` must contain 

79 ``vertex_project`` and ``vertex_location``. 

80 """ 

81 

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

83 """Initialise the Vertex AI client. 

84 

85 Args: 

86 config: LLM configuration with Vertex-specific ``extra`` keys. 

87 

88 Raises: 

89 ImportError: If ``google-auth`` is not installed. 

90 ValueError: If required ``extra`` keys are missing. 

91 """ 

92 super().__init__(config=config) 

93 

94 try: 

95 import google.auth # noqa: F401 

96 import google.auth.transport.requests # noqa: F401 

97 except ImportError as exc: 

98 raise ImportError( 

99 "VertexAIClient requires 'google-auth'. " 

100 "Install with: pip install google-cloud-aiplatform" 

101 ) from exc 

102 

103 extra: dict[str, Any] = config.extra or {} 

104 self._project: str = extra.get("vertex_project", "") 

105 self._location: str = extra.get("vertex_location", "") 

106 self._region: str = extra.get("vertex_region", "") or self._location 

107 self._credentials_file: str | None = extra.get("vertex_credentials_file") 

108 

109 if not self._project: 

110 raise ValueError( 

111 "VertexAIClient requires 'vertex_project' in ClientConfig.extra" 

112 ) 

113 if not self._location: 

114 raise ValueError( 

115 "VertexAIClient requires 'vertex_location' in ClientConfig.extra" 

116 ) 

117 

118 self._http: ResilientHTTPClient | None = None 

119 self._access_token: str | None = None 

120 

121 # ────────────────────────────────────────────────────────────────── 

122 # Token acquisition 

123 # ────────────────────────────────────────────────────────────────── 

124 

125 async def _get_access_token(self) -> str: 

126 """Return a valid OAuth2 access token, refreshing when expired. 

127 

128 Returns: 

129 Bearer token string. 

130 

131 Raises: 

132 AIError: On credential failure. 

133 """ 

134 import google.auth 

135 import google.auth.transport.requests 

136 

137 try: 

138 if self._credentials_file: 

139 from google.oauth2 import ( 

140 service_account, 

141 ) 

142 

143 creds = service_account.Credentials.from_service_account_file( 

144 self._credentials_file, scopes=[_SCOPE] 

145 ) 

146 else: 

147 creds, _ = google.auth.default(scopes=[_SCOPE]) 

148 

149 req = google.auth.transport.requests.Request() 

150 if not creds.valid: 

151 creds.refresh(req) 

152 return cast("str", creds.token) 

153 except Exception as exc: 

154 raise AIError(f"vertex: failed to acquire access token: {exc}") from exc 

155 

156 # ────────────────────────────────────────────────────────────────── 

157 # HTTP client 

158 # ────────────────────────────────────────────────────────────────── 

159 

160 async def _get_http(self) -> ResilientHTTPClient: 

161 """Return a lazily-created, token-refreshed HTTP client.""" 

162 token = await self._get_access_token() 

163 base_url = f"https://{self._region}-aiplatform.googleapis.com" 

164 if self._http is None: 

165 self._http = ResilientHTTPClient( 

166 base_url=base_url, 

167 headers={ 

168 "Authorization": f"Bearer {token}", 

169 "Content-Type": "application/json", 

170 }, 

171 timeout=self.config.timeout, 

172 name="vertex-ai-client", 

173 ) 

174 else: 

175 # Refresh token header in-place 

176 self._http.headers["Authorization"] = f"Bearer {token}" 

177 return self._http 

178 

179 def _model_path(self, model: str) -> str: 

180 """Build the Vertex AI model resource path. 

181 

182 Args: 

183 model: Model ID, e.g. ``gemini-1.5-pro``. 

184 

185 Returns: 

186 Full resource path string. 

187 """ 

188 return ( 

189 f"/v1/projects/{self._project}/locations/{self._location}" 

190 f"/publishers/google/models/{model}" 

191 ) 

192 

193 # ────────────────────────────────────────────────────────────────── 

194 # LLMClientProtocol implementation 

195 # ────────────────────────────────────────────────────────────────── 

196 

197 async def _do_complete( 

198 self, 

199 messages: list[Any], 

200 *, 

201 model: str | None = None, 

202 temperature: float = 0.2, 

203 max_tokens: int | None = None, 

204 **kwargs: Any, 

205 ) -> Result[Completion, LLMError]: 

206 """Generate a completion from Vertex AI. 

207 

208 Args: 

209 messages: OpenAI-compatible message list. 

210 model: Model override. 

211 temperature: Sampling temperature. 

212 max_tokens: Maximum output tokens. 

213 **kwargs: Ignored for protocol compatibility. 

214 

215 Returns: 

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

217 failures. 

218 

219 Raises: 

220 LLMAuthenticationError: When credentials are invalid. 

221 AIError: For unexpected infrastructure failures. 

222 """ 

223 active_model = model or self.config.model 

224 contents = messages_to_gemini(messages) 

225 payload: dict[str, Any] = { 

226 "contents": contents, 

227 "generationConfig": {"temperature": temperature}, 

228 } 

229 if max_tokens is not None: 

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

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

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

233 if tools: 

234 payload["tools"] = [ 

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

236 ] 

237 

238 path = f"{self._model_path(active_model)}:generateContent" 

239 try: 

240 http = await self._get_http() 

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

242 response.raise_for_status() 

243 except ( 

244 HttpStatusError, 

245 OSError, 

246 ConnectionError, 

247 TimeoutError, 

248 RuntimeError, 

249 ) as exc: 

250 return self._handle_error_as_result(exc) 

251 

252 return Ok(parse_gemini_response_with_tools(response.json, active_model)) 

253 

254 async def _do_stream_chat( 

255 self, 

256 messages: list[Any], 

257 *, 

258 model: str | None = None, 

259 temperature: float = 0.2, 

260 max_tokens: int | None = None, 

261 **kwargs: Any, 

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

263 """Stream completion tokens from Vertex AI. 

264 

265 Args: 

266 messages: OpenAI-compatible message list. 

267 model: Model override. 

268 temperature: Sampling temperature. 

269 max_tokens: Maximum output tokens. 

270 **kwargs: Ignored for protocol compatibility. 

271 

272 Returns: 

273 ``Ok(AsyncIterator[StreamChunk])`` on success. 

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

275 

276 Raises: 

277 LLMAuthenticationError: When credentials are invalid. 

278 AIError: For unexpected infrastructure failures. 

279 """ 

280 active_model = model or self.config.model 

281 contents = messages_to_gemini(messages) 

282 payload: dict[str, Any] = { 

283 "contents": contents, 

284 "generationConfig": {"temperature": temperature}, 

285 } 

286 if max_tokens is not None: 

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

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

289 

290 path = f"{self._model_path(active_model)}:streamGenerateContent?alt=sse" 

291 try: 

292 http = await self._get_http() 

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

294 response.raise_for_status() 

295 except ( 

296 HttpStatusError, 

297 OSError, 

298 ConnectionError, 

299 TimeoutError, 

300 RuntimeError, 

301 ) as exc: 

302 return self._handle_error_as_result(exc) 

303 

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

305 

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

307 for chunk in result: 

308 yield chunk 

309 

310 return Ok(_to_async()) 

311 

312 async def _do_chat( 

313 self, 

314 messages: list[Any], 

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

316 *, 

317 model: str | None = None, 

318 temperature: float = 0.2, 

319 max_tokens: int | None = None, 

320 **kwargs: Any, 

321 ) -> Result[Completion, LLMError]: 

322 """Generate completion with optional tool calling on Vertex AI. 

323 

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

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

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

327 

328 Args: 

329 messages: OpenAI-compatible message list. 

330 tools: Optional tool descriptors. 

331 model: Model override. 

332 temperature: Sampling temperature. 

333 max_tokens: Maximum output tokens. 

334 **kwargs: Ignored for protocol compatibility. 

335 

336 Returns: 

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

338 failures. 

339 """ 

340 return await self._do_complete( 

341 messages, 

342 model=model, 

343 temperature=temperature, 

344 max_tokens=max_tokens, 

345 tools=tools, 

346 **kwargs, 

347 ) 

348 

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

350 """Probe Vertex AI with a minimal generateContent request. 

351 

352 Args: 

353 timeout: Informational only. 

354 

355 Returns: 

356 Structured :class:`~lexigram.contracts.core.health.HealthCheckResult`. 

357 """ 

358 try: 

359 http = await self._get_http() 

360 path = ( 

361 f"/v1/projects/{self._project}/locations/{self._location}" 

362 "/publishers/google/models" 

363 ) 

364 resp = await http.get(path) 

365 resp.raise_for_status() 

366 except Exception as exc: # noqa: BLE001 - transport/auth SDK stack varies 

367 return HealthCheckResult( 

368 component="llm.vertex_ai", 

369 status=HealthStatus.UNHEALTHY, 

370 error=str(exc), 

371 details={ 

372 "project": self._project, 

373 "location": self._location, 

374 "model": self.config.model, 

375 }, 

376 ) 

377 

378 return HealthCheckResult( 

379 component="llm.vertex_ai", 

380 status=HealthStatus.HEALTHY, 

381 details={ 

382 "project": self._project, 

383 "location": self._location, 

384 "model": self.config.model, 

385 }, 

386 ) 

387 

388 async def close(self) -> None: 

389 """Close the underlying HTTP client.""" 

390 if self._http is not None: 

391 await self._http.close() 

392 self._http = None 

393 await super().close() 

394 

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

396 """Map a caught exception to ``Err`` or re-raise for infrastructure failures.""" 

397 status: int | None = None 

398 if isinstance(error, HttpStatusError): 

399 status = error.status 

400 

401 if status in (401, 403): 

402 raise LLMAuthenticationError( 

403 f"vertex: authentication failed ({status}): {error}" 

404 ) from error 

405 if status == 429: 

406 return Err(LLMRateLimitError(f"vertex: rate limit exceeded: {error}")) 

407 if status == 404: 

408 return Err(LLMModelNotFoundError(f"vertex: model not found: {error}")) 

409 raise AIError(f"vertex: infrastructure error: {error}") from error