Coverage for pydantic_ai_slim/pydantic_ai/models/gemini.py: 93.66%

1from __future__ import annotations as _annotations (empty)

2 

3import os (empty)

4import re (empty)

5from collections.abc import AsyncIterator, Sequence (empty)

6from contextlib import asynccontextmanager (empty)

7from copy import deepcopy (empty)

8from dataclasses import dataclass, field (empty)

9from datetime import datetime (empty)

10from typing import Annotated, Any, Literal, Protocol, Union, cast (empty)

11from uuid import uuid4 (empty)

12 

13import pydantic (empty)

14from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse (empty)

15from typing_extensions import NotRequired, TypedDict, assert_never (empty)

16 

17from .. import UnexpectedModelBehavior, _utils, exceptions, usage (empty)

18from ..messages import ( (empty)

19 ModelMessage, 

20 ModelRequest, 

21 ModelResponse, 

22 ModelResponsePart, 

23 ModelResponseStreamEvent, 

24 RetryPromptPart, 

25 SystemPromptPart, 

26 TextPart, 

27 ToolCallPart, 

28 ToolReturnPart, 

29 UserPromptPart, 

30) 

31from ..settings import ModelSettings (empty)

32from ..tools import ToolDefinition (empty)

33from . import ( (empty)

34 AgentModel, 

35 Model, 

36 StreamedResponse, 

37 cached_async_http_client, 

38 check_allow_model_requests, 

39 get_user_agent, 

40) 

41 

42GeminiModelName = Literal[ (empty)

43 'gemini-1.5-flash', 'gemini-1.5-flash-8b', 'gemini-1.5-pro', 'gemini-1.0-pro', 'gemini-2.0-flash-exp' 

44] 

45"""Named Gemini models. (empty)

46 

47See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list. 

48""" 

49 

50 

51class GeminiModelSettings(ModelSettings): (empty)

52 """Settings used for a Gemini model request.""" 

53 

54 # This class is a placeholder for any future gemini-specific settings 

55 

56 

57@dataclass(init=False) (empty)

58class GeminiModel(Model): (empty)

59 """A model that uses Gemini via `generativelanguage.googleapis.com` API. 

60 

61 This is implemented from scratch rather than using a dedicated SDK, good API documentation is 

62 available [here](https://ai.google.dev/api). 

63 

64 Apart from `__init__`, all methods are private or match those of the base class. 

65 """ 

66 

67 model_name: GeminiModelName (empty)

68 auth: AuthProtocol (empty)

69 http_client: AsyncHTTPClient (empty)

70 url: str (empty)

71 

72 def __init__( (empty)

73 self, 

74 model_name: GeminiModelName, 

75 *, 

76 api_key: str | None = None, 

77 http_client: AsyncHTTPClient | None = None, 

78 url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:', 

79 ): 

80 """Initialize a Gemini model. 

81 

82 Args: 

83 model_name: The name of the model to use. 

84 api_key: The API key to use for authentication, if not provided, the `GEMINI_API_KEY` environment variable 

85 will be used if available. 

86 http_client: An existing `httpx.AsyncClient` to use for making HTTP requests. 

87 url_template: The URL template to use for making requests, you shouldn't need to change this, 

88 docs [here](https://ai.google.dev/gemini-api/docs/quickstart?lang=rest#make-first-request), 

89 `model` is substituted with the model name, and `function` is added to the end of the URL. 

90 """ 

91 self.model_name = model_name (empty)

92 if api_key is None: (empty)

93 if env_api_key := os.getenv('GEMINI_API_KEY'): (empty)

94 api_key = env_api_key (empty)

95 else: 

96 raise exceptions.UserError('API key must be provided or set in the GEMINI_API_KEY environment variable') (empty)

97 self.auth = ApiKeyAuth(api_key) (empty)

98 self.http_client = http_client or cached_async_http_client() (empty)

99 self.url = url_template.format(model=model_name) (empty)

100 

101 async def agent_model( (empty)

102 self, 

103 *, 

104 function_tools: list[ToolDefinition], 

105 allow_text_result: bool, 

106 result_tools: list[ToolDefinition], 

107 ) -> GeminiAgentModel: 

