pydantic-ai-slim 0.8.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

pydantic_ai/providers/litellm.py

@@ -0,0 +1,134 @@
+from __future__ import annotations as _annotations
+
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.groq import groq_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.profiles.qwen import qwen_model_profile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the LiteLLM provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class LiteLLMProvider(Provider[AsyncOpenAI]):
+    """Provider for LiteLLM API."""
+
+    @property
+    def name(self) -> str:
+        return 'litellm'
+
+    @property
+    def base_url(self) -> str:
+        return str(self.client.base_url)
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        # Map provider prefixes to their profile functions
+        provider_to_profile = {
+            'anthropic': anthropic_model_profile,
+            'openai': openai_model_profile,
+            'google': google_model_profile,
+            'mistralai': mistral_model_profile,
+            'mistral': mistral_model_profile,
+            'cohere': cohere_model_profile,
+            'amazon': amazon_model_profile,
+            'bedrock': amazon_model_profile,
+            'meta-llama': meta_model_profile,
+            'meta': meta_model_profile,
+            'groq': groq_model_profile,
+            'deepseek': deepseek_model_profile,
+            'moonshotai': moonshotai_model_profile,
+            'x-ai': grok_model_profile,
+            'qwen': qwen_model_profile,
+        }
+
+        profile = None
+
+        # Check if model name contains a provider prefix (e.g., "anthropic/claude-3")
+        if '/' in model_name:
+            provider_prefix, model_suffix = model_name.split('/', 1)
+            if provider_prefix in provider_to_profile:
+                profile = provider_to_profile[provider_prefix](model_suffix)
+
+        # If no profile found, default to OpenAI profile
+        if profile is None:
+            profile = openai_model_profile(model_name)
+
+        # As LiteLLMProvider is used with OpenAIModel, which uses OpenAIJsonSchemaTransformer,
+        # we maintain that behavior
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        http_client: AsyncHTTPClient,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        """Initialize a LiteLLM provider.
+
+        Args:
+            api_key: API key for the model provider. If None, LiteLLM will try to get it from environment variables.
+            api_base: Base URL for the model provider. Use this for custom endpoints or self-hosted models.
+            openai_client: Pre-configured OpenAI client. If provided, other parameters are ignored.
+            http_client: Custom HTTP client to use.
+        """
+        if openai_client is not None:
+            self._client = openai_client
+            return
+
+        # Create OpenAI client that will be used with LiteLLM's completion function
+        # The actual API calls will be intercepted and routed through LiteLLM
+        if http_client is not None:
+            self._client = AsyncOpenAI(
+                base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )
+        else:
+            http_client = cached_async_http_client(provider='litellm')
+            self._client = AsyncOpenAI(
+                base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )

pydantic_ai/result.py

@@ -1,15 +1,13 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
 from copy import copy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Generic, cast
+from typing import Generic, cast, overload
 
 from pydantic import ValidationError
