diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index 8aca93c720..c78e92cc4c 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -710,6 +710,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: for part in parts: provider_details: dict[str, Any] | None = None + thought_signature: str | None = None if part.thought_signature: # Per https://ai.google.dev/gemini-api/docs/function-calling?example=meeting#thought-signatures: # - Always send the thought_signature back to the model inside its original Part. @@ -718,12 +719,28 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: thought_signature = base64.b64encode(part.thought_signature).decode('utf-8') provider_details = {'thought_signature': thought_signature} + # Google returns thought_signature on the part FOLLOWING the thinking part. + # Apply it to the previous ThinkingPart if this is a non-thinking part with a signature. + if thought_signature and not part.thought: + # Only apply signature if the latest part is a ThinkingPart + parts = self._parts_manager.get_parts() + if parts and isinstance(parts[-1], ThinkingPart): + for event in self._parts_manager.handle_thinking_delta( + vendor_part_id=None, + signature=thought_signature, + ): + yield event + if part.text is not None: if len(part.text) == 0 and not provider_details: continue if part.thought: for event in self._parts_manager.handle_thinking_delta( - vendor_part_id=None, content=part.text, provider_details=provider_details + vendor_part_id=None, + content=part.text, + signature=thought_signature, + provider_name=self._provider_name, + provider_details=provider_details, ): yield event else: @@ -878,8 +895,10 @@ def _process_response_from_parts( item: ModelResponsePart | None = None code_execution_tool_call_id: str | None = None + last_thinking_part: ThinkingPart | None = None for part in parts: provider_details: dict[str, Any] | None = None + thought_signature: str | None = None if part.thought_signature: # Per https://ai.google.dev/gemini-api/docs/function-calling?example=meeting#thought-signatures: # - Always send the thought_signature back to the model inside its original Part. @@ -899,7 +918,8 @@ def _process_response_from_parts( if len(part.text) == 0 and not provider_details: continue if part.thought: - item = ThinkingPart(content=part.text) + item = ThinkingPart(content=part.text, signature=thought_signature, provider_name=provider_name) + last_thinking_part = item else: item = TextPart(content=part.text) elif part.function_call: @@ -916,6 +936,12 @@ def _process_response_from_parts( else: # pragma: no cover raise UnexpectedModelBehavior(f'Unsupported response from Gemini: {part!r}') + # Google returns thought_signature on the part FOLLOWING the thinking part. + # Apply it to the previous ThinkingPart if this is a non-thinking part with a signature. + if thought_signature and last_thinking_part and not part.thought: + last_thinking_part.signature = thought_signature + last_thinking_part = None # Only apply once + if provider_details: item.provider_details = {**(item.provider_details or {}), **provider_details} diff --git a/tests/models/test_google.py b/tests/models/test_google.py index 3ef8cd5dda..0ba1025b37 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -1931,7 +1931,11 @@ def dummy() -> None: ... # pragma: no cover ), ModelResponse( parts=[ - ThinkingPart(content=IsStr()), + ThinkingPart( + content=IsStr(), + signature=IsStr(), + provider_name='google-gla', + ), TextPart( content=IsStr(), provider_details={'thought_signature': IsStr()}, @@ -1968,7 +1972,11 @@ def dummy() -> None: ... # pragma: no cover ), ModelResponse( parts=[ - ThinkingPart(content=IsStr()), + ThinkingPart( + content=IsStr(), + signature=IsStr(), + provider_name='google-gla', + ), TextPart( content=IsStr(), provider_details={'thought_signature': IsStr()}, @@ -2077,7 +2085,7 @@ def dummy() -> None: ... # pragma: no cover ), ModelResponse( parts=[ - ThinkingPart(content=IsStr()), + ThinkingPart(content=IsStr(), signature=IsStr(), provider_name='google-gla'), TextPart( content=IsStr(), provider_details={'thought_signature': IsStr()}, @@ -2133,7 +2141,7 @@ def dummy() -> None: ... # pragma: no cover ), ModelResponse( parts=[ - ThinkingPart(content=IsStr()), + ThinkingPart(content=IsStr(), signature=IsStr(), provider_name='google-gla'), TextPart( content=IsStr(), provider_details={'thought_signature': IsStr()}, @@ -2155,10 +2163,11 @@ def dummy() -> None: ... # pragma: no cover assert event_parts == snapshot( [ - PartStartEvent(index=0, part=ThinkingPart(content=IsStr())), - PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr())), - PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr())), - PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr())), + PartStartEvent(index=0, part=ThinkingPart(content=IsStr(), provider_name='google-gla')), + PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr(), provider_name='google-gla')), + PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr(), provider_name='google-gla')), + PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=IsStr(), provider_name='google-gla')), + PartDeltaEvent(index=0, delta=ThinkingPartDelta(signature_delta=IsStr())), PartEndEvent( index=0, part=ThinkingPart( @@ -2183,7 +2192,9 @@ def dummy() -> None: ... # pragma: no cover I've identified the core user intent: to learn safe street-crossing. Now, I'm focusing on crafting universally applicable steps. Finding safe crossing locations and looking-listening for traffic remain paramount. I'm prioritizing direct, clear language, addressing my limitations as an AI. I'm crafting advice that works generally, regardless of specific circumstances or locations. -""" +""", + signature=IsStr(), + provider_name='google-gla', ), next_part_kind='text', ),