108 check_allow_model_requests() (empty)

109 return GeminiAgentModel( (empty)

110 http_client=self.http_client, 

111 model_name=self.model_name, 

112 auth=self.auth, 

113 url=self.url, 

114 function_tools=function_tools, 

115 allow_text_result=allow_text_result, 

116 result_tools=result_tools, 

117 ) 

118 

119 def name(self) -> str: (empty)

120 return f'google-gla:{self.model_name}' (empty)

121 

122 

123class AuthProtocol(Protocol): (empty)

124 """Abstract definition for Gemini authentication.""" (empty)

125 

126 async def headers(self) -> dict[str, str]: ... (empty)

127 

128 

129@dataclass (empty)

130class ApiKeyAuth: (empty)

131 """Authentication using an API key for the `X-Goog-Api-Key` header.""" 

132 

133 api_key: str (empty)

134 

135 async def headers(self) -> dict[str, str]: (empty)

136 # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest 

137 return {'X-Goog-Api-Key': self.api_key} (empty)

138 

139 

140@dataclass(init=False) (empty)

141class GeminiAgentModel(AgentModel): (empty)

142 """Implementation of `AgentModel` for Gemini models.""" 

143 

144 http_client: AsyncHTTPClient (empty)

145 model_name: GeminiModelName (empty)

146 auth: AuthProtocol (empty)

147 tools: _GeminiTools | None (empty)

148 tool_config: _GeminiToolConfig | None (empty)

149 url: str (empty)

150 

151 def __init__( (empty)

152 self, 

153 http_client: AsyncHTTPClient, 

154 model_name: GeminiModelName, 

155 auth: AuthProtocol, 

156 url: str, 

157 function_tools: list[ToolDefinition], 

158 allow_text_result: bool, 

159 result_tools: list[ToolDefinition], 

160 ): 

161 tools = [_function_from_abstract_tool(t) for t in function_tools] (empty)

162 if result_tools: (empty)

163 tools += [_function_from_abstract_tool(t) for t in result_tools] (empty)

164 

165 if allow_text_result: (empty)

166 tool_config = None (empty)

167 else: 

168 tool_config = _tool_config([t['name'] for t in tools]) (empty)

169 

170 self.http_client = http_client (empty)

171 self.model_name = model_name (empty)

172 self.auth = auth (empty)

173 self.tools = _GeminiTools(function_declarations=tools) if tools else None (empty)

174 self.tool_config = tool_config (empty)

175 self.url = url (empty)

176 

177 async def request( (empty)

178 self, messages: list[ModelMessage], model_settings: ModelSettings | None 

179 ) -> tuple[ModelResponse, usage.Usage]: 

180 async with self._make_request( (empty)

181 messages, False, cast(GeminiModelSettings, model_settings or {}) 

182 ) as http_response: 

183 response = _gemini_response_ta.validate_json(await http_response.aread()) (empty)

184 return self._process_response(response), _metadata_as_usage(response) (empty)

185 

186 @asynccontextmanager (empty)

187 async def request_stream( (empty)

188 self, messages: list[ModelMessage], model_settings: ModelSettings | None 

189 ) -> AsyncIterator[StreamedResponse]: 

190 async with self._make_request(messages, True, cast(GeminiModelSettings, model_settings or {})) as http_response: (empty)

191 yield await self._process_streamed_response(http_response) (empty)

192 

193 @asynccontextmanager (empty)

194 async def _make_request( (empty)

195 self, messages: list[ModelMessage], streamed: bool, model_settings: GeminiModelSettings 

196 ) -> AsyncIterator[HTTPResponse]: 

197 sys_prompt_parts, contents = self._message_to_gemini_content(messages) (empty)

198 

199 request_data = _GeminiRequest(contents=contents) (empty)

200 if sys_prompt_parts: (empty)

201 request_data['system_instruction'] = _GeminiTextContent(role='user', parts=sys_prompt_parts) (empty)

202 if self.tools is not None: (empty)

203 request_data['tools'] = self.tools (empty)

204 if self.tool_config is not None: (empty)

205 request_data['tool_config'] = self.tool_config (empty)

