pydantic-ai-slim 0.4.6__tar.gz → 0.4.7__tar.gz

{pydantic_ai_slim-0.4.6 → pydantic_ai_slim-0.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.4.6
+Version: 0.4.7
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>, Douwe Maan <douwe@pydantic.dev>
 License-Expression: MIT
@@ -30,7 +30,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.4.6
+Requires-Dist: pydantic-graph==0.4.7
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -47,21 +47,21 @@ Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
 Requires-Dist: prompt-toolkit>=3; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
-Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == 'cohere'
+Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.4.6; extra == 'evals'
+Requires-Dist: pydantic-evals==0.4.7; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.24.0; extra == 'google'
 Provides-Extra: groq
 Requires-Dist: groq>=0.19.0; extra == 'groq'
 Provides-Extra: huggingface
-Requires-Dist: huggingface-hub[inference]>=0.33.2; extra == 'huggingface'
+Requires-Dist: huggingface-hub[inference]>=0.33.5; extra == 'huggingface'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Provides-Extra: mcp
-Requires-Dist: mcp>=1.9.4; (python_version >= '3.10') and extra == 'mcp'
+Requires-Dist: mcp>=1.10.0; (python_version >= '3.10') and extra == 'mcp'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
 Provides-Extra: openai

{pydantic_ai_slim-0.4.6 → pydantic_ai_slim-0.4.7}/pydantic_ai/_parts_manager.py

@@ -17,6 +17,7 @@ from collections.abc import Hashable
 from dataclasses import dataclass, field, replace
 from typing import Any, Union
 
+from pydantic_ai._thinking_part import END_THINK_TAG, START_THINK_TAG
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
     ModelResponsePart,
@@ -69,9 +70,10 @@ class ModelResponsePartsManager:
     def handle_text_delta(
         self,
         *,
-        vendor_part_id: Hashable | None,
+        vendor_part_id: VendorId | None,
         content: str,
-    ) -> ModelResponseStreamEvent:
+        extract_think_tags: bool = False,
+    ) -> ModelResponseStreamEvent | None:
         """Handle incoming text content, creating or updating a TextPart in the manager as appropriate.
 
         When `vendor_part_id` is None, the latest part is updated if it exists and is a TextPart;
@@ -83,9 +85,12 @@ class ModelResponsePartsManager:
                 of text. If None, a new part will be created unless the latest part is already
                 a TextPart.
             content: The text content to append to the appropriate TextPart.
+            extract_think_tags: Whether to extract `<think>` tags from the text content and handle them as thinking parts.
 
         Returns:
-            A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated.
+            - A `PartStartEvent` if a new part was created.
+            - A `PartDeltaEvent` if an existing part was updated.
+            - `None` if no new event is emitted (e.g., the first text part was all whitespace).
 
         Raises:
             UnexpectedModelBehavior: If attempting to apply text content to a part that is not a TextPart.
@@ -104,11 +109,32 @@ class ModelResponsePartsManager:
             part_index = self._vendor_id_to_part_index.get(vendor_part_id)
             if part_index is not None:
                 existing_part = self._parts[part_index]
-                if not isinstance(existing_part, TextPart):
+
+                if extract_think_tags and isinstance(existing_part, ThinkingPart):
+                    # We may be building a thinking part instead of a text part if we had previously seen a `<think>` tag
+                    if content == END_THINK_TAG:
+                        # When we see `</think>`, we're done with the thinking part and the next text delta will need a new part
+                        self._vendor_id_to_part_index.pop(vendor_part_id)
+                        return None
+                    else:
+                        return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)
+                elif isinstance(existing_part, TextPart):
+                    existing_text_part_and_index = existing_part, part_index
+                else:
                     raise UnexpectedModelBehavior(f'Cannot apply a text delta to {existing_part=}')
-                existing_text_part_and_index = existing_part, part_index
+
+        if extract_think_tags and content == START_THINK_TAG:
+            # When we see a `<think>` tag (which is a single token), we'll build a new thinking part instead
+            self._vendor_id_to_part_index.pop(vendor_part_id, None)
+            return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
 
         if existing_text_part_and_index is None:
+            # If the first text delta is all whitespace, don't emit a new part yet.
+            # This is a workaround for models that emit `<think>\n</think>\n\n` ahead of tool calls (e.g. Ollama + Qwen3),
+            # which we don't want to end up treating as a final result.
+            if content.isspace():
+                return None
+
             # There is no existing text part that should be updated, so create a new one
             new_part_index = len(self._parts)
             part = TextPart(content=content)

{pydantic_ai_slim-0.4.6 → pydantic_ai_slim-0.4.7}/pydantic_ai/ag_ui.py

@@ -291,12 +291,12 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             if isinstance(deps, StateHandler):
                 deps.state = run_input.state
 
-            history = _History.from_ag_ui(run_input.messages)
+            messages = _messages_from_ag_ui(run_input.messages)
 
             async with self.agent.iter(
                 user_prompt=None,
                 output_type=[output_type or self.agent.output_type, DeferredToolCalls],
-                message_history=history.messages,
+                message_history=messages,
                 model=model,
                 deps=deps,
                 model_settings=model_settings,
@@ -305,7 +305,7 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
                 infer_name=infer_name,
                 toolsets=toolsets,
             ) as run:
-                async for event in self._agent_stream(run, history):
+                async for event in self._agent_stream(run):
                     yield encoder.encode(event)
         except _RunError as e:
             yield encoder.encode(
@@ -327,20 +327,18 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
     async def _agent_stream(
         self,
         run: AgentRun[AgentDepsT, Any],
-        history: _History,
     ) -> AsyncGenerator[BaseEvent, None]:
         """Run the agent streaming responses using AG-UI protocol events.
 
         Args:
             run: The agent run to process.
