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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.4.7
+Version: 0.4.9
 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.9
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -51,7 +51,7 @@ Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'c
 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.9; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.24.0; extra == 'google'
 Provides-Extra: groq
@@ -66,6 +66,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.9}/pydantic_ai/_tool_manager.py

@@ -2,18 +2,17 @@ from __future__ import annotations
 
 import json
 from collections.abc import Iterable
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
 from pydantic import ValidationError
 from typing_extensions import assert_never
 
-from pydantic_ai.output import DeferredToolCalls
-
 from . import messages as _messages
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
 from .messages import ToolCallPart
+from .output import DeferredToolCalls
 from .tools import ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
@@ -28,6 +27,8 @@ class ToolManager(Generic[AgentDepsT]):
     """The toolset that provides the tools for this run step."""
     tools: dict[str, ToolsetTool[AgentDepsT]]
     """The cached tools for this run step."""
+    failed_tools: set[str] = field(default_factory=set)
+    """Names of tools that failed in this run step."""
 
     @classmethod
     async def build(cls, toolset: AbstractToolset[AgentDepsT], ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
@@ -40,7 +41,10 @@ class ToolManager(Generic[AgentDepsT]):
 
     async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
         """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
-        return await self.__class__.build(self.toolset, replace(ctx, retries=self.ctx.retries))
+        retries = {
+            failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1 for failed_tool_name in self.failed_tools
+        }
+        return await self.__class__.build(self.toolset, replace(ctx, retries=retries))
 
     @property
     def tool_defs(self) -> list[ToolDefinition]:
@@ -54,20 +58,25 @@ class ToolManager(Generic[AgentDepsT]):
         except KeyError:
             return None
 
-    async def handle_call(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    async def handle_call(
+        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> Any:
         """Handle a tool call by validating the arguments, calling the tool, and handling retries.
 
         Args:
             call: The tool call part to handle.
             allow_partial: Whether to allow partial validation of the tool arguments.
+            wrap_validation_errors: Whether to wrap validation errors in a retry prompt part.
         """
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
             # Output tool calls are not traced
-            return await self._call_tool(call, allow_partial)
+            return await self._call_tool(call, allow_partial, wrap_validation_errors)
         else:
-            return await self._call_tool_traced(call, allow_partial)
+            return await self._call_tool_traced(call, allow_partial, wrap_validation_errors)
 
-    async def _call_tool(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    async def _call_tool(
+        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> Any:
         name = call.tool_name
         tool = self.tools.get(name)
         try:
@@ -92,7 +101,7 @@ class ToolManager(Generic[AgentDepsT]):
             else:
                 args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
 
-            output = await self.toolset.call_tool(name, args_dict, ctx, tool)
+            return await self.toolset.call_tool(name, args_dict, ctx, tool)
         except (ValidationError, ModelRetry) as e:
             max_retries = tool.max_retries if tool is not None else 1
             current_retry = self.ctx.retries.get(name, 0)
@@ -100,30 +109,33 @@ class ToolManager(Generic[AgentDepsT]):
             if current_retry == max_retries:
                 raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
             else:
-                if isinstance(e, ValidationError):
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.errors(include_url=False, include_context=False),
-                        tool_call_id=call.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-                elif isinstance(e, ModelRetry):
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.message,
-                        tool_call_id=call.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-                else:
-                    assert_never(e)
+                if wrap_validation_errors:
+                    if isinstance(e, ValidationError):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.errors(include_url=False, include_context=False),
+                            tool_call_id=call.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+                    elif isinstance(e, ModelRetry):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.message,
+                            tool_call_id=call.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+                    else:
+                        assert_never(e)
+
+                if not allow_partial:
+                    # If we're validating partial arguments, we don't want to count this as a failed tool as it may still succeed once the full arguments are received.
+                    self.failed_tools.add(name)
 
-                self.ctx.retries[name] = current_retry + 1
                 raise e
-        else:
-            self.ctx.retries.pop(name, None)
-            return output
 
-    async def _call_tool_traced(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    async def _call_tool_traced(
+        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> Any:
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
         span_attributes = {
             'gen_ai.tool.name': call.tool_name,
@@ -152,7 +164,7 @@ class ToolManager(Generic[AgentDepsT]):
         }
         with self.ctx.tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
             try:
-                tool_result = await self._call_tool(call, allow_partial)
+                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
             except ToolRetryError as e:
                 part = e.tool_retry
                 if self.ctx.trace_include_content and span.is_recording():

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

@@ -9,11 +9,13 @@ from __future__ import annotations
 import json
 import uuid
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, field
+from dataclasses import Field, dataclass, field, replace
 from http import HTTPStatus
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Final,
     Generic,
     Protocol,
@@ -21,6 +23,11 @@ from typing import (
     runtime_checkable,
 )
 
+from pydantic_ai.exceptions import UserError
+
+if TYPE_CHECKING:
+    pass
+
 try:
     from ag_ui.core import (
         AssistantMessage,
@@ -288,8 +295,24 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             if not run_input.messages:
                 raise _NoMessagesError
 
+            raw_state: dict[str, Any] = run_input.state or {}
             if isinstance(deps, StateHandler):
-                deps.state = run_input.state
+                if isinstance(deps.state, BaseModel):
+                    try:
+                        state = type(deps.state).model_validate(raw_state)
+                    except ValidationError as e:  # pragma: no cover
+                        raise _InvalidStateError from e
+                else:
+                    state = raw_state
+
+                deps = replace(deps, state=state)
+            elif raw_state:
+                raise UserError(
+                    f'AG-UI state is provided but `deps` of type `{type(deps).__name__}` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.'
+                )
+            else:
+                # `deps` not being a `StateHandler` is OK if there is no state.
+                pass
 
             messages = _messages_from_ag_ui(run_input.messages)
 
@@ -311,7 +334,7 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             yield encoder.encode(
                 RunErrorEvent(message=e.message, code=e.code),
             )
-        except Exception as e:  # pragma: no cover
+        except Exception as e:
             yield encoder.encode(
                 RunErrorEvent(message=str(e)),
             )
@@ -486,6 +509,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 +528,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):
@@ -531,7 +554,11 @@ def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
 
 @runtime_checkable
 class StateHandler(Protocol):
-    """Protocol for state handlers in agent runs."""
+    """Protocol for state handlers in agent runs. Requires the class to be a dataclass with a `state` field."""
+
+    # Has to be a dataclass so we can use `replace` to update the state.
+    # From https://github.com/python/typeshed/blob/9ab7fde0a0cd24ed7a72837fcb21093b811b80d8/stdlib/_typeshed/__init__.pyi#L352
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
 
     @property
     def state(self) -> State:
@@ -558,6 +585,7 @@ StateT = TypeVar('StateT', bound=BaseModel)
 """Type variable for the state type, which must be a subclass of `BaseModel`."""
 
 
+@dataclass
 class StateDeps(Generic[StateT]):
     """Provides AG-UI state management.
 
@@ -570,42 +598,7 @@ class StateDeps(Generic[StateT]):
     Implements the `StateHandler` protocol.
     """
 
-    def __init__(self, default: StateT) -> None:
-        """Initialize the state with the provided state type."""
-        self._state = default
-
-    @property
-    def state(self) -> StateT:
-        """Get the current state of the agent run.
-
-        Returns:
-            The current run state.
-        """
-        return self._state
-
-    @state.setter
-    def state(self, state: State) -> None:
-        """Set the state of the agent run.
-
-        This method is called to update the state of the agent run with the
-        provided state.
-
-        Implements the `StateHandler` protocol.
-
-        Args:
-            state: The run state, which must be `None` or model validate for the state type.
-
-        Raises:
-            InvalidStateError: If `state` does not validate.
-        """
-        if state is None:
-            # If state is None, we keep the current state, which will be the default state.
-            return
-
-        try:
-            self._state = type(self._state).model_validate(state)
-        except ValidationError as e:  # pragma: no cover
-            raise _InvalidStateError from e
+    state: StateT
 
 
 @dataclass(repr=False)

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.9}/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.9}/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.9}/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.7 → pydantic_ai_slim-0.4.9}/pydantic_ai/profiles/google.py

@@ -49,7 +49,6 @@ class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
             )
 
         schema.pop('title', None)
-        schema.pop('default', None)
         schema.pop('$schema', None)
         if (const := schema.pop('const', None)) is not None:
             # Gemini doesn't support const, but it does support enum with a single value

{pydantic_ai_slim-0.4.7 → pydantic_ai_slim-0.4.9}/pydantic_ai/result.py

@@ -67,7 +67,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 except ValidationError:
                     pass
         if self._final_result_event is not None:  # pragma: no branch
-            yield await self._validate_response(self._raw_stream_response.get(), allow_partial=False)
+            yield await self._validate_response(self._raw_stream_response.get())
 
     async def stream_responses(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
         """Asynchronously stream the (unvalidated) model responses for the agent."""
@@ -128,7 +128,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         async for _ in self:
             pass
 
-        return await self._validate_response(self._raw_stream_response.get(), allow_partial=False)
+        return await self._validate_response(self._raw_stream_response.get())
 
     async def _validate_response(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
         """Validate a structured result message."""
@@ -150,7 +150,9 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                     f'Invalid response, unable to find tool call for {output_tool_name!r}'
                 )
-            return await self._tool_manager.handle_call(tool_call, allow_partial=allow_partial)
+            return await self._tool_manager.handle_call(
+                tool_call, allow_partial=allow_partial, wrap_validation_errors=False
+            )
         elif deferred_tool_calls := self._tool_manager.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
                 raise exceptions.UserError(

pydantic_ai_slim-0.4.9/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.9}/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.9}/pyproject.toml

@@ -84,6 +84,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 = [