206 

207 generation_config: _GeminiGenerationConfig = {} (empty)

208 if model_settings: (empty)

209 if (max_tokens := model_settings.get('max_tokens')) is not None: 209 ↛ 211line 209 didn't jump to line 211 because the condition on line 209 was always true(empty)

210 generation_config['max_output_tokens'] = max_tokens (empty)

211 if (temperature := model_settings.get('temperature')) is not None: 211 ↛ 213line 211 didn't jump to line 213 because the condition on line 211 was always true(empty)

212 generation_config['temperature'] = temperature (empty)

213 if (top_p := model_settings.get('top_p')) is not None: 213 ↛ 215line 213 didn't jump to line 215 because the condition on line 213 was always true(empty)

214 generation_config['top_p'] = top_p (empty)

215 if (presence_penalty := model_settings.get('presence_penalty')) is not None: 215 ↛ 217line 215 didn't jump to line 217 because the condition on line 215 was always true(empty)

216 generation_config['presence_penalty'] = presence_penalty (empty)

217 if (frequency_penalty := model_settings.get('frequency_penalty')) is not None: 217 ↛ 219line 217 didn't jump to line 219 because the condition on line 217 was always true(empty)

218 generation_config['frequency_penalty'] = frequency_penalty (empty)

219 if generation_config: (empty)

220 request_data['generation_config'] = generation_config (empty)

221 

222 url = self.url + ('streamGenerateContent' if streamed else 'generateContent') (empty)

223 

224 headers = { (empty)

225 'Content-Type': 'application/json', 

226 'User-Agent': get_user_agent(), 

227 **await self.auth.headers(), 

228 } 

229 

230 request_json = _gemini_request_ta.dump_json(request_data, by_alias=True) (empty)

231 

232 async with self.http_client.stream( (empty)

233 'POST', 

234 url, 

235 content=request_json, 

236 headers=headers, 

237 timeout=model_settings.get('timeout', USE_CLIENT_DEFAULT), 

238 ) as r: 

239 if r.status_code != 200: (empty)

240 await r.aread() (empty)

241 raise exceptions.UnexpectedModelBehavior(f'Unexpected response from gemini {r.status_code}', r.text) (empty)

242 yield r (empty)

243 

244 def _process_response(self, response: _GeminiResponse) -> ModelResponse: (empty)

245 if len(response['candidates']) != 1: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true(empty)

246 raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response') 

247 parts = response['candidates'][0]['content']['parts'] (empty)

248 return _process_response_from_parts(parts, model_name=self.model_name) (empty)

249 

250 async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse: (empty)

251 """Process a streamed response, and prepare a streaming response to return.""" 

252 aiter_bytes = http_response.aiter_bytes() (empty)

253 start_response: _GeminiResponse | None = None (empty)

254 content = bytearray() (empty)

255 

256 async for chunk in aiter_bytes: (empty)

257 content.extend(chunk) (empty)

258 responses = _gemini_streamed_response_ta.validate_json( (empty)

259 content, 

260 experimental_allow_partial='trailing-strings', 

261 ) 

262 if responses: 262 ↛ 256line 262 didn't jump to line 256 because the condition on line 262 was always true(empty)

263 last = responses[-1] (empty)

264 if last['candidates'] and last['candidates'][0]['content']['parts']: (empty)

265 start_response = last (empty)

266 break (empty)

267 

268 if start_response is None: (empty)

269 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') (empty)

270 

271 return GeminiStreamedResponse(_model_name=self.model_name, _content=content, _stream=aiter_bytes) (empty)

272 

273 @classmethod (empty)

274 def _message_to_gemini_content( (empty)

275 cls, messages: list[ModelMessage] 

276 ) -> tuple[list[_GeminiTextPart], list[_GeminiContent]]: 

277 sys_prompt_parts: list[_GeminiTextPart] = [] (empty)

278 contents: list[_GeminiContent] = [] (empty)

279 for m in messages: (empty)

280 if isinstance(m, ModelRequest): (empty)

281 message_parts: list[_GeminiPartUnion] = [] (empty)

282 

283 for part in m.parts: (empty)

