add more built-in-tools

Author: Kludex

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -419,7 +419,7 @@ async def stream(
         async for _event in stream:
             pass
 
-    async def _run_stream(
+    async def _run_stream(  # noqa C901
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> AsyncIterator[_messages.HandleResponseEvent]:
         if self._events_iterator is None:
@@ -435,6 +435,10 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                             texts.append(part.content)
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
+                    elif isinstance(part, _messages.ServerToolCallPart):
+                        yield _messages.ServerToolCallEvent(part)
+                    elif isinstance(part, _messages.ServerToolReturnPart):
+                        yield _messages.ServerToolResultEvent(part)
                     else:
                         assert_never(part)
 

pydantic_ai_slim/pydantic_ai/messages.py

@@ -914,6 +914,29 @@ class FunctionToolResultEvent:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@dataclass(repr=False)
+class ServerToolCallEvent:
+    """An event indicating the start to a call to a server tool."""
+
+    part: ServerToolCallPart
+    """The server tool call to make."""
+
+    event_kind: Literal['server_tool_call'] = 'server_tool_call'
+    """Event type identifier, used as a discriminator."""
+
+
+@dataclass(repr=False)
+class ServerToolResultEvent:
+    """An event indicating the result of a server tool call."""
+
+    result: ServerToolReturnPart
+    """The result of the call to the server tool."""
+
+    event_kind: Literal['server_tool_result'] = 'server_tool_result'
+    """Event type identifier, used as a discriminator."""
+
+
 HandleResponseEvent = Annotated[
-    Union[FunctionToolCallEvent, FunctionToolResultEvent], pydantic.Discriminator('event_kind')
+    Union[FunctionToolCallEvent, FunctionToolResultEvent, ServerToolCallEvent, ServerToolResultEvent],
+    pydantic.Discriminator('event_kind'),
 ]

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -8,6 +8,7 @@
 from typing import Any, Literal, Union, cast, overload
 
 from anthropic.types.beta import (
+    BetaCodeExecutionToolResultBlock,
     BetaCodeExecutionToolResultBlockParam,
     BetaServerToolUseBlockParam,
     BetaWebSearchToolResultBlockParam,
@@ -238,6 +239,7 @@ async def _messages_create(
         try:
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
+            extra_headers.setdefault('anthropic-beta', 'code-execution-2025-05-22')
             return await self.client.beta.messages.create(
                 max_tokens=model_settings.get('max_tokens', 1024),
                 system=system_prompt or NOT_GIVEN,
@@ -268,7 +270,7 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
             elif isinstance(item, BetaWebSearchToolResultBlock):
                 items.append(
                     ServerToolReturnPart(
-                        tool_name='web_search',
+                        tool_name=item.type,
                         content=item.content,
                         tool_call_id=item.tool_use_id,
                     )
@@ -282,6 +284,14 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
                         tool_call_id=item.id,
                     )
                 )
+            elif isinstance(item, BetaCodeExecutionToolResultBlock):
+                items.append(
+                    ServerToolReturnPart(
+                        tool_name=item.type,
+                        content=item.content,
+                        tool_call_id=item.tool_use_id,
+                    )
+                )
             else:
                 assert isinstance(item, BetaToolUseBlock), f'unexpected item type {type(item)}'
                 items.append(
@@ -326,7 +336,7 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
                         user_location=user_location,
                     )
                 )
-            if isinstance(tool, CodeExecutionTool):
+            elif isinstance(tool, CodeExecutionTool):
                 tools.append(BetaCodeExecutionTool20250522Param(name='code_execution', type='code_execution_20250522'))
         return tools
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -27,6 +27,8 @@
     ModelResponsePart,
     ModelResponseStreamEvent,
     RetryPromptPart,
+    ServerToolCallPart,
+    ServerToolReturnPart,
     SystemPromptPart,
     TextPart,
     ToolCallPart,
@@ -398,6 +400,9 @@ async def _map_messages(self, messages: list[ModelMessage]) -> list[chat.ChatCom
                         texts.append(item.content)
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
+                    # OpenAI doesn't return server tools calls.
+                    elif isinstance(item, (ServerToolCallPart, ServerToolReturnPart)):
+                        continue
                     else:
                         assert_never(item)
                 message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
@@ -779,6 +784,9 @@ async def _map_messages(
                         openai_messages.append(responses.EasyInputMessageParam(role='assistant', content=item.content))
                     elif isinstance(item, ToolCallPart):
                         openai_messages.append(self._map_tool_call(item))
+                    # OpenAI doesn't return server tools calls.
+                    elif isinstance(item, (ServerToolCallPart, ServerToolReturnPart)):
+                        continue
                     else:
                         assert_never(item)
             else:

tests/models/cassettes/test_anthropic/test_anthropic_code_execution_tool.yaml

@@ -8,7 +8,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '143'
+      - '250'
       content-type:
       - application/json
       host:
@@ -21,38 +21,64 @@ interactions:
         - text: How much is 3 * 12390?
           type: text
         role: user
-      model: claude-3-5-sonnet-latest
+      model: claude-sonnet-4-0
+      stream: false
+      tool_choice:
+        type: auto
+      tools:
+      - name: code_execution
+        type: code_execution_20250522
     uri: https://api.anthropic.com/v1/messages?beta=true
   response:
     headers:
       connection:
       - keep-alive
       content-length:
-      - '370'
+      - '927'
       content-type:
       - application/json
       strict-transport-security:
       - max-age=31536000; includeSubDomains; preload
       transfer-encoding:
       - chunked
     parsed_body:
+      container:
+        expires_at: '2025-05-26T13:30:45.703429+00:00'
+        id: container_011CPW9LpfbF8dmXMvVNCiQJ
       content:
-      - text: |-
-          Let me calculate that:
-
-          3 * 12390 = 37170
+      - text: I'll calculate 3 * 12390 for you.
         type: text
-      id: msg_01D3uRsKuEcst7jtMdwoqYUi
-      model: claude-3-5-sonnet-20241022
+      - id: srvtoolu_01CPfaeVC7ju4VsdzxjSLDrY
+        input:
+          code: |-
+            result = 3 * 12390
+            print(f"3 * 12390 = {result}")
+        name: code_execution
+        type: server_tool_use
+      - content:
+          content: []
+          return_code: 0
+          stderr: ''
+          stdout: |
+            3 * 12390 = 37170
+          type: code_execution_result
+        tool_use_id: srvtoolu_01CPfaeVC7ju4VsdzxjSLDrY
+        type: code_execution_tool_result
+      - text: The answer is **37,170**.
+        type: text
+      id: msg_015H6Emn2T8vZhE52mU2jF1U
+      model: claude-sonnet-4-20250514
       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
+        input_tokens: 1630
+        output_tokens: 105
+        server_tool_use:
+          web_search_requests: 0
         service_tier: standard
     status:
       code: 200