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/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
@@ -61,15 +60,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 +212,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:
@@ -229,22 +235,30 @@ async def _messages_create(
         try:
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
-            return await self.client.messages.create(
-                max_tokens=model_settings.get('max_tokens', 1024),
-                system=system_prompt or NOT_GIVEN,
-                messages=anthropic_messages,
-                model=self._model_name,
-                tools=tools or NOT_GIVEN,
-                tool_choice=tool_choice or NOT_GIVEN,
-                stream=stream,
-                stop_sequences=model_settings.get('stop_sequences', NOT_GIVEN),
-                temperature=model_settings.get('temperature', NOT_GIVEN),
-                top_p=model_settings.get('top_p', NOT_GIVEN),
-                timeout=model_settings.get('timeout', NOT_GIVEN),
-                metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
-                extra_headers=extra_headers,
-                extra_body=model_settings.get('extra_body'),
-            )
+            if beta_tools:
+                return await self.client.beta.messages.create(
+                    max_tokens=model_settings.get('max_tokens', 1024),
+                    system=system_prompt or NOT_GIVEN,
+                    messages=anthropic_messages,
+                    model=self._model_name,
+                )
+            else:
+                return await self.client.messages.create(
+                    max_tokens=model_settings.get('max_tokens', 1024),
+                    system=system_prompt or NOT_GIVEN,
+                    messages=anthropic_messages,
+                    model=self._model_name,
+                    tools=tools or NOT_GIVEN,
+                    tool_choice=tool_choice or NOT_GIVEN,
+                    stream=stream,
+                    stop_sequences=model_settings.get('stop_sequences', NOT_GIVEN),
+                    temperature=model_settings.get('temperature', NOT_GIVEN),
+                    top_p=model_settings.get('top_p', NOT_GIVEN),
+                    timeout=model_settings.get('timeout', NOT_GIVEN),
+                    metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
+                    extra_headers=extra_headers,
+                    extra_body=model_settings.get('extra_body'),
+                )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
@@ -254,10 +268,12 @@ def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
         for item in response.content:
+            breakpoint()
             if isinstance(item, TextBlock):
                 items.append(TextPart(content=item.text))
-            if isinstance(item, WebSearchToolResultBlock):
+            elif isinstance(item, WebSearchToolResultBlock):
                 # TODO(Marcelo): We should send something back to the user, because we need to send it back on the next request.
+                # Do we need a new part type here? All the content is encrypted anyway - but it needs to go back to Anthropic.
                 ...
             else:
                 assert isinstance(item, (ToolUseBlock, ServerToolUseBlock)), f'unexpected item type {type(item)}'
@@ -305,6 +321,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,51 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '257'
+      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
+      stream: false
+      tool_choice:
+        type: auto
+      tools:
+      - name: code_execution
+        type: code_execution_20250522
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '271'
+      content-type:
+      - application/json
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+    parsed_body:
+      error:
+        message: 'tools.0: Input tag ''code_execution_20250522'' found using ''type'' does not match any of the expected tags:
+          ''bash_20250124'', ''custom'', ''text_editor_20250124'', ''text_editor_20250429'', ''web_search_20250305'''
+        type: invalid_request_error
+      type: error
+    status:
+      code: 400
+      message: Bad Request
+version: 1