284 if isinstance(part, SystemPromptPart): (empty)

285 sys_prompt_parts.append(_GeminiTextPart(text=part.content)) (empty)

286 elif isinstance(part, UserPromptPart): (empty)

287 message_parts.append(_GeminiTextPart(text=part.content)) (empty)

288 elif isinstance(part, ToolReturnPart): (empty)

289 message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object())) (empty)

290 elif isinstance(part, RetryPromptPart): (empty)

291 if part.tool_name is None: 291 ↛ 292line 291 didn't jump to line 292 because the condition on line 291 was never true(empty)

292 message_parts.append(_GeminiTextPart(text=part.model_response())) 

293 else: 

294 response = {'call_error': part.model_response()} (empty)

295 message_parts.append(_response_part_from_response(part.tool_name, response)) (empty)

296 else: 

297 assert_never(part) 

298 

299 if message_parts: 299 ↛ 279line 299 didn't jump to line 279 because the condition on line 299 was always true(empty)

300 contents.append(_GeminiContent(role='user', parts=message_parts)) (empty)

301 elif isinstance(m, ModelResponse): (empty)

302 contents.append(_content_model_response(m)) (empty)

303 else: 

304 assert_never(m) 

305 

306 return sys_prompt_parts, contents (empty)

307 

308 

309@dataclass (empty)

310class GeminiStreamedResponse(StreamedResponse): (empty)

311 """Implementation of `StreamedResponse` for the Gemini model.""" 

312 

313 _content: bytearray (empty)

314 _stream: AsyncIterator[bytes] (empty)

315 _timestamp: datetime = field(default_factory=_utils.now_utc, init=False) (empty)

316 

317 async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: (empty)

318 async for gemini_response in self._get_gemini_responses(): (empty)

319 candidate = gemini_response['candidates'][0] (empty)

320 gemini_part: _GeminiPartUnion 

321 for gemini_part in candidate['content']['parts']: (empty)

322 if 'text' in gemini_part: (empty)

323 # Using vendor_part_id=None means we can produce multiple text parts if their deltas are sprinkled 

324 # amongst the tool call deltas 

325 yield self._parts_manager.handle_text_delta(vendor_part_id=None, content=gemini_part['text']) (empty)

326 

327 elif 'function_call' in gemini_part: 327 ↛ 341line 327 didn't jump to line 341 because the condition on line 327 was always true(empty)

328 # Here, we assume all function_call parts are complete and don't have deltas. 

329 # We do this by assigning a unique randomly generated "vendor_part_id". 

330 # We need to confirm whether this is actually true, but if it isn't, we can still handle it properly 

331 # it would just be a bit more complicated. And we'd need to confirm the intended semantics. 

332 maybe_event = self._parts_manager.handle_tool_call_delta( (empty)

333 vendor_part_id=uuid4(), 

334 tool_name=gemini_part['function_call']['name'], 

335 args=gemini_part['function_call']['args'], 

336 tool_call_id=None, 

337 ) 

338 if maybe_event is not None: 338 ↛ 321line 338 didn't jump to line 321 because the condition on line 338 was always true(empty)

339 yield maybe_event (empty)

340 else: 

341 assert 'function_response' in gemini_part, f'Unexpected part: {gemini_part}' 

342 

343 async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]: (empty)

344 # This method exists to ensure we only yield completed items, so we don't need to worry about 

345 # partial gemini responses, which would make everything more complicated 

346 

347 gemini_responses: list[_GeminiResponse] = [] (empty)

348 current_gemini_response_index = 0 (empty)

349 # Right now, there are some circumstances where we will have information that could be yielded sooner than it is 

350 # But changing that would make things a lot more complicated. 

351 async for chunk in self._stream: (empty)

352 self._content.extend(chunk) (empty)

353 

354 gemini_responses = _gemini_streamed_response_ta.validate_json( (empty)

355 self._content, 

356 experimental_allow_partial='trailing-strings', 

357 ) 

358 

359 # The idea: yield only up to the latest response, which might still be partial. 

360 # Note that if the latest response is complete, we could yield it immediately, but there's not a good 