-            history: The history of messages and tool calls to use for the run.
 
         Yields:
             AG-UI Server-Sent Events (SSE).
         """
         async for node in run:
+            stream_ctx = _RequestStreamContext()
             if isinstance(node, ModelRequestNode):
-                stream_ctx = _RequestStreamContext()
                 async with node.stream(run.ctx) as request_stream:
                     async for agent_event in request_stream:
                         async for msg in self._handle_model_request_event(stream_ctx, agent_event):
@@ -352,8 +350,8 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             elif isinstance(node, CallToolsNode):
                 async with node.stream(run.ctx) as handle_stream:
                     async for event in handle_stream:
-                        if isinstance(event, FunctionToolResultEvent) and isinstance(event.result, ToolReturnPart):
-                            async for msg in self._handle_tool_result_event(event.result, history.prompt_message_id):
+                        if isinstance(event, FunctionToolResultEvent):
+                            async for msg in self._handle_tool_result_event(stream_ctx, event):
                                 yield msg
 
     async def _handle_model_request_event(
@@ -382,19 +380,26 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
                 yield TextMessageStartEvent(
                     message_id=message_id,
                 )
-                stream_ctx.part_end = TextMessageEndEvent(
-                    message_id=message_id,
-                )
                 if part.content:  # pragma: no branch
                     yield TextMessageContentEvent(
                         message_id=message_id,
                         delta=part.content,
                     )
+                stream_ctx.part_end = TextMessageEndEvent(
+                    message_id=message_id,
+                )
             elif isinstance(part, ToolCallPart):  # pragma: no branch
+                message_id = stream_ctx.message_id or stream_ctx.new_message_id()
                 yield ToolCallStartEvent(
                     tool_call_id=part.tool_call_id,
                     tool_call_name=part.tool_name,
+                    parent_message_id=message_id,
                 )
+                if part.args:
+                    yield ToolCallArgsEvent(
+                        tool_call_id=part.tool_call_id,
+                        delta=part.args if isinstance(part.args, str) else json.dumps(part.args),
+                    )
                 stream_ctx.part_end = ToolCallEndEvent(
                     tool_call_id=part.tool_call_id,
                 )
@@ -407,7 +412,7 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
                 # used to indicate the start of thinking.
                 yield ThinkingTextMessageContentEvent(
                     type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
-                    delta=part.content or '',
+                    delta=part.content,
                 )
                 stream_ctx.part_end = ThinkingTextMessageEndEvent(
                     type=EventType.THINKING_TEXT_MESSAGE_END,
@@ -435,20 +440,25 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
 
     async def _handle_tool_result_event(
         self,
-        result: ToolReturnPart,
-        prompt_message_id: str,
+        stream_ctx: _RequestStreamContext,
+        event: FunctionToolResultEvent,
     ) -> AsyncGenerator[BaseEvent, None]:
         """Convert a tool call result to AG-UI events.
 
         Args:
