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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.4.7
+Version: 0.4.8
 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.7
+Requires-Dist: pydantic-graph==0.4.8
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -48,10 +48,11 @@ Requires-Dist: prompt-toolkit>=3; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'cohere'
+Requires-Dist: tokenizers<=0.21.2; extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.4.7; extra == 'evals'
+Requires-Dist: pydantic-evals==0.4.8; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.24.0; extra == 'google'
 Provides-Extra: groq
@@ -66,6 +67,8 @@ Provides-Extra: mistral
 Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
 Provides-Extra: openai
 Requires-Dist: openai>=1.92.0; extra == 'openai'
+Provides-Extra: retries
+Requires-Dist: tenacity>=8.2.3; extra == 'retries'
 Provides-Extra: tavily
 Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
 Provides-Extra: vertexai

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.8}/pydantic_ai/agent.py

@@ -1792,9 +1792,11 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         """
         async with self._enter_lock:
             if self._entered_count == 0:
-                self._exit_stack = AsyncExitStack()
-                toolset = self._get_toolset()
-                await self._exit_stack.enter_async_context(toolset)
+                async with AsyncExitStack() as exit_stack:
+                    toolset = self._get_toolset()
+                    await exit_stack.enter_async_context(toolset)
+
+                    self._exit_stack = exit_stack.pop_all()
             self._entered_count += 1
         return self
 

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

@@ -201,25 +201,24 @@ class MCPServer(AbstractToolset[Any], ABC):
         """
         async with self._enter_lock:
             if self._running_count == 0:
-                self._exit_stack = AsyncExitStack()
-
-                self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(
-                    self.client_streams()
-                )
-                client = ClientSession(
-                    read_stream=self._read_stream,
-                    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)
-
-                with anyio.fail_after(self.timeout):
-                    await self._client.initialize()
-
-                    if log_level := self.log_level:
-                        await self._client.set_logging_level(log_level)
+                async with AsyncExitStack() as exit_stack:
+                    self._read_stream, self._write_stream = await exit_stack.enter_async_context(self.client_streams())
+                    client = ClientSession(
+                        read_stream=self._read_stream,
+                        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 exit_stack.enter_async_context(client)
+
+                    with anyio.fail_after(self.timeout):
+                        await self._client.initialize()
+
+                        if log_level := self.log_level:
+                            await self._client.set_logging_level(log_level)
+
+                    self._exit_stack = exit_stack.pop_all()
             self._running_count += 1
         return self
 
@@ -544,6 +543,7 @@ class _MCPServerHTTP(MCPServer):
         self.max_retries = max_retries
         self.sampling_model = sampling_model
         self.read_timeout = read_timeout
+        self.__post_init__()
 
     @property
     @abstractmethod

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.8}/pydantic_ai/messages.py

@@ -815,11 +815,16 @@ class ModelResponse:
                         },
                     }
                 )
-            elif isinstance(part, TextPart):
-                if body.get('content'):
-                    body = new_event_body()
-                if settings.include_content:
-                    body['content'] = part.content
+            elif isinstance(part, (TextPart, ThinkingPart)):
+                kind = part.part_kind
+                body.setdefault('content', []).append(
+                    {'kind': kind, **({'text': part.content} if settings.include_content else {})}
+                )
+
+        if content := body.get('content'):
+            text_content = content[0].get('text')
+            if content == [{'kind': 'text', 'text': text_content}]:
+                body['content'] = text_content
 
         return result
 

pydantic_ai_slim-0.4.8/pydantic_ai/retries.py

