Add builtin tools

Author: Kludex

pydantic_ai_slim/pydantic_ai/builtin_tools.py

@@ -1,7 +1,7 @@
 from __future__ import annotations as _annotations
 
 from abc import ABC
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Literal
 
 from typing_extensions import TypedDict
@@ -19,20 +19,6 @@ class AbstractBuiltinTool(ABC):
     """
 
 
-class UserLocation(TypedDict, total=False):
-    """Allows you to localize search results based on a user's location.
-
-    Supported by:
-    * Anthropic
-    * OpenAI
-    """
-
-    city: str
-    country: str
-    region: str
-    timezone: str
-
-
 @dataclass
 class WebSearchTool(AbstractBuiltinTool):
     """A builtin tool that allows your agent to search the web for information.
@@ -47,7 +33,7 @@ class WebSearchTool(AbstractBuiltinTool):
     * OpenAI
     """
 
-    user_location: UserLocation = field(default_factory=UserLocation)
+    user_location: UserLocation | None = None
     """The `user_location` parameter allows you to localize search results based on a user's location.
 
     Supported by:
@@ -82,3 +68,26 @@ class WebSearchTool(AbstractBuiltinTool):
     Supported by:
     * Anthropic
     """
+
+
+class UserLocation(TypedDict, total=False):
+    """Allows you to localize search results based on a user's location.
+
+    Supported by:
+    * Anthropic
+    * OpenAI
+    """
+
+    city: str
+    country: str
+    region: str
+    timezone: str
+
+
+class CodeExecutionTool(AbstractBuiltinTool):
+    """A builtin tool that allows your agent to execute code.
+
+    Supported by:
+    * Anthropic
+    * OpenAI
+    """

pydantic_ai_slim/pydantic_ai/messages.py

@@ -349,8 +349,8 @@ def otel_event(self, settings: InstrumentationSettings) -> Event:
 
 
 @dataclass(repr=False)
-class ToolReturnPart:
-    """A tool return message, this encodes the result of running a tool."""
+class BaseToolReturnPart:
+    """Base class for tool return parts."""
 
     tool_name: str
     """The name of the "tool" was called."""
@@ -364,9 +364,6 @@ class ToolReturnPart:
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp, when the tool returned."""
 
-    part_kind: Literal['tool-return'] = 'tool-return'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
     def model_response_str(self) -> str:
         """Return a string representation of the content for the model."""
         if isinstance(self.content, str):
@@ -391,6 +388,22 @@ def otel_event(self, _settings: InstrumentationSettings) -> Event:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@dataclass(repr=False)
+class ToolReturnPart(BaseToolReturnPart):
+    """A tool return message, this encodes the result of running a tool."""
+
+    part_kind: Literal['tool-return'] = 'tool-return'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+@dataclass(repr=False)
+class ServerToolReturnPart(BaseToolReturnPart):
+    """A tool return message from a server tool."""
+
+    part_kind: Literal['server-tool-return'] = 'server-tool-return'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
 
 
@@ -503,7 +516,7 @@ def has_content(self) -> bool:
 
 
 @dataclass(repr=False)
-class ToolCallPart:
+class BaseToolCallPart:
     """A tool call from a model."""
 
     tool_name: str
@@ -521,9 +534,6 @@ class ToolCallPart:
     In case the tool call id is not provided by the model, PydanticAI will generate a random one.
     """
 
-    part_kind: Literal['tool-call'] = 'tool-call'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
     def args_as_dict(self) -> dict[str, Any]:
         """Return the arguments as a Python dictionary.
 
@@ -560,7 +570,28 @@ def has_content(self) -> bool:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
+@dataclass(repr=False)
+class ToolCallPart(BaseToolCallPart):
+    """A tool call from a model."""
+
+    part_kind: Literal['tool-call'] = 'tool-call'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+@dataclass(repr=False)
+class ServerToolCallPart(BaseToolCallPart):
+    """A tool call from a server tool."""
+
+    model_name: str | None = None
+    """The name of the model that generated the response."""
+
+    part_kind: Literal['server-tool-call'] = 'server-tool-call'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+ModelResponsePart = Annotated[
+    Union[TextPart, ToolCallPart, ServerToolCallPart, ServerToolReturnPart], pydantic.Discriminator('part_kind')
+]
 """A message part returned by a model."""
 
 

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -7,11 +7,10 @@
 from datetime import datetime, timezone
 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 anthropic.types.beta import BetaMessage, BetaRawMessageStreamEvent, BetaToolUnionParam
 from typing_extensions import assert_never
 