361 # allow_partial API to determine if the last item in the list is complete. 

362 responses_to_yield = gemini_responses[:-1] (empty)

363 for r in responses_to_yield[current_gemini_response_index:]: (empty)

364 current_gemini_response_index += 1 (empty)

365 self._usage += _metadata_as_usage(r) (empty)

366 yield r (empty)

367 

368 # Now yield the final response, which should be complete 

369 if gemini_responses: 369 ↛ exitline 369 didn't return from function '_get_gemini_responses' because the condition on line 369 was always true(empty)

370 r = gemini_responses[-1] (empty)

371 self._usage += _metadata_as_usage(r) (empty)

372 yield r (empty)

373 

374 def timestamp(self) -> datetime: (empty)

375 return self._timestamp (empty)

376 

377 

378# We use typed dicts to define the Gemini API response schema 

379# once Pydantic partial validation supports, dataclasses, we could revert to using them 

380# TypeAdapters take care of validation and serialization 

381 

382 

383@pydantic.with_config(pydantic.ConfigDict(defer_build=True)) (empty)

384class _GeminiRequest(TypedDict): (empty)

385 """Schema for an API request to the Gemini API. 

386 

387 See <https://ai.google.dev/api/generate-content#request-body> for API docs. 

388 """ 

389 

390 contents: list[_GeminiContent] (empty)

391 tools: NotRequired[_GeminiTools] (empty)

392 tool_config: NotRequired[_GeminiToolConfig] (empty)

393 # we don't implement `generationConfig`, instead we use a named tool for the response 

394 system_instruction: NotRequired[_GeminiTextContent] (empty)

395 """ (empty)

396 Developer generated system instructions, see 

397 <https://ai.google.dev/gemini-api/docs/system-instructions?lang=rest> 

398 """ 

399 generation_config: NotRequired[_GeminiGenerationConfig] (empty)

400 

401 

402class _GeminiGenerationConfig(TypedDict, total=False): (empty)

403 """Schema for an API request to the Gemini API. 

404 

405 Note there are many additional fields available that have not been added yet. 

406 

407 See <https://ai.google.dev/api/generate-content#generationconfig> for API docs. 

408 """ 

409 

410 max_output_tokens: int (empty)

411 temperature: float (empty)

412 top_p: float (empty)

413 presence_penalty: float (empty)

414 frequency_penalty: float (empty)

415 

416 

417class _GeminiContent(TypedDict): (empty)

418 role: Literal['user', 'model'] (empty)

419 parts: list[_GeminiPartUnion] (empty)

420 

421 

422def _content_model_response(m: ModelResponse) -> _GeminiContent: (empty)

423 parts: list[_GeminiPartUnion] = [] (empty)

424 for item in m.parts: (empty)

425 if isinstance(item, ToolCallPart): (empty)

426 parts.append(_function_call_part_from_call(item)) (empty)

427 elif isinstance(item, TextPart): (empty)

428 if item.content: (empty)

429 parts.append(_GeminiTextPart(text=item.content)) (empty)

430 else: 

431 assert_never(item) 

432 return _GeminiContent(role='model', parts=parts) (empty)

433 

434 

435class _GeminiTextPart(TypedDict): (empty)

436 text: str (empty)

437 

438 

439class _GeminiFunctionCallPart(TypedDict): (empty)

440 function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')] (empty)

441 

442 

443def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart: (empty)

444 return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict())) (empty)

445 

446 

447def _process_response_from_parts( (empty)

448 parts: Sequence[_GeminiPartUnion], model_name: GeminiModelName, timestamp: datetime | None = None 

449) -> ModelResponse: 

450 items: list[ModelResponsePart] = [] (empty)

451 for part in parts: (empty)

452 if 'text' in part: (empty)

453 items.append(TextPart(content=part['text'])) (empty)

454 elif 'function_call' in part: 454 ↛ 461line 454 didn't jump to line 461 because the condition on line 454 was always true(empty)

455 items.append( (empty)

456 ToolCallPart( 

457 tool_name=part['function_call']['name'], 

458 args=part['function_call']['args'], 

459 ) 

460 ) 

461 elif 'function_response' in part: 