@@ -0,0 +1,249 @@
+"""Retries utilities based on tenacity, especially for HTTP requests.
+
+This module provides HTTP transport wrappers and wait strategies that integrate with
+the tenacity library to add retry capabilities to HTTP requests. The transports can be
+used with HTTP clients that support custom transports (such as httpx), while the wait
+strategies can be used with any tenacity retry decorator.
+
+The module includes:
+- TenacityTransport: Synchronous HTTP transport with retry capabilities
+- AsyncTenacityTransport: Asynchronous HTTP transport with retry capabilities
+- wait_retry_after: Wait strategy that respects HTTP Retry-After headers
+"""
+
+from __future__ import annotations
+
+from httpx import AsyncBaseTransport, AsyncHTTPTransport, BaseTransport, HTTPTransport, Request, Response
+
+try:
+    from tenacity import AsyncRetrying, Retrying
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `tenacity` to use the retries utilities, '
+        'you can use the `retries` optional group — `pip install "pydantic-ai-slim[retries]"`'
+    ) from _import_error
+
+
+__all__ = ['TenacityTransport', 'AsyncTenacityTransport', 'wait_retry_after']
+
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+from typing import Callable, cast
+
+from httpx import HTTPStatusError
+from tenacity import RetryCallState, wait_exponential
+
+
+class TenacityTransport(BaseTransport):
+    """Synchronous HTTP transport with tenacity-based retry functionality.
+
+    This transport wraps another BaseTransport and adds retry capabilities using the tenacity library.
+    It can be configured to retry requests based on various conditions such as specific exception types,
+    response status codes, or custom validation logic.
+
+    The transport works by intercepting HTTP requests and responses, allowing the tenacity controller
+    to determine when and how to retry failed requests. The validate_response function can be used
+    to convert HTTP responses into exceptions that trigger retries.
+
+    Args:
+        wrapped: The underlying transport to wrap and add retry functionality to.
+        controller: The tenacity Retrying instance that defines the retry behavior
+                   (retry conditions, wait strategy, stop conditions, etc.).
+        validate_response: Optional callable that takes a Response and can raise an exception
+            to be handled by the controller if the response should trigger a retry.
+            Common use case is to raise exceptions for certain HTTP status codes.
+            If None, no response validation is performed.
+
+    Example:
+        ```python
+        from httpx import Client, HTTPTransport, HTTPStatusError
+        from tenacity import Retrying, stop_after_attempt, retry_if_exception_type
+        from pydantic_ai.retries import TenacityTransport, wait_retry_after
+
+        transport = TenacityTransport(
+            HTTPTransport(),
+            Retrying(
+                retry=retry_if_exception_type(HTTPStatusError),
+                wait=wait_retry_after(max_wait=300),
+                stop=stop_after_attempt(5),
+                reraise=True
+            ),
+            validate_response=lambda r: r.raise_for_status()
+        )
+        client = Client(transport=transport)
+        ```
+    """
+
+    def __init__(
+        self,
+        controller: Retrying,
+        wrapped: BaseTransport | None = None,
+        validate_response: Callable[[Response], None] | None = None,
+    ):
+        self.controller = controller
+        self.wrapped = wrapped or HTTPTransport()
+        self.validate_response = validate_response
+
+    def handle_request(self, request: Request) -> Response:
+        """Handle an HTTP request with retry logic.
+
+        Args:
+            request: The HTTP request to handle.
+
+        Returns:
+            The HTTP response.
+
+        Raises:
+            RuntimeError: If the retry controller did not make any attempts.
+            Exception: Any exception raised by the wrapped transport or validation function.
+        """
+        for attempt in self.controller:
+            with attempt:
+                response = self.wrapped.handle_request(request)
+                if self.validate_response:
+                    self.validate_response(response)
+                return response
+        raise RuntimeError('The retry controller did not make any attempts')  # pragma: no cover
+
+
+class AsyncTenacityTransport(AsyncBaseTransport):
+    """Asynchronous HTTP transport with tenacity-based retry functionality.
+
+    This transport wraps another AsyncBaseTransport and adds retry capabilities using the tenacity library.
+    It can be configured to retry requests based on various conditions such as specific exception types,
+    response status codes, or custom validation logic.
+
+    The transport works by intercepting HTTP requests and responses, allowing the tenacity controller
+    to determine when and how to retry failed requests. The validate_response function can be used
+    to convert HTTP responses into exceptions that trigger retries.
+
+    Args:
+        wrapped: The underlying async transport to wrap and add retry functionality to.
+        controller: The tenacity AsyncRetrying instance that defines the retry behavior
+                   (retry conditions, wait strategy, stop conditions, etc.).
+        validate_response: Optional callable that takes a Response and can raise an exception
+            to be handled by the controller if the response should trigger a retry.
+            Common use case is to raise exceptions for certain HTTP status codes.
+            If None, no response validation is performed.
+
+    Example:
+        ```python
+        from httpx import AsyncClient, HTTPStatusError
+        from tenacity import AsyncRetrying, stop_after_attempt, retry_if_exception_type
+        from pydantic_ai.retries import AsyncTenacityTransport, wait_retry_after
+
+        transport = AsyncTenacityTransport(
+            AsyncRetrying(
+                retry=retry_if_exception_type(HTTPStatusError),
+                wait=wait_retry_after(max_wait=300),
+                stop=stop_after_attempt(5),
+                reraise=True
+            ),
+            validate_response=lambda r: r.raise_for_status()
+        )
+        client = AsyncClient(transport=transport)
+        ```
+    """
+
+    def __init__(
+        self,
+        controller: AsyncRetrying,
+        wrapped: AsyncBaseTransport | None = None,
+        validate_response: Callable[[Response], None] | None = None,
+    ):
+        self.controller = controller
+        self.wrapped = wrapped or AsyncHTTPTransport()
+        self.validate_response = validate_response
+
+    async def handle_async_request(self, request: Request) -> Response:
+        """Handle an async HTTP request with retry logic.
+
+        Args:
+            request: The HTTP request to handle.
+
+        Returns:
+            The HTTP response.
+
+        Raises:
+            RuntimeError: If the retry controller did not make any attempts.
+            Exception: Any exception raised by the wrapped transport or validation function.
+        """
+        async for attempt in self.controller:
+            with attempt:
+                response = await self.wrapped.handle_async_request(request)
+                if self.validate_response:
+                    self.validate_response(response)
+                return response
+        raise RuntimeError('The retry controller did not make any attempts')  # pragma: no cover
+
+
+def wait_retry_after(
+    fallback_strategy: Callable[[RetryCallState], float] | None = None, max_wait: float = 300
+) -> Callable[[RetryCallState], float]:
+    """Create a tenacity-compatible wait strategy that respects HTTP Retry-After headers.
+
+    This wait strategy checks if the exception contains an HTTPStatusError with a
+    Retry-After header, and if so, waits for the time specified in the header.
+    If no header is present or parsing fails, it falls back to the provided strategy.
+
+    The Retry-After header can be in two formats:
+    - An integer representing seconds to wait
+    - An HTTP date string representing when to retry
+
+    Args:
+        fallback_strategy: Wait strategy to use when no Retry-After header is present
+                          or parsing fails. Defaults to exponential backoff with max 60s.
+        max_wait: Maximum time to wait in seconds, regardless of header value.
+                 Defaults to 300 (5 minutes).
+
+    Returns:
+        A wait function that can be used with tenacity retry decorators.
+
+    Example:
+        ```python
+        from httpx import AsyncClient, HTTPStatusError
+        from tenacity import AsyncRetrying, stop_after_attempt, retry_if_exception_type
+        from pydantic_ai.retries import AsyncTenacityTransport, wait_retry_after
+
+        transport = AsyncTenacityTransport(
+            AsyncRetrying(
+                retry=retry_if_exception_type(HTTPStatusError),
+                wait=wait_retry_after(max_wait=120),
+                stop=stop_after_attempt(5),
+                reraise=True
+            ),
+            validate_response=lambda r: r.raise_for_status()
+        )
+        client = AsyncClient(transport=transport)
+        ```
+    """
+    if fallback_strategy is None:
+        fallback_strategy = wait_exponential(multiplier=1, max=60)
+
+    def wait_func(state: RetryCallState) -> float:
+        exc = state.outcome.exception() if state.outcome else None
+        if isinstance(exc, HTTPStatusError):
+            retry_after = exc.response.headers.get('retry-after')
+            if retry_after:
+                try:
+                    # Try parsing as seconds first
+                    wait_seconds = int(retry_after)
+                    return min(float(wait_seconds), max_wait)
+                except ValueError:
+                    # Try parsing as HTTP date
+                    try:
+                        retry_time = cast(datetime, parsedate_to_datetime(retry_after))
+                        assert isinstance(retry_time, datetime)
+                        now = datetime.now(timezone.utc)
+                        wait_seconds = (retry_time - now).total_seconds()
+
+                        if wait_seconds > 0:
+                            return min(wait_seconds, max_wait)
+                    except (ValueError, TypeError, AssertionError):
+                        # If date parsing fails, fall back to fallback strategy
+                        pass
+
+        # Use fallback strategy
+        return fallback_strategy(state)
+
+    return wait_func

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.8}/pydantic_ai/toolsets/combined.py