-from pydantic_ai.builtin_tools import WebSearchTool
+from pydantic_ai.builtin_tools import CodeExecutionTool, WebSearchTool
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
@@ -25,6 +24,8 @@
     ModelResponsePart,
     ModelResponseStreamEvent,
     RetryPromptPart,
+    ServerToolCallPart,
+    ServerToolReturnPart,
     SystemPromptPart,
     TextPart,
     ToolCallPart,
@@ -61,15 +62,21 @@
         RawMessageStartEvent,
         RawMessageStopEvent,
         RawMessageStreamEvent,
+        ServerToolUseBlock,
         TextBlock,
         TextBlockParam,
         TextDelta,
         ToolChoiceParam,
         ToolParam,
         ToolResultBlockParam,
+        ToolUnionParam,
         ToolUseBlock,
         ToolUseBlockParam,
+        WebSearchTool20250305Param,
+        WebSearchToolResultBlock,
     )
+    from anthropic.types.beta.beta_code_execution_tool_20250522_param import BetaCodeExecutionTool20250522Param
+    from anthropic.types.web_search_tool_20250305_param import UserLocation
 except ImportError as _import_error:
     raise ImportError(
         'Please install `anthropic` to use the Anthropic model, '
@@ -207,10 +214,11 @@ async def _messages_create(
         stream: bool,
         model_settings: AnthropicModelSettings,
         model_request_parameters: ModelRequestParameters,
-    ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]:
+    ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent] | BetaMessage | AsyncStream[BetaRawMessageStreamEvent]:
         # standalone function to make it easier to override
         tools = self._get_tools(model_request_parameters)
         tools += self._get_builtin_tools(model_request_parameters)
+        beta_tools = self._get_beta_tools(model_request_parameters)
         tool_choice: ToolChoiceParam | None
 
         if not tools:
@@ -256,11 +264,25 @@ 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.
-                ...
+            elif isinstance(item, WebSearchToolResultBlock):
+                items.append(
+                    ServerToolReturnPart(
+                        tool_name='web_search',
+                        content=item.content,
+                        tool_call_id=item.tool_use_id,
+                    )
+                )
+            elif isinstance(item, ServerToolUseBlock):
+                items.append(
+                    ServerToolCallPart(
+                        model_name='anthropic',
+                        tool_name=item.name,
+                        args=cast(dict[str, Any], item.input),
+                        tool_call_id=item.id,
+                    )
+                )
             else:
-                assert isinstance(item, (ToolUseBlock, ServerToolUseBlock)), f'unexpected item type {type(item)}'
+                assert isinstance(item, ToolUseBlock), f'unexpected item type {type(item)}'
                 items.append(
                     ToolCallPart(
                         tool_name=item.name,
@@ -305,6 +327,13 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
                 )
         return tools
 
+    def _get_beta_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolUnionParam]:
+        tools: list[BetaToolUnionParam] = []
+        for tool in model_request_parameters.builtin_tools:
+            if isinstance(tool, CodeExecutionTool):
+                tools.append(BetaCodeExecutionTool20250522Param(name='code_execution', type='code_execution_20250522'))
+        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/pyproject.toml

@@ -65,7 +65,7 @@ openai = ["openai>=1.75.0"]
 cohere = ["cohere>=5.13.11; platform_system != 'Emscripten'"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
 google = ["google-genai>=1.15.0"]
-anthropic = ["anthropic>=0.49.0"]
+anthropic = ["anthropic>=0.52.0"]
 groq = ["groq>=0.15.0"]
 mistral = ["mistralai>=1.2.5"]
 bedrock = ["boto3>=1.35.74"]

tests/models/cassettes/test_anthropic/test_anthropic_code_execution_tool.yaml

@@ -0,0 +1,60 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '143'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      max_tokens: 1024
+      messages:
+      - content:
+        - text: How much is 3 * 12390?
+          type: text
+        role: user
+      model: claude-3-5-sonnet-latest
+    uri: https://api.anthropic.com/v1/messages?beta=true
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '370'
+      content-type:
+      - application/json
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      content:
+      - text: |-
+          Let me calculate that:
+
+          3 * 12390 = 37170
+        type: text
+      id: msg_01D3uRsKuEcst7jtMdwoqYUi
+      model: claude-3-5-sonnet-20241022
+      role: assistant
+      stop_reason: end_turn
+      stop_sequence: null
+      type: message
+      usage:
+        cache_creation_input_tokens: 0
+        cache_read_input_tokens: 0
+        input_tokens: 18
+        output_tokens: 20
+        service_tier: standard
+    status:
+      code: 200
+      message: OK
+version: 1