pydantic-ai-slim 1.0.8__py3-none-any.whl → 1.0.9__py3-none-any.whl

pydantic_ai/models/google.py

@@ -51,9 +51,12 @@ from . import (
 try:
     from google.genai import Client
     from google.genai.types import (
+        CodeExecutionResult,
+        CodeExecutionResultDict,
         ContentDict,
         ContentUnionDict,
         CountTokensConfigDict,
+        ExecutableCode,
         ExecutableCodeDict,
         FinishReason as GoogleFinishReason,
         FunctionCallDict,
@@ -64,6 +67,7 @@ try:
         GenerateContentResponse,
         GenerationConfigDict,
         GoogleSearchDict,
+        GroundingMetadata,
         HttpOptionsDict,
         MediaResolution,
         Part,
@@ -434,6 +438,7 @@ class GoogleModel(Model):
         usage = _metadata_as_usage(response)
         return _process_response_from_parts(
             parts,
+            candidate.grounding_metadata,
             response.model_version or self._model_name,
             self._provider.name,
             usage,
@@ -569,6 +574,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
+        code_execution_tool_call_id: str | None = None
         async for chunk in self._response:
             self._usage = _metadata_as_usage(chunk)
 
@@ -582,6 +588,19 @@ class GeminiStreamedResponse(StreamedResponse):
                 self.provider_details = {'finish_reason': raw_finish_reason.value}
                 self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
 
+            # Google streams the grounding metadata (including the web search queries and results)
+            # _after_ the text that was generated using it, so it would show up out of order in the stream,
+            # and cause issues with the logic that doesn't consider text ahead of built-in tool calls as output.
+            # If that gets fixed (or we have a workaround), we can uncomment this:
+            # web_search_call, web_search_return = _map_grounding_metadata(
+            #     candidate.grounding_metadata, self.provider_name
+            # )
+            # if web_search_call and web_search_return:
+            #     yield self._parts_manager.handle_builtin_tool_call_part(vendor_part_id=uuid4(), part=web_search_call)
+            #     yield self._parts_manager.handle_builtin_tool_return_part(
+            #         vendor_part_id=uuid4(), part=web_search_return
+            #     )
+
             if candidate.content is None or candidate.content.parts is None:
                 if candidate.finish_reason == 'STOP':  # pragma: no cover
                     # Normal completion - skip this chunk
@@ -590,6 +609,7 @@ class GeminiStreamedResponse(StreamedResponse):
                     raise UnexpectedModelBehavior('Safety settings triggered', str(chunk))
                 else:  # pragma: no cover
                     raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
+
             parts = candidate.content.parts or []
             for part in parts:
                 if part.thought_signature:
@@ -617,9 +637,21 @@ class GeminiStreamedResponse(StreamedResponse):
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 elif part.executable_code is not None:
-                    pass
+                    code_execution_tool_call_id = _utils.generate_tool_call_id()
+                    yield self._parts_manager.handle_builtin_tool_call_part(
+                        vendor_part_id=uuid4(),
+                        part=_map_executable_code(
+                            part.executable_code, self.provider_name, code_execution_tool_call_id
+                        ),
+                    )
                 elif part.code_execution_result is not None:
-                    pass
+                    assert code_execution_tool_call_id is not None
+                    yield self._parts_manager.handle_builtin_tool_return_part(
+                        vendor_part_id=uuid4(),
+                        part=_map_code_execution_result(
+                            part.code_execution_result, self.provider_name, code_execution_tool_call_id
+                        ),
+                    )
                 else:
                     assert part.function_response is not None, f'Unexpected part: {part}'  # pragma: no cover
 
@@ -639,7 +671,7 @@ class GeminiStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict:
+def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict:  # noqa: C901
     parts: list[PartDict] = []
     thought_signature: bytes | None = None
     for item in m.parts:
@@ -663,12 +695,18 @@ def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict
                 part['thought'] = True
         elif isinstance(item, BuiltinToolCallPart):
             if item.provider_name == provider_name:
-                if item.tool_name == 'code_execution':  # pragma: no branch
-                    part['executable_code'] = cast(ExecutableCodeDict, item.args)
+                if item.tool_name == CodeExecutionTool.kind:
+                    part['executable_code'] = cast(ExecutableCodeDict, item.args_as_dict())
+                elif item.tool_name == WebSearchTool.kind:
+                    # Web search calls are not sent back
+                    pass
         elif isinstance(item, BuiltinToolReturnPart):
             if item.provider_name == provider_name:
-                if item.tool_name == 'code_execution':  # pragma: no branch
-                    part['code_execution_result'] = item.content
+                if item.tool_name == CodeExecutionTool.kind and isinstance(item.content, dict):
+                    part['code_execution_result'] = cast(CodeExecutionResultDict, item.content)  # pyright: ignore[reportUnknownMemberType]
+                elif item.tool_name == WebSearchTool.kind:
+                    # Web search results are not sent back
+                    pass
         else:
             assert_never(item)
 
@@ -679,6 +717,7 @@ def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict
 
 def _process_response_from_parts(
     parts: list[Part],
+    grounding_metadata: GroundingMetadata | None,
     model_name: GoogleModelName,
     provider_name: str,
     usage: usage.RequestUsage,
@@ -687,7 +726,17 @@ def _process_response_from_parts(
     finish_reason: FinishReason | None = None,
 ) -> ModelResponse:
     items: list[ModelResponsePart] = []
+
+    # We don't currently turn `candidate.url_context_metadata` into BuiltinToolCallPart and BuiltinToolReturnPart for UrlContextTool.
+    # Please file an issue if you need this.
+
+    web_search_call, web_search_return = _map_grounding_metadata(grounding_metadata, provider_name)
+    if web_search_call and web_search_return:
+        items.append(web_search_call)
+        items.append(web_search_return)
+
     item: ModelResponsePart | None = None
+    code_execution_tool_call_id: str | None = None
     for part in parts:
         if part.thought_signature:
             signature = base64.b64encode(part.thought_signature).decode('utf-8')
@@ -698,16 +747,11 @@ def _process_response_from_parts(
             item.provider_name = provider_name
 
         if part.executable_code is not None:
-            item = BuiltinToolCallPart(
-                provider_name=provider_name, args=part.executable_code.model_dump(), tool_name='code_execution'
-            )
+            code_execution_tool_call_id = _utils.generate_tool_call_id()
+            item = _map_executable_code(part.executable_code, provider_name, code_execution_tool_call_id)
         elif part.code_execution_result is not None:
-            item = BuiltinToolReturnPart(
-                provider_name=provider_name,
-                tool_name='code_execution',
-                content=part.code_execution_result,
-                tool_call_id='not_provided',
-            )
+            assert code_execution_tool_call_id is not None
+            item = _map_code_execution_result(part.code_execution_result, provider_name, code_execution_tool_call_id)
         elif part.text is not None:
             if part.thought:
                 item = ThinkingPart(content=part.text)
@@ -799,3 +843,48 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
         cache_audio_read_tokens=cache_audio_read_tokens,
         details=details,
     )
+
+
+def _map_executable_code(executable_code: ExecutableCode, provider_name: str, tool_call_id: str) -> BuiltinToolCallPart:
+    return BuiltinToolCallPart(
+        provider_name=provider_name,
+        tool_name=CodeExecutionTool.kind,
+        args=executable_code.model_dump(mode='json'),
+        tool_call_id=tool_call_id,
+    )
+
+
+def _map_code_execution_result(
+    code_execution_result: CodeExecutionResult, provider_name: str, tool_call_id: str
+) -> BuiltinToolReturnPart:
+    return BuiltinToolReturnPart(
+        provider_name=provider_name,
+        tool_name=CodeExecutionTool.kind,
+        content=code_execution_result.model_dump(mode='json'),
+        tool_call_id=tool_call_id,
+    )
+
+
+def _map_grounding_metadata(
+    grounding_metadata: GroundingMetadata | None, provider_name: str
+) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart] | tuple[None, None]:
+    if grounding_metadata and (web_search_queries := grounding_metadata.web_search_queries):
+        tool_call_id = _utils.generate_tool_call_id()
+        return (
+            BuiltinToolCallPart(
+                provider_name=provider_name,
+                tool_name=WebSearchTool.kind,
+                tool_call_id=tool_call_id,
+                args={'queries': web_search_queries},
+            ),
+            BuiltinToolReturnPart(
+                provider_name=provider_name,
+                tool_name=WebSearchTool.kind,
+                tool_call_id=tool_call_id,
+                content=[chunk.web.model_dump(mode='json') for chunk in grounding_chunks if chunk.web]
+                if (grounding_chunks := grounding_metadata.grounding_chunks)
+                else None,
+            ),
+        )
+    else:
+        return None, None

pydantic_ai/models/groq.py

@@ -8,11 +8,11 @@ from datetime import datetime
 from typing import Any, Literal, cast, overload
 
 from pydantic import BaseModel, Json, ValidationError
+from pydantic_core import from_json
 from typing_extensions import assert_never
 
-from pydantic_ai._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
-
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
+from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
 from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import generate_tool_call_id, guard_tool_call_id as _guard_tool_call_id, number_to_datetime
@@ -55,6 +55,7 @@ try:
     from groq import NOT_GIVEN, APIError, APIStatusError, AsyncGroq, AsyncStream
     from groq.types import chat
     from groq.types.chat.chat_completion_content_part_image_param import ImageURL
+    from groq.types.chat.chat_completion_message import ExecutedTool
 except ImportError as _import_error:
     raise ImportError(
         'Please install `groq` to use the Groq model, '
@@ -308,22 +309,15 @@ class GroqModel(Model):
         timestamp = number_to_datetime(response.created)
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
-        if choice.message.executed_tools:
-            for tool in choice.message.executed_tools:
-                tool_call_id = generate_tool_call_id()
-                items.append(
-                    BuiltinToolCallPart(
-                        tool_name=tool.type, args=tool.arguments, provider_name=self.system, tool_call_id=tool_call_id
-                    )
-                )
-                items.append(
-                    BuiltinToolReturnPart(
-                        provider_name=self.system, tool_name=tool.type, content=tool.output, tool_call_id=tool_call_id
-                    )
-                )
         if choice.message.reasoning is not None:
             # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`.
             items.append(ThinkingPart(content=choice.message.reasoning))
+        if choice.message.executed_tools:
+            for tool in choice.message.executed_tools:
+                call_part, return_part = _map_executed_tool(tool, self.system)
+                if call_part and return_part:  # pragma: no branch
+                    items.append(call_part)
+                    items.append(return_part)
         if choice.message.content is not None:
             # NOTE: The `<think>` tag is only present if `groq_reasoning_format` is set to `raw`.
             items.extend(split_content_into_text_and_thinking(choice.message.content, self.profile.thinking_tags))
@@ -400,7 +394,7 @@ class GroqModel(Model):
                         start_tag, end_tag = self.profile.thinking_tags
                         texts.append('\n'.join([start_tag, item.content, end_tag]))
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
-                        # This is currently never returned from groq
+                        # These are not currently sent back
                         pass
                     else:
                         assert_never(item)
@@ -513,8 +507,9 @@ class GroqStreamedResponse(StreamedResponse):
     _timestamp: datetime
     _provider_name: str
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         try:
+            executed_tool_call_id: str | None = None
             async for chunk in self._response:
                 self._usage += _map_usage(chunk)
 
@@ -530,6 +525,28 @@ class GroqStreamedResponse(StreamedResponse):
                     self.provider_details = {'finish_reason': raw_finish_reason}
                     self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
 
+                if choice.delta.reasoning is not None:
+                    # NOTE: The `reasoning` field is only present if `groq_reasoning_format` is set to `parsed`.
+                    yield self._parts_manager.handle_thinking_delta(
+                        vendor_part_id='reasoning', content=choice.delta.reasoning
+                    )
+
+                if choice.delta.executed_tools:
+                    for tool in choice.delta.executed_tools:
+                        call_part, return_part = _map_executed_tool(
+                            tool, self.provider_name, streaming=True, tool_call_id=executed_tool_call_id
+                        )
+                        if call_part:
+                            executed_tool_call_id = call_part.tool_call_id
+                            yield self._parts_manager.handle_builtin_tool_call_part(
+                                vendor_part_id=f'executed_tools-{tool.index}-call', part=call_part
+                            )
+                        if return_part:
+                            executed_tool_call_id = None
+                            yield self._parts_manager.handle_builtin_tool_return_part(
+                                vendor_part_id=f'executed_tools-{tool.index}-return', part=return_part
+                            )
+
                 # Handle the text part of the response
                 content = choice.delta.content
                 if content is not None:
@@ -626,3 +643,37 @@ class _GroqToolUseFailedError(BaseModel):
     # }
 
     error: _GroqToolUseFailedInnerError
+
+
+def _map_executed_tool(
+    tool: ExecutedTool, provider_name: str, streaming: bool = False, tool_call_id: str | None = None
+) -> tuple[BuiltinToolCallPart | None, BuiltinToolReturnPart | None]:
+    if tool.type == 'search':
+        if tool.search_results and (tool.search_results.images or tool.search_results.results):
+            results = tool.search_results.model_dump(mode='json')
+        else:
+            results = tool.output
+
+        tool_call_id = tool_call_id or generate_tool_call_id()
+        call_part = BuiltinToolCallPart(
+            tool_name=WebSearchTool.kind,
+            args=from_json(tool.arguments),
+            provider_name=provider_name,
+            tool_call_id=tool_call_id,
+        )
+        return_part = BuiltinToolReturnPart(
+            tool_name=WebSearchTool.kind,
+            content=results,
+            provider_name=provider_name,
+            tool_call_id=tool_call_id,
+        )
+
+        if streaming:
+            if results:
+                return None, return_part
+            else:
+                return call_part, None
+        else:
+            return call_part, return_part
+    else:  # pragma: no cover
+        return None, None