@@ -43,9 +43,10 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
     async def __aenter__(self) -> Self:
         async with self._enter_lock:
             if self._entered_count == 0:
-                self._exit_stack = AsyncExitStack()
-                for toolset in self.toolsets:
-                    await self._exit_stack.enter_async_context(toolset)
+                async with AsyncExitStack() as exit_stack:
+                    for toolset in self.toolsets:
+                        await exit_stack.enter_async_context(toolset)
+                    self._exit_stack = exit_stack.pop_all()
             self._entered_count += 1
         return self
 

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.8}/pyproject.toml

@@ -63,7 +63,12 @@ dependencies = [
 logfire = ["logfire>=3.11.0"]
 # Models
 openai = ["openai>=1.92.0"]
-cohere = ["cohere>=5.16.0; platform_system != 'Emscripten'"]
+cohere = [
+    "cohere>=5.16.0; platform_system != 'Emscripten'",
+    # Remove once all wheels for 0.21.4+ are built successfully
+    # https://github.com/huggingface/tokenizers/actions/runs/16570140346/job/46860152621
+    "tokenizers<=0.21.2",
+]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
 google = ["google-genai>=1.24.0"]
 anthropic = ["anthropic>=0.52.0"]
@@ -84,6 +89,8 @@ evals = ["pydantic-evals=={{ version }}"]
 a2a = ["fasta2a>=0.4.1"]
 # AG-UI
 ag-ui = ["ag-ui-protocol>=0.1.8", "starlette>=0.45.3"]
+# Retries
+retries = ["tenacity>=8.2.3"]
 
 [dependency-groups]
 dev = [

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

@@ -486,6 +486,9 @@ def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
         if isinstance(msg, UserMessage):
             result.append(ModelRequest(parts=[UserPromptPart(content=msg.content)]))
         elif isinstance(msg, AssistantMessage):
+            if msg.content:
+                result.append(ModelResponse(parts=[TextPart(content=msg.content)]))
+
             if msg.tool_calls:
                 for tool_call in msg.tool_calls:
                     tool_calls[tool_call.id] = tool_call.function.name
@@ -502,9 +505,6 @@ def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
                         ]
                     )
                 )
-
-            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):