-            result: The tool call result to process.
-            prompt_message_id: The message ID of the prompt that initiated the tool call.
+            stream_ctx: The request stream context to manage state.
+            event: The tool call result event to process.
 
         Yields:
             AG-UI Server-Sent Events (SSE).
         """
+        result = event.result
+        if not isinstance(result, ToolReturnPart):
+            return
+
+        message_id = stream_ctx.new_message_id()
         yield ToolCallResultEvent(
-            message_id=prompt_message_id,
+            message_id=message_id,
             type=EventType.TOOL_CALL_RESULT,
             role='tool',
             tool_call_id=result.tool_call_id,
@@ -468,75 +478,55 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
                     yield item
 
 
-@dataclass
-class _History:
-    """A simple history representation for AG-UI protocol."""
-
-    prompt_message_id: str  # The ID of the last user message.
-    messages: list[ModelMessage]
-
-    @classmethod
-    def from_ag_ui(cls, messages: list[Message]) -> _History:
-        """Convert a AG-UI history to a Pydantic AI one.
-
-        Args:
-            messages: List of AG-UI messages to convert.
-
-        Returns:
-            List of Pydantic AI model messages.
-        """
-        prompt_message_id = ''
-        result: list[ModelMessage] = []
-        tool_calls: dict[str, str] = {}  # Tool call ID to tool name mapping.
-        for msg in messages:
-            if isinstance(msg, UserMessage):
-                prompt_message_id = msg.id
-                result.append(ModelRequest(parts=[UserPromptPart(content=msg.content)]))
-            elif isinstance(msg, AssistantMessage):
-                if msg.tool_calls:
-                    for tool_call in msg.tool_calls:
-                        tool_calls[tool_call.id] = tool_call.function.name
-
-                    result.append(
-                        ModelResponse(
-                            parts=[
-                                ToolCallPart(
-                                    tool_name=tool_call.function.name,
-                                    tool_call_id=tool_call.id,
-                                    args=tool_call.function.arguments,
-                                )
-                                for tool_call in msg.tool_calls
-                            ]
-                        )
-                    )
-
-                if msg.content:
-                    result.append(ModelResponse(parts=[TextPart(content=msg.content)]))
-            elif isinstance(msg, SystemMessage):
-                result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
-            elif isinstance(msg, ToolMessage):
-                tool_name = tool_calls.get(msg.tool_call_id)
-                if tool_name is None:  # pragma: no cover
-                    raise _ToolCallNotFoundError(tool_call_id=msg.tool_call_id)
+def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
+    """Convert a AG-UI history to a Pydantic AI one."""
+    result: list[ModelMessage] = []
+    tool_calls: dict[str, str] = {}  # Tool call ID to tool name mapping.
+    for msg in messages:
+        if isinstance(msg, UserMessage):
+            result.append(ModelRequest(parts=[UserPromptPart(content=msg.content)]))
+        elif isinstance(msg, AssistantMessage):
+            if msg.tool_calls:
+                for tool_call in msg.tool_calls:
+                    tool_calls[tool_call.id] = tool_call.function.name
 
                 result.append(
-                    ModelRequest(
+                    ModelResponse(
                         parts=[
-                            ToolReturnPart(
-                                tool_name=tool_name,
-                                content=msg.content,
-                                tool_call_id=msg.tool_call_id,
+                            ToolCallPart(
+                                tool_name=tool_call.function.name,
+                                tool_call_id=tool_call.id,
+                                args=tool_call.function.arguments,
                             )
+                            for tool_call in msg.tool_calls
                         ]
                     )
                 )
-            elif isinstance(msg, DeveloperMessage):  # pragma: no branch
-                result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
 
-        return cls(
-            prompt_message_id=prompt_message_id,
-            messages=result,
-        )
+            if msg.content:
+                result.append(ModelResponse(parts=[TextPart(content=msg.content)]))
+        elif isinstance(msg, SystemMessage):
+            result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
+        elif isinstance(msg, ToolMessage):
+            tool_name = tool_calls.get(msg.tool_call_id)
+            if tool_name is None:  # pragma: no cover
+                raise _ToolCallNotFoundError(tool_call_id=msg.tool_call_id)
+
+            result.append(
+                ModelRequest(
+                    parts=[
+                        ToolReturnPart(
+                            tool_name=tool_name,
+                            content=msg.content,
+                            tool_call_id=msg.tool_call_id,
+                        )
+                    ]
+                )
+            )
+        elif isinstance(msg, DeveloperMessage):  # pragma: no branch
+            result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
+
+    return result
 
 
 @runtime_checkable

{pydantic_ai_slim-0.4.6 → pydantic_ai_slim-0.4.7}/pydantic_ai/mcp.py

@@ -2,11 +2,13 @@ from __future__ import annotations
 
 import base64
 import functools
+import warnings
 from abc import ABC, abstractmethod
 from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field, replace
+from datetime import timedelta
 from pathlib import Path
 from typing import Any, Callable
 
@@ -37,7 +39,7 @@ except ImportError as _import_error:
     ) from _import_error
 
 # after mcp imports so any import error maps to this file, not _mcp.py
-from . import _mcp, exceptions, messages, models
+from . import _mcp, _utils, exceptions, messages, models
 
 __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
 
@@ -59,6 +61,7 @@ class MCPServer(AbstractToolset[Any], ABC):
     log_level: mcp_types.LoggingLevel | None = None
     log_handler: LoggingFnT | None = None
     timeout: float = 5
+    read_timeout: float = 5 * 60
     process_tool_call: ProcessToolCallback | None = None
     allow_sampling: bool = True
     max_retries: int = 1
@@ -148,7 +151,7 @@ class MCPServer(AbstractToolset[Any], ABC):
             except McpError as e:
                 raise exceptions.ModelRetry(e.error.message)
 
-        content = [self._map_tool_result_part(part) for part in result.content]
+        content = [await self._map_tool_result_part(part) for part in result.content]
 
         if result.isError:
             text = '\n'.join(str(part) for part in content)
@@ -208,6 +211,7 @@ class MCPServer(AbstractToolset[Any], ABC):
                     write_stream=self._write_stream,
                     sampling_callback=self._sampling_callback if self.allow_sampling else None,
                     logging_callback=self.log_handler,
+                    read_timeout_seconds=timedelta(seconds=self.read_timeout),
                 )
                 self._client = await self._exit_stack.enter_async_context(client)
 
@@ -258,8 +262,8 @@ class MCPServer(AbstractToolset[Any], ABC):
             model=self.sampling_model.model_name,
         )
 
-    def _map_tool_result_part(
-        self, part: mcp_types.Content
+    async def _map_tool_result_part(
+        self, part: mcp_types.ContentBlock
     ) -> str | messages.BinaryContent | dict[str, Any] | list[Any]:
         # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values
 
@@ -281,18 +285,29 @@ class MCPServer(AbstractToolset[Any], ABC):
             )  # pragma: no cover
         elif isinstance(part, mcp_types.EmbeddedResource):
             resource = part.resource
-            if isinstance(resource, mcp_types.TextResourceContents):
-                return resource.text
-            elif isinstance(resource, mcp_types.BlobResourceContents):
-                return messages.BinaryContent(
-                    data=base64.b64decode(resource.blob),
-                    media_type=resource.mimeType or 'application/octet-stream',
-                )
-            else:
-                assert_never(resource)
+            return self._get_content(resource)
+        elif isinstance(part, mcp_types.ResourceLink):
+            resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri)
+            return (
+                self._get_content(resource_result.contents[0])
+                if len(resource_result.contents) == 1
+                else [self._get_content(resource) for resource in resource_result.contents]
+            )
         else:
             assert_never(part)
 
+    def _get_content(
+        self, resource: mcp_types.TextResourceContents | mcp_types.BlobResourceContents
+    ) -> str | messages.BinaryContent:
+        if isinstance(resource, mcp_types.TextResourceContents):
+            return resource.text
+        elif isinstance(resource, mcp_types.BlobResourceContents):
+            return messages.BinaryContent(
+                data=base64.b64decode(resource.blob), media_type=resource.mimeType or 'application/octet-stream'
+            )
+        else:
+            assert_never(resource)
+
 
 @dataclass
 class MCPServerStdio(MCPServer):
@@ -401,7 +416,7 @@ class MCPServerStdio(MCPServer):
         return f'MCPServerStdio(command={self.command!r}, args={self.args!r}, tool_prefix={self.tool_prefix!r})'
 
 
-@dataclass
+@dataclass(init=False)
 class _MCPServerHTTP(MCPServer):
     url: str
     """The URL of the endpoint on the MCP server."""
@@ -438,10 +453,10 @@ class _MCPServerHTTP(MCPServer):
         ```
     """
 