462 raise exceptions.UnexpectedModelBehavior( 

463 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}' 

464 ) 

465 return ModelResponse(parts=items, model_name=model_name, timestamp=timestamp or _utils.now_utc()) (empty)

466 

467 

468class _GeminiFunctionCall(TypedDict): (empty)

469 """See <https://ai.google.dev/api/caching#FunctionCall>.""" 

470 

471 name: str (empty)

472 args: dict[str, Any] (empty)

473 

474 

475class _GeminiFunctionResponsePart(TypedDict): (empty)

476 function_response: Annotated[_GeminiFunctionResponse, pydantic.Field(alias='functionResponse')] (empty)

477 

478 

479def _response_part_from_response(name: str, response: dict[str, Any]) -> _GeminiFunctionResponsePart: (empty)

480 return _GeminiFunctionResponsePart(function_response=_GeminiFunctionResponse(name=name, response=response)) (empty)

481 

482 

483class _GeminiFunctionResponse(TypedDict): (empty)

484 """See <https://ai.google.dev/api/caching#FunctionResponse>.""" 

485 

486 name: str (empty)

487 response: dict[str, Any] (empty)

488 

489 

490def _part_discriminator(v: Any) -> str: (empty)

491 if isinstance(v, dict): 491 ↛ 498line 491 didn't jump to line 498 because the condition on line 491 was always true(empty)

492 if 'text' in v: (empty)

493 return 'text' (empty)

494 elif 'functionCall' in v or 'function_call' in v: (empty)

495 return 'function_call' (empty)

496 elif 'functionResponse' in v or 'function_response' in v: (empty)

497 return 'function_response' (empty)

498 return 'text' (empty)

499 

500 

501# See <https://ai.google.dev/api/caching#Part> 

502# we don't currently support other part types 

503# TODO discriminator 

504_GeminiPartUnion = Annotated[ (empty)

505 Union[ 

506 Annotated[_GeminiTextPart, pydantic.Tag('text')], 

507 Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')], 

508 Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')], 

509 ], 

510 pydantic.Discriminator(_part_discriminator), 

511] 

512 

513 

514class _GeminiTextContent(TypedDict): (empty)

515 role: Literal['user', 'model'] (empty)

516 parts: list[_GeminiTextPart] (empty)

517 

518 

519class _GeminiTools(TypedDict): (empty)

520 function_declarations: list[Annotated[_GeminiFunction, pydantic.Field(alias='functionDeclarations')]] (empty)

521 

522 

523class _GeminiFunction(TypedDict): (empty)

524 name: str (empty)

525 description: str (empty)

526 parameters: NotRequired[dict[str, Any]] (empty)

527 """ (empty)

528 ObjectJsonSchema isn't really true since Gemini only accepts a subset of JSON Schema 

529 <https://ai.google.dev/gemini-api/docs/function-calling#function_declarations> 

530 and 

531 <https://ai.google.dev/api/caching#FunctionDeclaration> 

532 """ 

533 

534 

535def _function_from_abstract_tool(tool: ToolDefinition) -> _GeminiFunction: (empty)

536 json_schema = _GeminiJsonSchema(tool.parameters_json_schema).simplify() (empty)

537 f = _GeminiFunction( (empty)

538 name=tool.name, 

539 description=tool.description, 

540 ) 

541 if json_schema.get('properties'): 541 ↛ 543line 541 didn't jump to line 543 because the condition on line 541 was always true(empty)

542 f['parameters'] = json_schema (empty)

543 return f (empty)

544 

545 

546class _GeminiToolConfig(TypedDict): (empty)

547 function_calling_config: _GeminiFunctionCallingConfig (empty)

548 

549 

550def _tool_config(function_names: list[str]) -> _GeminiToolConfig: (empty)

551 return _GeminiToolConfig( (empty)

552 function_calling_config=_GeminiFunctionCallingConfig(mode='ANY', allowed_function_names=function_names) 

553 ) 

554 

555 

556class _GeminiFunctionCallingConfig(TypedDict): (empty)

557 mode: Literal['ANY', 'AUTO'] (empty)

558 allowed_function_names: list[str] (empty)

