Add more work on it

Author: Kludex

pydantic_ai_slim/pydantic_ai/builtin_tools.py

@@ -14,6 +14,8 @@ class AbstractBuiltinTool(ABC):
     """A builtin tool that can be used by an agent.
 
     This class is abstract and cannot be instantiated directly.
+
+    The builtin tools are are passed to the model as part of the `ModelRequestParameters`.
     """
 
 
@@ -73,3 +75,10 @@ class WebSearchTool(AbstractBuiltinTool):
     * Anthropic (https://docs.anthropic.com/en/docs/build-with-claude/tool-use/web-search-tool#domain-filtering)
     * Groq (https://console.groq.com/docs/agentic-tooling#search-settings)
     """
+
+    max_uses: int | None = None
+    """If provided, the tool will stop searching the web after the given number of uses.
+
+    Supported by:
+    * Anthropic
+    """

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -8,8 +8,12 @@
 from json import JSONDecodeError, loads as json_loads
 from typing import Any, Literal, Union, cast, overload
 
+from anthropic.types import ServerToolUseBlock, ToolUnionParam, WebSearchTool20250305Param, WebSearchToolResultBlock
+from anthropic.types.web_search_tool_20250305_param import UserLocation
 from typing_extensions import assert_never
 
+from pydantic_ai.builtin_tools import WebSearchTool
+
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
@@ -207,6 +211,7 @@ async def _messages_create(
     ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]:
         # standalone function to make it easier to override
         tools = self._get_tools(model_request_parameters)
+        tools += self._get_builtin_tools(model_request_parameters)
         tool_choice: ToolChoiceParam | None
 
         if not tools:
@@ -252,8 +257,11 @@ def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         for item in response.content:
             if isinstance(item, TextBlock):
                 items.append(TextPart(content=item.text))
+            if isinstance(item, WebSearchToolResultBlock):
+                # TODO(Marcelo): We should send something back to the user, because we need to send it back on the next request.
+                ...
             else:
-                assert isinstance(item, ToolUseBlock), 'unexpected item type'
+                assert isinstance(item, (ToolUseBlock, ServerToolUseBlock)), f'unexpected item type {type(item)}'
                 items.append(
                     ToolCallPart(
                         tool_name=item.name,
@@ -282,6 +290,22 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[T
             tools += [self._map_tool_definition(r) for r in model_request_parameters.output_tools]
         return tools
 
+    def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolUnionParam]:
+        tools: list[ToolUnionParam] = []
+        for tool in model_request_parameters.builtin_tools:
+            if isinstance(tool, WebSearchTool):
+                user_location = UserLocation(type='approximate', **tool.user_location) if tool.user_location else None
+                tools.append(
+                    WebSearchTool20250305Param(
+                        name='web_search',
+                        type='web_search_20250305',
+                        allowed_domains=tool.allowed_domains,
+                        blocked_domains=tool.blocked_domains,
+                        user_location=user_location,
+                    )
+                )
+        return tools
+
     async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
         """Just maps a `pydantic_ai.Message` to a `anthropic.types.MessageParam`."""
         system_prompt_parts: list[str] = []

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -173,6 +173,7 @@ def _customize_tool_def(t: ToolDefinition):
 
         return ModelRequestParameters(
             function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
+            builtin_tools=model_request_parameters.builtin_tools,
             allow_text_output=model_request_parameters.allow_text_output,
             output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
         )

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -234,6 +234,7 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
+        # TODO(Marcelo): The `choice.message.executed_tools` has the server tools executed by Groq.
         if choice.message.content is not None:
             items.append(TextPart(content=choice.message.content))
         if choice.message.tool_calls is not None:

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -278,7 +278,6 @@ async def _completions_create(
             return await self.client.chat.completions.create(
                 model=self._model_name,
                 messages=openai_messages,
-                n=1,
                 parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
                 tools=tools or NOT_GIVEN,
                 tool_choice=tool_choice or NOT_GIVEN,
@@ -338,13 +337,15 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[c
     def _get_web_search_options(self, model_request_parameters: ModelRequestParameters) -> WebSearchOptions | None:
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
-                return WebSearchOptions(
-                    search_context_size=tool.search_context_size,
-                    user_location=WebSearchOptionsUserLocation(
-                        type='approximate',
-                        approximate=WebSearchOptionsUserLocationApproximate(**tool.user_location),
-                    ),
-                )
+                if tool.user_location:
+                    return WebSearchOptions(
+                        search_context_size=tool.search_context_size,
+                        user_location=WebSearchOptionsUserLocation(
+                            type='approximate',
+                            approximate=WebSearchOptionsUserLocationApproximate(**tool.user_location),
+                        ),
+                    )
+                return WebSearchOptions(search_context_size=tool.search_context_size)
 
     async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
@@ -677,13 +678,14 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
         tools: list[responses.ToolParam] = []
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
-                tools.append(
-                    responses.WebSearchToolParam(
-                        type='web_search_preview',
-                        search_context_size=tool.search_context_size,
-                        user_location={'type': 'approximate', **tool.user_location},
-                    )
+                web_search_tool = responses.WebSearchToolParam(
+                    type='web_search_preview', search_context_size=tool.search_context_size
                 )
+                if tool.user_location:
+                    web_search_tool['user_location'] = responses.web_search_tool_param.UserLocation(
+                        type='approximate', **tool.user_location
+                    )
+                tools.append(web_search_tool)
         return tools
 
     @staticmethod
@@ -1112,6 +1114,7 @@ def _customize_tool_def(t: ToolDefinition):
 
     return ModelRequestParameters(
         function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
+        builtin_tools=model_request_parameters.builtin_tools,
         allow_text_output=model_request_parameters.allow_text_output,
         output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
     )