-    sse_read_timeout: float = 5 * 60
-    """Maximum time in seconds to wait for new SSE messages before timing out.
+    read_timeout: float = 5 * 60
+    """Maximum time in seconds to wait for new messages before timing out.
 
-    This timeout applies to the long-lived SSE connection after it's established.
+    This timeout applies to the long-lived connection after it's established.
     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).
     """
@@ -485,6 +500,51 @@ class _MCPServerHTTP(MCPServer):
     sampling_model: models.Model | None = None
     """The model to use for sampling."""
 
+    def __init__(
+        self,
+        *,
+        url: str,
+        headers: dict[str, str] | None = None,
+        http_client: httpx.AsyncClient | None = None,
+        read_timeout: float | None = None,
+        tool_prefix: str | None = None,
+        log_level: mcp_types.LoggingLevel | None = None,
+        log_handler: LoggingFnT | None = None,
+        timeout: float = 5,
+        process_tool_call: ProcessToolCallback | None = None,
+        allow_sampling: bool = True,
+        max_retries: int = 1,
+        sampling_model: models.Model | None = None,
+        **kwargs: Any,
+    ):
+        # Handle deprecated sse_read_timeout parameter
+        if 'sse_read_timeout' in kwargs:
+            if read_timeout is not None:
+                raise TypeError("'read_timeout' and 'sse_read_timeout' cannot be set at the same time.")
+
+            warnings.warn(
+                "'sse_read_timeout' is deprecated, use 'read_timeout' instead.", DeprecationWarning, stacklevel=2
+            )
+            read_timeout = kwargs.pop('sse_read_timeout')
+
+        _utils.validate_empty_kwargs(kwargs)
+
+        if read_timeout is None:
+            read_timeout = 5 * 60
+
+        self.url = url
+        self.headers = headers
+        self.http_client = http_client
+        self.tool_prefix = tool_prefix
+        self.log_level = log_level
+        self.log_handler = log_handler
+        self.timeout = timeout
+        self.process_tool_call = process_tool_call
+        self.allow_sampling = allow_sampling
+        self.max_retries = max_retries
+        self.sampling_model = sampling_model
+        self.read_timeout = read_timeout
+
     @property
     @abstractmethod
     def _transport_client(
@@ -522,7 +582,7 @@ class _MCPServerHTTP(MCPServer):
             self._transport_client,
             url=self.url,
             timeout=self.timeout,
-            sse_read_timeout=self.sse_read_timeout,
+            sse_read_timeout=self.read_timeout,
         )
 
         if self.http_client is not None:
@@ -549,7 +609,7 @@ class _MCPServerHTTP(MCPServer):
         return f'{self.__class__.__name__}(url={self.url!r}, tool_prefix={self.tool_prefix!r})'
 
 
-@dataclass
+@dataclass(init=False)
 class MCPServerSSE(_MCPServerHTTP):
     """An MCP server that connects over streamable HTTP connections.