559 

560 

561@pydantic.with_config(pydantic.ConfigDict(defer_build=True)) (empty)

562class _GeminiResponse(TypedDict): (empty)

563 """Schema for the response from the Gemini API. 

564 

565 See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse> 

566 and <https://cloud.google.com/vertex-ai/docs/reference/rest/v1/GenerateContentResponse> 

567 """ 

568 

569 candidates: list[_GeminiCandidates] (empty)

570 # usageMetadata appears to be required by both APIs but is omitted when streaming responses until the last response 

571 usage_metadata: NotRequired[Annotated[_GeminiUsageMetaData, pydantic.Field(alias='usageMetadata')]] (empty)

572 prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, pydantic.Field(alias='promptFeedback')]] (empty)

573 

574 

575class _GeminiCandidates(TypedDict): (empty)

576 """See <https://ai.google.dev/api/generate-content#v1beta.Candidate>.""" 

577 

578 content: _GeminiContent (empty)

579 finish_reason: NotRequired[Annotated[Literal['STOP', 'MAX_TOKENS'], pydantic.Field(alias='finishReason')]] (empty)

580 """ (empty)

581 See <https://ai.google.dev/api/generate-content#FinishReason>, lots of other values are possible, 

582 but let's wait until we see them and know what they mean to add them here. 

583 """ 

584 avg_log_probs: NotRequired[Annotated[float, pydantic.Field(alias='avgLogProbs')]] (empty)

585 index: NotRequired[int] (empty)

