Add new _GeminiThoughtPart to adhere to Gemini response (#1897)

Author: dnouri

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -431,7 +431,8 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 else:
-                    assert 'function_response' in gemini_part, f'Unexpected part: {gemini_part}'  # pragma: no cover
+                    if not any([key in gemini_part for key in ['function_response', 'thought']]):
+                        raise AssertionError(f'Unexpected part: {gemini_part}')  # pragma: no cover
 
     async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]:
         # This method exists to ensure we only yield completed items, so we don't need to worry about
@@ -605,6 +606,11 @@ class _GeminiFileDataPart(TypedDict):
     file_data: Annotated[_GeminiFileData, pydantic.Field(alias='fileData')]
 
 
+class _GeminiThoughtPart(TypedDict):
+    thought: bool
+    thought_signature: Annotated[str, pydantic.Field(alias='thoughtSignature')]
+
+
 class _GeminiFunctionCallPart(TypedDict):
     function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')]
 
@@ -665,6 +671,8 @@ def _part_discriminator(v: Any) -> str:
             return 'inline_data'  # pragma: no cover
         elif 'fileData' in v:
             return 'file_data'  # pragma: no cover
+        elif 'thought' in v:
+            return 'thought'
         elif 'functionCall' in v or 'function_call' in v:
             return 'function_call'
         elif 'functionResponse' in v or 'function_response' in v:
@@ -682,6 +690,7 @@ def _part_discriminator(v: Any) -> str:
         Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')],
         Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')],
         Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')],
+        Annotated[_GeminiThoughtPart, pydantic.Tag('thought')],
     ],
     pydantic.Discriminator(_part_discriminator),
 ]

tests/models/test_gemini.py

@@ -41,9 +41,13 @@
     _GeminiCandidates,
     _GeminiContent,
     _GeminiFunction,
+    _GeminiFunctionCall,
     _GeminiFunctionCallingConfig,
+    _GeminiFunctionCallPart,
     _GeminiResponse,
     _GeminiSafetyRating,
+    _GeminiTextPart,
+    _GeminiThoughtPart,
     _GeminiToolConfig,
     _GeminiTools,
     _GeminiUsageMetaData,
@@ -898,8 +902,11 @@ async def test_stream_text_heterogeneous(get_gemini_client: GetGeminiClient):
             _GeminiContent(
                 role='model',
                 parts=[
-                    {'text': 'foo'},
-                    {'function_call': {'name': 'get_location', 'args': {'loc_name': 'San Fransisco'}}},
+                    _GeminiThoughtPart(thought=True, thought_signature='test-signature-value'),
+                    _GeminiTextPart(text='foo'),
+                    _GeminiFunctionCallPart(
+                        function_call=_GeminiFunctionCall(name='get_location', args={'loc_name': 'San Fransisco'})
+                    ),
                 ],
             )
         ),
@@ -1327,3 +1334,23 @@ async def test_gemini_no_finish_reason(get_gemini_client: GetGeminiClient):
     for message in result.all_messages():
         if isinstance(message, ModelResponse):
             assert message.vendor_details is None
+
+
+async def test_response_with_thought_part(get_gemini_client: GetGeminiClient):
+    """Tests that a response containing a 'thought' part can be parsed."""
+    content_with_thought = _GeminiContent(
+        role='model',
+        parts=[
+            _GeminiThoughtPart(thought=True, thought_signature='test-signature-value'),
+            _GeminiTextPart(text='Hello from thought test'),
+        ],
+    )
+    response = gemini_response(content_with_thought)
+    gemini_client = get_gemini_client(response)
+    m = GeminiModel('gemini-1.5-flash', provider=GoogleGLAProvider(http_client=gemini_client))
+    agent = Agent(m)
+
+    result = await agent.run('Test with thought')
+
+    assert result.output == 'Hello from thought test'
+    assert result.usage() == snapshot(Usage(requests=1, request_tokens=1, response_tokens=2, total_tokens=3))