-from typing_extensions import TypeVar
-
-from pydantic_ai._tool_manager import ToolManager
+from typing_extensions import TypeVar, deprecated
 
 from . import _utils, exceptions, messages as _messages, models
 from ._output import (
@@ -22,11 +20,14 @@ from ._output import (
     ToolOutputSchema,
 )
 from ._run_context import AgentDepsT, RunContext
+from ._tool_manager import ToolManager
 from .messages import ModelResponseStreamEvent
 from .output import (
+    DeferredToolRequests,
     OutputDataT,
     ToolOutput,
 )
+from .run import AgentRunResult
 from .usage import RunUsage, UsageLimits
 
 __all__ = (
@@ -41,7 +42,7 @@ T = TypeVar('T')
 """An invariant TypeVar."""
 
 
-@dataclass
+@dataclass(kw_only=True)
 class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _raw_stream_response: models.StreamedResponse
     _output_schema: OutputSchema[OutputDataT]
@@ -62,11 +63,11 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         async for response in self.stream_responses(debounce_by=debounce_by):
             if self._raw_stream_response.final_result_event is not None:
                 try:
-                    yield await self._validate_response(response, allow_partial=True)
+                    yield await self.validate_response_output(response, allow_partial=True)
                 except ValidationError:
                     pass
         if self._raw_stream_response.final_result_event is not None:  # pragma: no branch
-            yield await self._validate_response(self._raw_stream_response.get())
+            yield await self.validate_response_output(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."""
@@ -127,9 +128,11 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         async for _ in self:
             pass
 
-        return await self._validate_response(self._raw_stream_response.get())
+        return await self.validate_response_output(self._raw_stream_response.get())
 
-    async def _validate_response(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
+    async def validate_response_output(
+        self, message: _messages.ModelResponse, *, allow_partial: bool = False
+    ) -> OutputDataT:
         """Validate a structured result message."""
         final_result_event = self._raw_stream_response.final_result_event
         if final_result_event is None:
@@ -153,12 +156,12 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
             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:
+        elif deferred_tool_requests := _get_deferred_tool_requests(message.parts, self._tool_manager):
+            if not self._output_schema.allows_deferred_tools:
                 raise exceptions.UserError(
-                    'A deferred tool call was present, but `DeferredToolCalls` is not among output types. To resolve this, add `DeferredToolCalls` to the list of output types for this agent.'
+                    'A deferred tool call was present, but `DeferredToolRequests` is not among output types. To resolve this, add `DeferredToolRequests` to the list of output types for this agent.'
                 )
-            return cast(OutputDataT, deferred_tool_calls)
+            return cast(OutputDataT, deferred_tool_requests)
         elif isinstance(self._output_schema, TextOutputSchema):
             text = '\n\n'.join(x.content for x in message.parts if isinstance(x, _messages.TextPart))
 
@@ -231,26 +234,61 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         return self._agent_stream_iterator
 
 
-@dataclass
+@dataclass(init=False)
 class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
     """Result of a streamed run that returns structured data via a tool call."""
 
     _all_messages: list[_messages.ModelMessage]
     _new_message_index: int
 
-    _stream_response: AgentStream[AgentDepsT, OutputDataT]
-    _on_complete: Callable[[], Awaitable[None]]
+    _stream_response: AgentStream[AgentDepsT, OutputDataT] | None = None
+    _on_complete: Callable[[], Awaitable[None]] | None = None
+
+    _run_result: AgentRunResult[OutputDataT] | None = None
 
     is_complete: bool = field(default=False, init=False)
     """Whether the stream has all been received.
 
     This is set to `True` when one of
-    [`stream`][pydantic_ai.result.StreamedRunResult.stream],
+    [`stream_output`][pydantic_ai.result.StreamedRunResult.stream_output],
     [`stream_text`][pydantic_ai.result.StreamedRunResult.stream_text],
-    [`stream_structured`][pydantic_ai.result.StreamedRunResult.stream_structured] or
+    [`stream_responses`][pydantic_ai.result.StreamedRunResult.stream_responses] or
     [`get_output`][pydantic_ai.result.StreamedRunResult.get_output] completes.
     """
 
+    @overload
+    def __init__(
+        self,
+        all_messages: list[_messages.ModelMessage],
+        new_message_index: int,
+        stream_response: AgentStream[AgentDepsT, OutputDataT] | None,
+        on_complete: Callable[[], Awaitable[None]] | None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        all_messages: list[_messages.ModelMessage],
+        new_message_index: int,
+        *,
+        run_result: AgentRunResult[OutputDataT],
+    ) -> None: ...
+
+    def __init__(
+        self,
+        all_messages: list[_messages.ModelMessage],
+        new_message_index: int,
+        stream_response: AgentStream[AgentDepsT, OutputDataT] | None = None,
+        on_complete: Callable[[], Awaitable[None]] | None = None,
+        run_result: AgentRunResult[OutputDataT] | None = None,
+    ) -> None:
+        self._all_messages = all_messages
+        self._new_message_index = new_message_index
+
+        self._stream_response = stream_response
+        self._on_complete = on_complete
+        self._run_result = run_result
+
     def all_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
         """Return the history of _messages.
 
@@ -318,24 +356,35 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
             self.new_messages(output_tool_return_content=output_tool_return_content)
         )
 
+    @deprecated('`StreamedRunResult.stream` is deprecated, use `stream_output` instead.')
     async def stream(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
-        """Stream the response as an async iterable.
+        async for output in self.stream_output(debounce_by=debounce_by):
+            yield output
+
+    async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
+        """Stream the output as an async iterable.
 
         The pydantic validator for structured data will be called in
         [partial mode](https://docs.pydantic.dev/dev/concepts/experimental/#partial-validation)
         on each iteration.
 
         Args:
-            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
-                Debouncing is particularly important for long structured responses to reduce the overhead of
+            debounce_by: by how much (if at all) to debounce/group the output chunks by. `None` means no debouncing.
+                Debouncing is particularly important for long structured outputs to reduce the overhead of
                 performing validation as each token is received.
 
         Returns:
             An async iterable of the response data.
         """
-        async for output in self._stream_response.stream_output(debounce_by=debounce_by):
-            yield output
-        await self._marked_completed(self._stream_response.get())
+        if self._run_result is not None:
+            yield self._run_result.output
+            await self._marked_completed()
+        elif self._stream_response is not None:
+            async for output in self._stream_response.stream_output(debounce_by=debounce_by):
+                yield output
+            await self._marked_completed(self._stream_response.get())
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
     async def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> AsyncIterator[str]:
         """Stream the text result as an async iterable.
@@ -350,12 +399,30 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
                 Debouncing is particularly important for long structured responses to reduce the overhead of
                 performing validation as each token is received.
         """
-        async for text in self._stream_response.stream_text(delta=delta, debounce_by=debounce_by):
-            yield text
-        await self._marked_completed(self._stream_response.get())
+        if self._run_result is not None:  # pragma: no cover
+            # We can't really get here, as `_run_result` is only set in `run_stream` when `CallToolsNode` produces `DeferredToolRequests` output
+            # as a result of a tool function raising `CallDeferred` or `ApprovalRequired`.
+            # That'll change if we ever support something like `raise EndRun(output: OutputT)` where `OutputT` could be `str`.
+            if not isinstance(self._run_result.output, str):
+                raise exceptions.UserError('stream_text() can only be used with text responses')
+            yield self._run_result.output
+            await self._marked_completed()
+        elif self._stream_response is not None:
+            async for text in self._stream_response.stream_text(delta=delta, debounce_by=debounce_by):
+                yield text
+            await self._marked_completed(self._stream_response.get())
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    @deprecated('`StreamedRunResult.stream_structured` is deprecated, use `stream_responses` instead.')
     async def stream_structured(
         self, *, debounce_by: float | None = 0.1
+    ) -> AsyncIterator[tuple[_messages.ModelResponse, bool]]:
+        async for msg, last in self.stream_responses(debounce_by=debounce_by):
+            yield msg, last
+
+    async def stream_responses(
+        self, *, debounce_by: float | None = 0.1
     ) -> AsyncIterator[tuple[_messages.ModelResponse, bool]]:
         """Stream the response as an async iterable of Structured LLM Messages.
 
@@ -367,20 +434,34 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
         Returns:
             An async iterable of the structured response message and whether that is the last message.
         """
-        # if the message currently has any parts with content, yield before streaming
-        async for msg in self._stream_response.stream_responses(debounce_by=debounce_by):
-            yield msg, False
-
-        msg = self._stream_response.get()
-        yield msg, True
-
-        await self._marked_completed(msg)
+        if self._run_result is not None:
+            model_response = cast(_messages.ModelResponse, self.all_messages()[-1])
+            yield model_response, True
+            await self._marked_completed()
+        elif self._stream_response is not None:
+            # if the message currently has any parts with content, yield before streaming
+            async for msg in self._stream_response.stream_responses(debounce_by=debounce_by):
+                yield msg, False
+
+            msg = self._stream_response.get()
+            yield msg, True
+
+            await self._marked_completed(msg)
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
     async def get_output(self) -> OutputDataT:
         """Stream the whole response, validate and return it."""
-        output = await self._stream_response.get_output()
-        await self._marked_completed(self._stream_response.get())
-        return output
+        if self._run_result is not None:
+            output = self._run_result.output
+            await self._marked_completed()
+            return output
+        elif self._stream_response is not None:
+            output = await self._stream_response.get_output()
+            await self._marked_completed(self._stream_response.get())
+            return output
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
     def usage(self) -> RunUsage:
         """Return the usage of the whole run.
@@ -388,24 +469,45 @@ class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
         !!! note
             This won't return the full usage until the stream is finished.
         """
-        return self._stream_response.usage()
+        if self._run_result is not None:
+            return self._run_result.usage()
+        elif self._stream_response is not None:
+            return self._stream_response.usage()
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
-        return self._stream_response.timestamp()
+        if self._run_result is not None:
+            return self._run_result.timestamp()
+        elif self._stream_response is not None:
+            return self._stream_response.timestamp()
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    @deprecated('`validate_structured_output` is deprecated, use `validate_response_output` instead.')
     async def validate_structured_output(
         self, message: _messages.ModelResponse, *, allow_partial: bool = False
+    ) -> OutputDataT:
+        return await self.validate_response_output(message, allow_partial=allow_partial)
+
+    async def validate_response_output(
+        self, message: _messages.ModelResponse, *, allow_partial: bool = False
     ) -> OutputDataT:
         """Validate a structured result message."""
-        return await self._stream_response._validate_response(  # pyright: ignore[reportPrivateUsage]
-            message, allow_partial=allow_partial
-        )
+        if self._run_result is not None:
+            return self._run_result.output
+        elif self._stream_response is not None:
+            return await self._stream_response.validate_response_output(message, allow_partial=allow_partial)
+        else:
+            raise ValueError('No stream response or run result provided')  # pragma: no cover
 
-    async def _marked_completed(self, message: _messages.ModelResponse) -> None:
+    async def _marked_completed(self, message: _messages.ModelResponse | None = None) -> None:
         self.is_complete = True
-        self._all_messages.append(message)
-        await self._on_complete()
+        if message is not None:
+            self._all_messages.append(message)
+        if self._on_complete is not None:
+            await self._on_complete()
 
 
 @dataclass(repr=False)
@@ -414,8 +516,10 @@ class FinalResult(Generic[OutputDataT]):
 
     output: OutputDataT
     """The final result data."""
+
     tool_name: str | None = None
     """Name of the final output tool; `None` if the output came from unstructured text content."""
+
     tool_call_id: str | None = None
     """ID of the tool call that produced the final output; `None` if the output came from unstructured text content."""
 
@@ -436,9 +540,26 @@ def _get_usage_checking_stream_response(
 
         return _usage_checking_iterator()
     else:
-        # TODO: Use `return aiter(stream_response)` once we drop support for Python 3.9
-        async def _iterator():
-            async for item in stream_response:
-                yield item
+        return aiter(stream_response)
+
+
+def _get_deferred_tool_requests(
+    parts: Iterable[_messages.ModelResponsePart], tool_manager: ToolManager[AgentDepsT]
+) -> DeferredToolRequests | None:
+    """Get the deferred tool requests from the model response parts."""
+    approvals: list[_messages.ToolCallPart] = []
+    calls: list[_messages.ToolCallPart] = []
+
+    for part in parts:
+        if isinstance(part, _messages.ToolCallPart):
+            tool_def = tool_manager.get_tool_def(part.tool_name)
+            if tool_def is not None:  # pragma: no branch
+                if tool_def.kind == 'unapproved':
+                    approvals.append(part)
+                elif tool_def.kind == 'external':
+                    calls.append(part)
+
+    if not calls and not approvals:
+        return None
 
-        return _iterator()
+    return DeferredToolRequests(calls=calls, approvals=approvals)

pydantic_ai/retries.py

@@ -13,6 +13,8 @@ The module includes:
 
 from __future__ import annotations
 
+from types import TracebackType
+
 from httpx import (
     AsyncBaseTransport,
     AsyncHTTPTransport,
@@ -24,17 +26,17 @@ from httpx import (
 )
 
 try:
-    from tenacity import AsyncRetrying, RetryCallState, RetryError, Retrying, retry, wait_exponential
+    from tenacity import RetryCallState, RetryError, retry, wait_exponential
 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
 
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from datetime import datetime, timezone
 from email.utils import parsedate_to_datetime
-from typing import TYPE_CHECKING, Any, Callable, NoReturn, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from typing_extensions import TypedDict
 
@@ -134,8 +136,9 @@ class TenacityTransport(BaseTransport):
 
     Example:
         ```python
-        from httpx import Client, HTTPTransport, HTTPStatusError
-        from tenacity import stop_after_attempt, retry_if_exception_type
+        from httpx import Client, HTTPStatusError, HTTPTransport
+        from tenacity import retry_if_exception_type, stop_after_attempt
+
         from pydantic_ai.retries import RetryConfig, TenacityTransport, wait_retry_after
 
         transport = TenacityTransport(
@@ -157,18 +160,7 @@ class TenacityTransport(BaseTransport):
         config: RetryConfig,
         wrapped: BaseTransport | None = None,
         validate_response: Callable[[Response], Any] | None = None,
-        **kwargs: NoReturn,
     ):
-        # TODO: Remove the following checks (and **kwargs) during v1 release
-        if 'controller' in kwargs:  # pragma: no cover
-            raise TypeError('The `controller` argument has been renamed to `config`, and now requires a `RetryConfig`.')
-        if kwargs:  # pragma: no cover
-            raise TypeError(f'Unexpected keyword arguments: {", ".join(kwargs)}')
-        if isinstance(config, Retrying):  # pragma: no cover
-            raise ValueError(
-                'Passing a Retrying instance is no longer supported; the `config` argument must be a `pydantic_ai.retries.RetryConfig`.'
-            )
-
         self.config = config
         self.wrapped = wrapped or HTTPTransport()
         self.validate_response = validate_response
@@ -195,11 +187,30 @@ class TenacityTransport(BaseTransport):
             response.request = req
 
             if self.validate_response:
-                self.validate_response(response)
+                try:
+                    self.validate_response(response)
+                except Exception:
+                    response.close()
+                    raise
             return response
 
         return handle_request(request)
 
+    def __enter__(self) -> TenacityTransport:
+        self.wrapped.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        self.wrapped.__exit__(exc_type, exc_value, traceback)
+
+    def close(self) -> None:
+        self.wrapped.close()  # pragma: no cover
+
 
 class AsyncTenacityTransport(AsyncBaseTransport):
     """Asynchronous HTTP transport with tenacity-based retry functionality.
@@ -224,7 +235,8 @@ class AsyncTenacityTransport(AsyncBaseTransport):
     Example:
         ```python
         from httpx import AsyncClient, HTTPStatusError
-        from tenacity import stop_after_attempt, retry_if_exception_type
+        from tenacity import retry_if_exception_type, stop_after_attempt
+
         from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
 
         transport = AsyncTenacityTransport(
@@ -245,18 +257,7 @@ class AsyncTenacityTransport(AsyncBaseTransport):
         config: RetryConfig,
         wrapped: AsyncBaseTransport | None = None,
         validate_response: Callable[[Response], Any] | None = None,
-        **kwargs: NoReturn,
     ):
-        # TODO: Remove the following checks (and **kwargs) during v1 release
-        if 'controller' in kwargs:  # pragma: no cover
-            raise TypeError('The `controller` argument has been renamed to `config`, and now requires a `RetryConfig`.')
-        if kwargs:  # pragma: no cover
-            raise TypeError(f'Unexpected keyword arguments: {", ".join(kwargs)}')
-        if isinstance(config, AsyncRetrying):  # pragma: no cover
-            raise ValueError(
-                'Passing an AsyncRetrying instance is no longer supported; the `config` argument must be a `pydantic_ai.retries.RetryConfig`.'
-            )
-
         self.config = config
         self.wrapped = wrapped or AsyncHTTPTransport()
         self.validate_response = validate_response
@@ -283,11 +284,30 @@ class AsyncTenacityTransport(AsyncBaseTransport):
             response.request = req
 
             if self.validate_response:
-                self.validate_response(response)
+                try:
+                    self.validate_response(response)
+                except Exception:
+                    await response.aclose()
+                    raise
             return response
 
         return await handle_async_request(request)
 
+    async def __aenter__(self) -> AsyncTenacityTransport:
+        await self.wrapped.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        await self.wrapped.__aexit__(exc_type, exc_value, traceback)
+
+    async def aclose(self) -> None:
+        await self.wrapped.aclose()
+
 
 def wait_retry_after(
     fallback_strategy: Callable[[RetryCallState], float] | None = None, max_wait: float = 300
@@ -314,7 +334,8 @@ def wait_retry_after(
     Example:
         ```python
         from httpx import AsyncClient, HTTPStatusError
-        from tenacity import stop_after_attempt, retry_if_exception_type
+        from tenacity import retry_if_exception_type, stop_after_attempt
+
         from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
 
         transport = AsyncTenacityTransport(