586 safety_ratings: NotRequired[Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')]] (empty)

587 

588 

589class _GeminiUsageMetaData(TypedDict, total=False): (empty)

590 """See <https://ai.google.dev/api/generate-content#FinishReason>. 

591 

592 The docs suggest all fields are required, but some are actually not required, so we assume they are all optional. 

593 """ 

594 

595 prompt_token_count: Annotated[int, pydantic.Field(alias='promptTokenCount')] (empty)

596 candidates_token_count: NotRequired[Annotated[int, pydantic.Field(alias='candidatesTokenCount')]] (empty)

597 total_token_count: Annotated[int, pydantic.Field(alias='totalTokenCount')] (empty)

598 cached_content_token_count: NotRequired[Annotated[int, pydantic.Field(alias='cachedContentTokenCount')]] (empty)

599 

600 

601def _metadata_as_usage(response: _GeminiResponse) -> usage.Usage: (empty)

602 metadata = response.get('usage_metadata') (empty)

603 if metadata is None: 603 ↛ 604line 603 didn't jump to line 604 because the condition on line 603 was never true(empty)

604 return usage.Usage() 

605 details: dict[str, int] = {} (empty)

606 if cached_content_token_count := metadata.get('cached_content_token_count'): 606 ↛ 607line 606 didn't jump to line 607 because the condition on line 606 was never true(empty)

607 details['cached_content_token_count'] = cached_content_token_count 

608 return usage.Usage( (empty)

609 request_tokens=metadata.get('prompt_token_count', 0), 

610 response_tokens=metadata.get('candidates_token_count', 0), 

611 total_tokens=metadata.get('total_token_count', 0), 

612 details=details, 

613 ) 

614 

615 

616class _GeminiSafetyRating(TypedDict): (empty)

617 """See <https://ai.google.dev/gemini-api/docs/safety-settings#safety-filters>.""" 

618 

619 category: Literal[ (empty)

620 'HARM_CATEGORY_HARASSMENT', 

621 'HARM_CATEGORY_HATE_SPEECH', 

622 'HARM_CATEGORY_SEXUALLY_EXPLICIT', 

623 'HARM_CATEGORY_DANGEROUS_CONTENT', 

624 'HARM_CATEGORY_CIVIC_INTEGRITY', 

625 ] 

626 probability: Literal['NEGLIGIBLE', 'LOW', 'MEDIUM', 'HIGH'] (empty)

627 

628 

629class _GeminiPromptFeedback(TypedDict): (empty)

630 """See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse>.""" 

631 

632 block_reason: Annotated[str, pydantic.Field(alias='blockReason')] (empty)

633 safety_ratings: Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')] (empty)

634 

635 

636_gemini_request_ta = pydantic.TypeAdapter(_GeminiRequest) (empty)

637_gemini_response_ta = pydantic.TypeAdapter(_GeminiResponse) (empty)

638 

639# steam requests return a list of https://ai.google.dev/api/generate-content#method:-models.streamgeneratecontent 

640_gemini_streamed_response_ta = pydantic.TypeAdapter(list[_GeminiResponse], config=pydantic.ConfigDict(defer_build=True)) (empty)

641 

642 

643class _GeminiJsonSchema: (empty)

644 """Transforms the JSON Schema from Pydantic to be suitable for Gemini. 

645 

646 Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations) 

647 a subset of OpenAPI v3.0.3. 

648 

649 Specifically: 

650 * gemini doesn't allow the `title` keyword to be set 

651 * gemini doesn't allow `$defs` — we need to inline the definitions where possible 

652 """ 

653 

654 def __init__(self, schema: _utils.ObjectJsonSchema): (empty)

655 self.schema = deepcopy(schema) (empty)

656 self.defs = self.schema.pop('$defs', {}) (empty)

657 

658 def simplify(self) -> dict[str, Any]: (empty)

659 self._simplify(self.schema, refs_stack=()) (empty)

660 return self.schema (empty)

661 

662 def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None: (empty)

663 schema.pop('title', None) (empty)

664 schema.pop('default', None) (empty)

665 if ref := schema.pop('$ref', None): (empty)

666 # noinspection PyTypeChecker 

667 key = re.sub(r'^#/\$defs/', '', ref) (empty)

668 if key in refs_stack: (empty)

669 raise exceptions.UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini') (empty)

670 refs_stack += (key,) (empty)

671 schema_def = self.defs[key] (empty)

672 self._simplify(schema_def, refs_stack) (empty)

673 schema.update(schema_def) (empty)

674 return (empty)

675 

676 if any_of := schema.get('anyOf'): (empty)

677 for item_schema in any_of: (empty)

678 self._simplify(item_schema, refs_stack) (empty)

679 if len(any_of) == 2 and {'type': 'null'} in any_of: 679 ↛ 687line 679 didn't jump to line 687 because the condition on line 679 was always true(empty)

680 for item_schema in any_of: 680 ↛ 687line 680 didn't jump to line 687 because the loop on line 680 didn't complete(empty)

681 if item_schema != {'type': 'null'}: 681 ↛ 680line 681 didn't jump to line 680 because the condition on line 681 was always true(empty)

682 schema.clear() (empty)

683 schema.update(item_schema) (empty)

684 schema['nullable'] = True (empty)

685 return (empty)

686 

687 type_ = schema.get('type') (empty)

688 

689 if type_ == 'object': (empty)

690 self._object(schema, refs_stack) (empty)

691 elif type_ == 'array': (empty)

692 return self._array(schema, refs_stack) (empty)

693 elif type_ == 'string' and (fmt := schema.pop('format', None)): (empty)

694 description = schema.get('description') (empty)

695 if description: (empty)

696 schema['description'] = f'{description} (format: {fmt})' (empty)

697 else: 

698 schema['description'] = f'Format: {fmt}' (empty)

699 

700 def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None: (empty)

701 ad_props = schema.pop('additionalProperties', None) (empty)

702 if ad_props: 702 ↛ 703line 702 didn't jump to line 703 because the condition on line 702 was never true(empty)

703 raise exceptions.UserError('Additional properties in JSON Schema are not supported by Gemini') 

704 

705 if properties := schema.get('properties'): # pragma: no branch (empty)

706 for value in properties.values(): (empty)

707 self._simplify(value, refs_stack) (empty)

708 

709 def _array(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None: (empty)

710 if prefix_items := schema.get('prefixItems'): (empty)

711 # TODO I think this not is supported by Gemini, maybe we should raise an error? 

712 for prefix_item in prefix_items: (empty)

713 self._simplify(prefix_item, refs_stack) (empty)

714 

715 if items_schema := schema.get('items'): # pragma: no branch (empty)

716 self._simplify(items_schema, refs_stack) (empty)