Skip to content

Commit 2f2c757

Browse files
authored
Fix bug with google model safety handling (#2066)
1 parent 10eb5b8 commit 2f2c757

File tree

1 file changed

+11
-6
lines changed

1 file changed

+11
-6
lines changed

pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -374,17 +374,18 @@ async def _build_content_and_config(
374374
def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
375375
if not response.candidates or len(response.candidates) != 1:
376376
raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response') # pragma: no cover
377-
if response.candidates[0].content is None or response.candidates[0].content.parts is None:
378-
if response.candidates[0].finish_reason == 'SAFETY':
377+
candidate = response.candidates[0]
378+
if candidate.content is None or candidate.content.parts is None:
379+
if candidate.finish_reason == 'SAFETY':
379380
raise UnexpectedModelBehavior('Safety settings triggered', str(response))
380381
else:
381382
raise UnexpectedModelBehavior(
382383
'Content field missing from Gemini response', str(response)
383384
) # pragma: no cover
384-
parts = response.candidates[0].content.parts or []
385+
parts = candidate.content.parts or []
385386
vendor_id = response.response_id or None
386387
vendor_details: dict[str, Any] | None = None
387-
finish_reason = response.candidates[0].finish_reason
388+
finish_reason = candidate.finish_reason
388389
if finish_reason: # pragma: no branch
389390
vendor_details = {'finish_reason': finish_reason.value}
390391
usage = _metadata_as_usage(response)
@@ -526,8 +527,12 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
526527

527528
assert chunk.candidates is not None
528529
candidate = chunk.candidates[0]
529-
if candidate.content is None:
530-
raise UnexpectedModelBehavior('Streamed response has no content field') # pragma: no cover
530+
if candidate.content is None or candidate.content.parts is None:
531+
if candidate.finish_reason == 'SAFETY': # pragma: no cover
532+
raise UnexpectedModelBehavior('Safety settings triggered', str(chunk))
533+
else: # pragma: no cover
534+
raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
535+
531536
assert candidate.content.parts is not None
532537
for part in candidate.content.parts:
533538
if part.text is not None:

0 commit comments

Comments
 (0)