Merge branch 'main' into multimodal-tool-output

Author: DouweM

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -291,13 +291,4 @@ async def __anext__(self) -> T:
 
 
 def get_traceparent(x: AgentRun | AgentRunResult | GraphRun | GraphRunResult) -> str:
-    import logfire
-    import logfire_api
-    from logfire.experimental.annotations import get_traceparent
-
-    span: AbstractSpan | None = x._span(required=False)  # type: ignore[reportPrivateUsage]
-    if not span:  # pragma: no cover
-        return ''
-    if isinstance(span, logfire_api.LogfireSpan):  # pragma: no cover
-        assert isinstance(span, logfire.LogfireSpan)
-    return get_traceparent(span)
+    return x._traceparent(required=False) or ''  # type: ignore[reportPrivateUsage]

pydantic_ai_slim/pydantic_ai/agent.py

@@ -27,7 +27,6 @@
     result,
     usage as _usage,
 )
-from ._utils import AbstractSpan
 from .models.instrumented import InstrumentationSettings, InstrumentedModel
 from .result import FinalResult, OutputDataT, StreamedRunResult, ToolOutput
 from .settings import ModelSettings, merge_model_settings
@@ -1683,14 +1682,14 @@ async def main():
     ]
 
     @overload
-    def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:
-        span = self._graph_run._span(required=False)  # type: ignore[reportPrivateUsage]
-        if span is None and required:  # pragma: no cover
-            raise AttributeError('Span is not available for this agent run')
-        return span
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        traceparent = self._graph_run._traceparent(required=False)  # type: ignore[reportPrivateUsage]
+        if traceparent is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return traceparent
 
     @property
     def ctx(self) -> GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any]]:
@@ -1729,7 +1728,7 @@ def result(self) -> AgentRunResult[OutputDataT] | None:
             graph_run_result.output.tool_name,
             graph_run_result.state,
             self._graph_run.deps.new_message_index,
-            self._graph_run._span(required=False),  # type: ignore[reportPrivateUsage]
+            self._traceparent(required=False),
         )
 
     def __aiter__(
@@ -1847,16 +1846,16 @@ class AgentRunResult(Generic[OutputDataT]):
     _output_tool_name: str | None = dataclasses.field(repr=False)
     _state: _agent_graph.GraphAgentState = dataclasses.field(repr=False)
     _new_message_index: int = dataclasses.field(repr=False)
-    _span_value: AbstractSpan | None = dataclasses.field(repr=False)
+    _traceparent_value: str | None = dataclasses.field(repr=False)
 
     @overload
-    def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
     @overload
-    def _span(self) -> AbstractSpan: ...
-    def _span(self, *, required: bool = True) -> AbstractSpan | None:
-        if self._span_value is None and required:  # pragma: no cover
-            raise AttributeError('Span is not available for this agent run')
-        return self._span_value
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        if self._traceparent_value is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return self._traceparent_value
 
     @property
     @deprecated('`result.data` is deprecated, use `result.output` instead.')

pydantic_ai_slim/pydantic_ai/common_tools/duckduckgo.py

@@ -54,8 +54,6 @@ async def __call__(self, query: str) -> list[DuckDuckGoResult]:
         """
         search = functools.partial(self.client.text, max_results=self.max_results)
         results = await anyio.to_thread.run_sync(search, query)
-        if len(results) == 0:
-            raise RuntimeError('No search results found.')
         return duckduckgo_ta.validate_python(results)
 
 

pydantic_ai_slim/pydantic_ai/common_tools/tavily.py

@@ -63,8 +63,6 @@ async def __call__(
             The search results.
         """
         results = await self.client.search(query, search_depth=search_deep, topic=topic, time_range=time_range)  # type: ignore[reportUnknownMemberType]
-        if not results['results']:
-            raise RuntimeError('No search results found.')
         return tavily_search_ta.validate_python(results['results'])
 
 

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -9,7 +9,7 @@
 from typing import Any
 
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from mcp.types import JSONRPCMessage
+from mcp.types import JSONRPCMessage, LoggingLevel
 from typing_extensions import Self
 
 from pydantic_ai.tools import ToolDefinition
@@ -52,6 +52,11 @@ async def client_streams(
         raise NotImplementedError('MCP Server subclasses must implement this method.')
         yield
 
+    @abstractmethod
+    def _get_log_level(self) -> LoggingLevel | None:
+        """Get the log level for the MCP server."""
+        raise NotImplementedError('MCP Server subclasses must implement this method.')
+
     async def list_tools(self) -> list[ToolDefinition]:
         """Retrieve tools that are currently active on the server.
 
@@ -89,6 +94,8 @@ async def __aenter__(self) -> Self:
         self._client = await self._exit_stack.enter_async_context(client)
 
         await self._client.initialize()
+        if log_level := self._get_log_level():
+            await self._client.set_logging_level(log_level)
         self.is_running = True
         return self
 
@@ -150,6 +157,13 @@ async def main():
     By default the subprocess will not inherit any environment variables from the parent process.
     If you want to inherit the environment variables from the parent process, use `env=os.environ`.
     """
+    log_level: LoggingLevel | None = None
+    """The log level to set when connecting to the server, if any.
+
+    See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
+
+    If `None`, no log level will be set.
+    """
 
     cwd: str | Path | None = None
     """The working directory to use when spawning the process."""
@@ -164,6 +178,9 @@ async def client_streams(
         async with stdio_client(server=server) as (read_stream, write_stream):
             yield read_stream, write_stream
 
+    def _get_log_level(self) -> LoggingLevel | None:
+        return self.log_level
+
 
 @dataclass
 class MCPServerHTTP(MCPServer):
@@ -223,6 +240,13 @@ async def main():
     If no new messages are received within this time, the connection will be considered stale
     and may be closed. Defaults to 5 minutes (300 seconds).
     """
+    log_level: LoggingLevel | None = None
+    """The log level to set when connecting to the server, if any.
+
+    See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
+
+    If `None`, no log level will be set.
+    """
 
     @asynccontextmanager
     async def client_streams(
@@ -234,3 +258,6 @@ async def client_streams(
             url=self.url, headers=self.headers, timeout=self.timeout, sse_read_timeout=self.sse_read_timeout
         ) as (read_stream, write_stream):
             yield read_stream, write_stream
+
+    def _get_log_level(self) -> LoggingLevel | None:
+        return self.log_level

pydantic_ai_slim/pydantic_ai/messages.py

@@ -516,6 +516,8 @@ def args_as_dict(self) -> dict[str, Any]:
         """
         if isinstance(self.args, dict):
             return self.args
+        if isinstance(self.args, str) and not self.args:
+            return {}
         args = pydantic_core.from_json(self.args)
         assert isinstance(args, dict), 'args should be a dict'
         return cast(dict[str, Any], args)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -106,6 +106,7 @@
         'google-gla:gemini-2.0-flash',
         'google-gla:gemini-2.0-flash-lite-preview-02-05',
         'google-gla:gemini-2.0-pro-exp-02-05',
+        'google-gla:gemini-2.5-flash-preview-04-17',
         'google-gla:gemini-2.5-pro-exp-03-25',
         'google-gla:gemini-2.5-pro-preview-03-25',
         'google-vertex:gemini-1.0-pro',
@@ -118,6 +119,7 @@
         'google-vertex:gemini-2.0-flash',
         'google-vertex:gemini-2.0-flash-lite-preview-02-05',
         'google-vertex:gemini-2.0-pro-exp-02-05',
+        'google-vertex:gemini-2.5-flash-preview-04-17',
         'google-vertex:gemini-2.5-pro-exp-03-25',
         'google-vertex:gemini-2.5-pro-preview-03-25',
         'gpt-3.5-turbo',
@@ -192,6 +194,8 @@
         'o1-mini-2024-09-12',
         'o1-preview',
         'o1-preview-2024-09-12',
+        'o3',
+        'o3-2025-04-16',
         'o3-mini',
         'o3-mini-2025-01-31',
         'openai:chatgpt-4o-latest',
@@ -241,8 +245,12 @@
         'openai:o1-mini-2024-09-12',
         'openai:o1-preview',
         'openai:o1-preview-2024-09-12',
+        'openai:o3',
+        'openai:o3-2025-04-16',
         'openai:o3-mini',
         'openai:o3-mini-2025-01-31',
+        'openai:o4-mini',
+        'openai:o4-mini-2025-04-16',
         'test',
     ],
 )

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -239,6 +239,7 @@ async def _messages_create(
                 timeout=model_settings.get('timeout', NOT_GIVEN),
                 metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
                 extra_headers={'User-Agent': get_user_agent()},
+                extra_body=model_settings.get('extra_body'),
             )
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -2,10 +2,11 @@
 
 import functools
 import typing
-from collections.abc import AsyncIterator, Iterable, Mapping
+from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
+from itertools import count
 from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast, overload
 
 import anyio
@@ -369,13 +370,14 @@ async def _map_messages(
         """Just maps a `pydantic_ai.Message` to the Bedrock `MessageUnionTypeDef`."""
         system_prompt: list[SystemContentBlockTypeDef] = []
         bedrock_messages: list[MessageUnionTypeDef] = []
+        document_count: Iterator[int] = count(1)
         for m in messages:
             if isinstance(m, ModelRequest):
                 for part in m.parts:
                     if isinstance(part, SystemPromptPart):
                         system_prompt.append({'text': part.content})
                     elif isinstance(part, UserPromptPart):
-                        bedrock_messages.extend(await self._map_user_prompt(part))
+                        bedrock_messages.extend(await self._map_user_prompt(part, document_count))
                     elif isinstance(part, ToolReturnPart):
                         assert part.tool_call_id is not None
                         bedrock_messages.append(
@@ -430,20 +432,18 @@ async def _map_messages(
         return system_prompt, bedrock_messages
 
     @staticmethod
-    async def _map_user_prompt(part: UserPromptPart) -> list[MessageUnionTypeDef]:
+    async def _map_user_prompt(part: UserPromptPart, document_count: Iterator[int]) -> list[MessageUnionTypeDef]:
         content: list[ContentBlockUnionTypeDef] = []
         if isinstance(part.content, str):
             content.append({'text': part.content})
         else:
-            document_count = 0
             for item in part.content:
                 if isinstance(item, str):
                     content.append({'text': item})
                 elif isinstance(item, BinaryContent):
                     format = item.format
                     if item.is_document:
-                        document_count += 1
-                        name = f'Document {document_count}'
+                        name = f'Document {next(document_count)}'
                         assert format in ('pdf', 'txt', 'csv', 'doc', 'docx', 'xls', 'xlsx', 'html', 'md')
                         content.append({'document': {'name': name, 'format': format, 'source': {'bytes': item.data}}})
                     elif item.is_image:
@@ -464,8 +464,7 @@ async def _map_user_prompt(part: UserPromptPart) -> list[MessageUnionTypeDef]:
                         content.append({'image': image})
 
                     elif item.kind == 'document-url':
-                        document_count += 1
-                        name = f'Document {document_count}'
+                        name = f'Document {next(document_count)}'
                         data = response.content
                         content.append({'document': {'name': name, 'format': item.format, 'source': {'bytes': data}}})
 

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -58,6 +58,7 @@
     'gemini-2.0-flash',
     'gemini-2.0-flash-lite-preview-02-05',
     'gemini-2.0-pro-exp-02-05',
+    'gemini-2.5-flash-preview-04-17',
     'gemini-2.5-pro-exp-03-25',
     'gemini-2.5-pro-preview-03-25',
 ]