pydantic-ai-slim 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

pydantic_ai/result.py

@@ -1,19 +1,20 @@
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic, TypeVar, cast
 
 import logfire_api
 
-from . import _result, _utils, exceptions, messages, models
-from .tools import AgentDeps
+from . import _result, _utils, exceptions, messages as _messages, models
+from .settings import UsageLimits
+from .tools import AgentDeps, RunContext
 
 __all__ = (
     'ResultData',
-    'Cost',
+    'Usage',
     'RunResult',
     'StreamedRunResult',
 )
@@ -26,30 +27,32 @@ _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
 
 @dataclass
-class Cost:
-    """Cost of a request or run.
+class Usage:
+    """LLM usage associated to a request or run.
 
-    Responsibility for calculating costs is on the model used, PydanticAI simply sums the cost of requests.
+    Responsibility for calculating usage is on the model; PydanticAI simply sums the usage information across requests.
 
-    You'll need to look up the documentation of the model you're using to convent "token count" costs to monetary costs.
+    You'll need to look up the documentation of the model you're using to convert usage to monetary costs.
     """
 
+    requests: int = 0
+    """Number of requests made."""
     request_tokens: int | None = None
-    """Tokens used in processing the request."""
+    """Tokens used in processing requests."""
     response_tokens: int | None = None
-    """Tokens used in generating the response."""
+    """Tokens used in generating responses."""
     total_tokens: int | None = None
     """Total tokens used in the whole run, should generally be equal to `request_tokens + response_tokens`."""
     details: dict[str, int] | None = None
     """Any extra details returned by the model."""
 
-    def __add__(self, other: Cost) -> Cost:
-        """Add two costs together.
+    def __add__(self, other: Usage) -> Usage:
+        """Add two Usages together.
 
-        This is provided so it's trivial to sum costs from multiple requests and runs.
+        This is provided so it's trivial to sum usage information from multiple requests and runs.
         """
         counts: dict[str, int] = {}
-        for f in 'request_tokens', 'response_tokens', 'total_tokens':
+        for f in 'requests', 'request_tokens', 'response_tokens', 'total_tokens':
             self_value = getattr(self, f)
             other_value = getattr(other, f)
             if self_value is not None or other_value is not None:
@@ -61,7 +64,7 @@ class Cost:
             for key, value in other.details.items():
                 details[key] = details.get(key, 0) + value
 
-        return Cost(**counts, details=details or None)
+        return Usage(**counts, details=details or None)
 
 
 @dataclass
@@ -71,19 +74,19 @@ class _BaseRunResult(ABC, Generic[ResultData]):
     You should not import or use this type directly, instead use its subclasses `RunResult` and `StreamedRunResult`.
     """
 
-    _all_messages: list[messages.Message]
+    _all_messages: list[_messages.ModelMessage]
     _new_message_index: int
 
-    def all_messages(self) -> list[messages.Message]:
-        """Return the history of messages."""
+    def all_messages(self) -> list[_messages.ModelMessage]:
+        """Return the history of _messages."""
         # this is a method to be consistent with the other methods
         return self._all_messages
 
     def all_messages_json(self) -> bytes:
-        """Return all messages from [`all_messages`][..all_messages] as JSON bytes."""
-        return messages.MessagesTypeAdapter.dump_json(self.all_messages())
+        """Return all messages from [`all_messages`][pydantic_ai.result._BaseRunResult.all_messages] as JSON bytes."""
+        return _messages.ModelMessagesTypeAdapter.dump_json(self.all_messages())
 
-    def new_messages(self) -> list[messages.Message]:
+    def new_messages(self) -> list[_messages.ModelMessage]:
         """Return new messages associated with this run.
 
         System prompts and any messages from older runs are excluded.
@@ -91,11 +94,11 @@ class _BaseRunResult(ABC, Generic[ResultData]):
         return self.all_messages()[self._new_message_index :]
 
     def new_messages_json(self) -> bytes:
-        """Return new messages from [`new_messages`][..new_messages] as JSON bytes."""
-        return messages.MessagesTypeAdapter.dump_json(self.new_messages())
+        """Return new messages from [`new_messages`][pydantic_ai.result._BaseRunResult.new_messages] as JSON bytes."""
+        return _messages.ModelMessagesTypeAdapter.dump_json(self.new_messages())
 
     @abstractmethod
-    def cost(self) -> Cost:
+    def usage(self) -> Usage:
         raise NotImplementedError()
 
 
@@ -105,24 +108,26 @@ class RunResult(_BaseRunResult[ResultData]):
 
     data: ResultData
     """Data from the final response in the run."""
-    _cost: Cost
+    _usage: Usage
 
-    def cost(self) -> Cost:
-        """Return the cost of the whole run."""
-        return self._cost
+    def usage(self) -> Usage:
+        """Return the usage of the whole run."""
+        return self._usage
 
 
 @dataclass
 class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultData]):
     """Result of a streamed run that returns structured data via a tool call."""
 
-    cost_so_far: Cost
-    """Cost of the run up until the last request."""
+    usage_so_far: Usage
+    """Usage of the run up until the last request."""
+    _usage_limits: UsageLimits | None
     _stream_response: models.EitherStreamedResponse
     _result_schema: _result.ResultSchema[ResultData] | None
-    _deps: AgentDeps
+    _run_ctx: RunContext[AgentDeps]
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]]
-    _on_complete: Callable[[list[messages.Message]], None]
+    _result_tool_name: str | None
+    _on_complete: Callable[[], Awaitable[None]]
     is_complete: bool = field(default=False, init=False)
     """Whether the stream has all been received.
 
@@ -172,11 +177,15 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                 Debouncing is particularly important for long structured responses to reduce the overhead of
                 performing validation as each token is received.
         """
+        usage_checking_stream = _get_usage_checking_stream_response(
+            self._stream_response, self._usage_limits, self.usage
+        )
+
         with _logfire.span('response stream text') as lf_span:
             if isinstance(self._stream_response, models.StreamStructuredResponse):
                 raise exceptions.UserError('stream_text() can only be used with text responses')
             if delta:
-                async with _utils.group_by_temporal(self._stream_response, debounce_by) as group_iter:
+                async with _utils.group_by_temporal(usage_checking_stream, debounce_by) as group_iter:
                     async for _ in group_iter:
                         yield ''.join(self._stream_response.get())
                 final_delta = ''.join(self._stream_response.get(final=True))
@@ -187,7 +196,7 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                 # yielding at each step
                 chunks: list[str] = []
                 combined = ''
-                async with _utils.group_by_temporal(self._stream_response, debounce_by) as group_iter:
+                async with _utils.group_by_temporal(usage_checking_stream, debounce_by) as group_iter:
                     async for _ in group_iter:
                         new = False
                         for chunk in self._stream_response.get():
@@ -205,11 +214,11 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
                     combined = await self._validate_text_result(''.join(chunks))
                     yield combined
                 lf_span.set_attribute('combined_text', combined)
-                self._marked_completed(text=combined)
+                await self._marked_completed(_messages.ModelResponse.from_text(combined))
 
     async def stream_structured(
         self, *, debounce_by: float | None = 0.1
-    ) -> AsyncIterator[tuple[messages.ModelStructuredResponse, bool]]:
+    ) -> AsyncIterator[tuple[_messages.ModelResponse, bool]]:
         """Stream the response as an async iterable of Structured LLM Messages.
 
         !!! note
@@ -224,61 +233,75 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         Returns:
             An async iterable of the structured response message and whether that is the last message.
         """
+        usage_checking_stream = _get_usage_checking_stream_response(
+            self._stream_response, self._usage_limits, self.usage
+        )
+
         with _logfire.span('response stream structured') as lf_span:
             if isinstance(self._stream_response, models.StreamTextResponse):
                 raise exceptions.UserError('stream_structured() can only be used with structured responses')
             else:
                 # we should already have a message at this point, yield that first if it has any content
                 msg = self._stream_response.get()
-                if any(call.has_content() for call in msg.calls):
-                    yield msg, False
-                async with _utils.group_by_temporal(self._stream_response, debounce_by) as group_iter:
+                for item in msg.parts:
+                    if isinstance(item, _messages.ToolCallPart) and item.has_content():
+                        yield msg, False
+                        break
+                async with _utils.group_by_temporal(usage_checking_stream, debounce_by) as group_iter:
                     async for _ in group_iter:
                         msg = self._stream_response.get()
-                        if any(call.has_content() for call in msg.calls):
-                            yield msg, False
+                        for item in msg.parts:
+                            if isinstance(item, _messages.ToolCallPart) and item.has_content():
+                                yield msg, False
+                                break
                 msg = self._stream_response.get(final=True)
                 yield msg, True
                 lf_span.set_attribute('structured_response', msg)
-                self._marked_completed(structured_message=msg)
+                await self._marked_completed(msg)
 
     async def get_data(self) -> ResultData:
         """Stream the whole response, validate and return it."""
-        async for _ in self._stream_response:
+        usage_checking_stream = _get_usage_checking_stream_response(
+            self._stream_response, self._usage_limits, self.usage
+        )
+
+        async for _ in usage_checking_stream:
             pass
+
         if isinstance(self._stream_response, models.StreamTextResponse):
             text = ''.join(self._stream_response.get(final=True))
             text = await self._validate_text_result(text)
-            self._marked_completed(text=text)
+            await self._marked_completed(_messages.ModelResponse.from_text(text))
             return cast(ResultData, text)
         else:
-            structured_message = self._stream_response.get(final=True)
-            self._marked_completed(structured_message=structured_message)
-            return await self.validate_structured_result(structured_message)
+            message = self._stream_response.get(final=True)
+            await self._marked_completed(message)
+            return await self.validate_structured_result(message)
 
     @property
     def is_structured(self) -> bool:
         """Return whether the stream response contains structured data (as opposed to text)."""
         return isinstance(self._stream_response, models.StreamStructuredResponse)
 
-    def cost(self) -> Cost:
-        """Return the cost of the whole run.
+    def usage(self) -> Usage:
+        """Return the usage of the whole run.
 
         !!! note
-            This won't return the full cost until the stream is finished.
+            This won't return the full usage until the stream is finished.
         """
-        return self.cost_so_far + self._stream_response.cost()
+        return self.usage_so_far + self._stream_response.usage()
 
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
         return self._stream_response.timestamp()
 
     async def validate_structured_result(
-        self, message: messages.ModelStructuredResponse, *, allow_partial: bool = False
+        self, message: _messages.ModelResponse, *, allow_partial: bool = False
     ) -> ResultData:
         """Validate a structured result message."""
         assert self._result_schema is not None, 'Expected _result_schema to not be None'
-        match = self._result_schema.find_tool(message)
+        assert self._result_tool_name is not None, 'Expected _result_tool_name to not be None'
+        match = self._result_schema.find_named_tool(message.parts, self._result_tool_name)
         if match is None:
             raise exceptions.UnexpectedModelBehavior(
                 f'Invalid message, unable to find tool: {self._result_schema.tool_names()}'
@@ -288,29 +311,34 @@ class StreamedRunResult(_BaseRunResult[ResultData], Generic[AgentDeps, ResultDat
         result_data = result_tool.validate(call, allow_partial=allow_partial, wrap_validation_errors=False)
 
         for validator in self._result_validators:
-            result_data = await validator.validate(result_data, self._deps, 0, call)
+            result_data = await validator.validate(result_data, call, self._run_ctx)
         return result_data
 
     async def _validate_text_result(self, text: str) -> str:
         for validator in self._result_validators:
             text = await validator.validate(  # pyright: ignore[reportAssignmentType]
                 text,  # pyright: ignore[reportArgumentType]
-                self._deps,
-                0,
                 None,
+                self._run_ctx,
             )
         return text
 
-    def _marked_completed(
-        self, *, text: str | None = None, structured_message: messages.ModelStructuredResponse | None = None
-    ) -> None:
+    async def _marked_completed(self, message: _messages.ModelResponse) -> None:
         self.is_complete = True
-        if text is not None:
-            assert structured_message is None, 'Either text or structured_message should provided, not both'
-            self._all_messages.append(
-                messages.ModelTextResponse(content=text, timestamp=self._stream_response.timestamp())
-            )
-        else:
-            assert structured_message is not None, 'Either text or structured_message should provided, not both'
-            self._all_messages.append(structured_message)
-        self._on_complete(self._all_messages)
+        self._all_messages.append(message)
+        await self._on_complete()
+
+
+def _get_usage_checking_stream_response(
+    stream_response: AsyncIterator[ResultData], limits: UsageLimits | None, get_usage: Callable[[], Usage]
+) -> AsyncIterator[ResultData]:
+    if limits is not None and limits.has_token_limits():
+
+        async def _usage_checking_iterator():
+            async for item in stream_response:
+                limits.check_tokens(get_usage())
+                yield item
+
+        return _usage_checking_iterator()
+    else:
+        return stream_response

pydantic_ai/settings.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from httpx import Timeout
+from typing_extensions import TypedDict
+
+from .exceptions import UsageLimitExceeded
+
+if TYPE_CHECKING:
+    from .result import Usage
+
+
+class ModelSettings(TypedDict, total=False):
+    """Settings to configure an LLM.
+
+    Here we include only settings which apply to multiple models / model providers.
+    """
+
+    max_tokens: int
+    """The maximum number of tokens to generate before stopping.
+
+    Supported by:
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    temperature: float
+    """Amount of randomness injected into the response.
+
+    Use `temperature` closer to `0.0` for analytical / multiple choice, and closer to a model's
+    maximum `temperature` for creative and generative tasks.
+
+    Note that even with `temperature` of `0.0`, the results will not be fully deterministic.
+
+    Supported by:
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    top_p: float
+    """An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass.
+
+    So 0.1 means only the tokens comprising the top 10% probability mass are considered.
+
+    You should either alter `temperature` or `top_p`, but not both.
+
+    Supported by:
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+    timeout: float | Timeout
+    """Override the client-level default timeout for a request, in seconds.
+
+    Supported by:
+    * Gemini
+    * Anthropic
+    * OpenAI
+    * Groq
+    """
+
+
+def merge_model_settings(base: ModelSettings | None, overrides: ModelSettings | None) -> ModelSettings | None:
+    """Merge two sets of model settings, preferring the overrides.
+
+    A common use case is: merge_model_settings(<agent settings>, <run settings>)
+    """
+    # Note: we may want merge recursively if/when we add non-primitive values
+    if base and overrides:
+        return base | overrides
+    else:
+        return base or overrides
+
+
+@dataclass
+class UsageLimits:
+    """Limits on model usage.
+
+    The request count is tracked by pydantic_ai, and the request limit is checked before each request to the model.
+    Token counts are provided in responses from the model, and the token limits are checked after each response.
+
+    Each of the limits can be set to `None` to disable that limit.
+    """
+
+    request_limit: int | None = 50
+    """The maximum number of requests allowed to the model."""
+    request_tokens_limit: int | None = None
+    """The maximum number of tokens allowed in requests to the model."""
+    response_tokens_limit: int | None = None
+    """The maximum number of tokens allowed in responses from the model."""
+    total_tokens_limit: int | None = None
+    """The maximum number of tokens allowed in requests and responses combined."""
+
+    def has_token_limits(self) -> bool:
+        """Returns `True` if this instance places any limits on token counts.
+
+        If this returns `False`, the `check_tokens` method will never raise an error.
+
+        This is useful because if we have token limits, we need to check them after receiving each streamed message.
+        If there are no limits, we can skip that processing in the streaming response iterator.
+        """
+        return any(
+            limit is not None
+            for limit in (self.request_tokens_limit, self.response_tokens_limit, self.total_tokens_limit)
+        )
+
+    def check_before_request(self, usage: Usage) -> None:
+        """Raises a `UsageLimitExceeded` exception if the next request would exceed the request_limit."""
+        request_limit = self.request_limit
+        if request_limit is not None and usage.requests >= request_limit:
+            raise UsageLimitExceeded(f'The next request would exceed the request_limit of {request_limit}')
+
+    def check_tokens(self, usage: Usage) -> None:
+        """Raises a `UsageLimitExceeded` exception if the usage exceeds any of the token limits."""
+        request_tokens = usage.request_tokens or 0
+        if self.request_tokens_limit is not None and request_tokens > self.request_tokens_limit:
+            raise UsageLimitExceeded(
+                f'Exceeded the request_tokens_limit of {self.request_tokens_limit} ({request_tokens=})'
+            )
+
+        response_tokens = usage.response_tokens or 0
+        if self.response_tokens_limit is not None and response_tokens > self.response_tokens_limit:
+            raise UsageLimitExceeded(
+                f'Exceeded the response_tokens_limit of {self.response_tokens_limit} ({response_tokens=})'
+            )
+
+        total_tokens = request_tokens + response_tokens
+        if self.total_tokens_limit is not None and total_tokens > self.total_tokens_limit:
+            raise UsageLimitExceeded(f'Exceeded the total_tokens_limit of {self.total_tokens_limit} ({total_tokens=})')

pydantic_ai/tools.py

@@ -1,23 +1,18 @@
 from __future__ import annotations as _annotations
 
+import dataclasses
 import inspect
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast
+from typing import Any, Callable, Generic, TypeVar, Union, cast
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
 from typing_extensions import Concatenate, ParamSpec, TypeAlias
 
-from . import _pydantic, _utils, messages
+from . import _pydantic, _utils, messages as _messages, models
 from .exceptions import ModelRetry, UnexpectedModelBehavior
 
-if TYPE_CHECKING:
-    from .result import ResultData
-else:
-    ResultData = Any
-
-
 __all__ = (
     'AgentDeps',
     'RunContext',
@@ -37,7 +32,7 @@ AgentDeps = TypeVar('AgentDeps')
 """Type variable for agent dependencies."""
 
 
-@dataclass
+@dataclasses.dataclass
 class RunContext(Generic[AgentDeps]):
     """Information about the current call."""
 
@@ -45,8 +40,23 @@ class RunContext(Generic[AgentDeps]):
     """Dependencies for the agent."""
     retry: int
     """Number of retries so far."""
-    tool_name: str | None = None
+    messages: list[_messages.ModelMessage]
+    """Messages exchanged in the conversation so far."""
+    tool_name: str | None
     """Name of the tool being called."""
+    model: models.Model
+    """The model used in this run."""
+
+    def replace_with(
+        self, retry: int | None = None, tool_name: str | None | _utils.Unset = _utils.UNSET
+    ) -> RunContext[AgentDeps]:
+        # Create a new `RunContext` a new `retry` value and `tool_name`.
+        kwargs = {}
+        if retry is not None:
+            kwargs['retry'] = retry
+        if tool_name is not _utils.UNSET:
+            kwargs['tool_name'] = tool_name
+        return dataclasses.replace(self, **kwargs)
 
 
 ToolParams = ParamSpec('ToolParams')
@@ -63,6 +73,8 @@ SystemPromptFunc = Union[
 Usage `SystemPromptFunc[AgentDeps]`.
 """
 
+ResultData = TypeVar('ResultData')
+
 ResultValidatorFunc = Union[
     Callable[[RunContext[AgentDeps], ResultData], ResultData],
     Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]],
@@ -87,7 +99,7 @@ ToolFuncPlain = Callable[ToolParams, Any]
 Usage `ToolPlainFunc[ToolParams]`.
 """
 ToolFuncEither = Union[ToolFuncContext[AgentDeps, ToolParams], ToolFuncPlain[ToolParams]]
-"""Either kind of tool function.
+"""Either part_kind of tool function.
 
 This is just a union of [`ToolFuncContext`][pydantic_ai.tools.ToolFuncContext] and
 [`ToolFuncPlain`][pydantic_ai.tools.ToolFuncPlain].
@@ -97,11 +109,11 @@ Usage `ToolFuncEither[AgentDeps, ToolParams]`.
 ToolPrepareFunc: TypeAlias = 'Callable[[RunContext[AgentDeps], ToolDefinition], Awaitable[ToolDefinition | None]]'
 """Definition of a function that can prepare a tool definition at call time.
 
-See [tool docs](../agents.md#tool-prepare) for more information.
+See [tool docs](../tools.md#tool-prepare) for more information.
 
 Example — here `only_if_42` is valid as a `ToolPrepareFunc`:
 
-```py
+```python {lint="not-imports"}
 from typing import Union
 
 from pydantic_ai import RunContext, Tool
@@ -157,7 +169,7 @@ class Tool(Generic[AgentDeps]):
 
         Example usage:
 
-        ```py
+        ```python {lint="not-imports"}
         from pydantic_ai import Agent, RunContext, Tool
 
         async def my_tool(ctx: RunContext[int], x: int, y: int) -> str:
@@ -168,7 +180,7 @@ class Tool(Generic[AgentDeps]):
 
         or with a custom prepare method:
 
-        ```py
+        ```python {lint="not-imports"}
         from typing import Union
 
         from pydantic_ai import Agent, RunContext, Tool
@@ -235,17 +247,19 @@ class Tool(Generic[AgentDeps]):
         else:
             return tool_def
 
-    async def run(self, deps: AgentDeps, message: messages.ToolCall) -> messages.Message:
+    async def run(
+        self, message: _messages.ToolCallPart, run_context: RunContext[AgentDeps]
+    ) -> _messages.ModelRequestPart:
         """Run the tool function asynchronously."""
         try:
-            if isinstance(message.args, messages.ArgsJson):
+            if isinstance(message.args, _messages.ArgsJson):
                 args_dict = self._validator.validate_json(message.args.args_json)
             else:
                 args_dict = self._validator.validate_python(message.args.args_dict)
         except ValidationError as e:
             return self._on_error(e, message)
 
-        args, kwargs = self._call_args(deps, args_dict, message)
+        args, kwargs = self._call_args(args_dict, message, run_context)
         try:
             if self._is_async:
                 function = cast(Callable[[Any], Awaitable[str]], self.function)
@@ -257,19 +271,23 @@ class Tool(Generic[AgentDeps]):
             return self._on_error(e, message)
 
         self.current_retry = 0
-        return messages.ToolReturn(
+        return _messages.ToolReturnPart(
             tool_name=message.tool_name,
             content=response_content,
-            tool_id=message.tool_id,
+            tool_call_id=message.tool_call_id,
         )
 
     def _call_args(
-        self, deps: AgentDeps, args_dict: dict[str, Any], message: messages.ToolCall
+        self,
+        args_dict: dict[str, Any],
+        message: _messages.ToolCallPart,
+        run_context: RunContext[AgentDeps],
     ) -> tuple[list[Any], dict[str, Any]]:
         if self._single_arg_name:
             args_dict = {self._single_arg_name: args_dict}
 
-        args = [RunContext(deps, self.current_retry, message.tool_name)] if self.takes_ctx else []
+        ctx = dataclasses.replace(run_context, retry=self.current_retry, tool_name=message.tool_name)
+        args = [ctx] if self.takes_ctx else []
         for positional_field in self._positional_fields:
             args.append(args_dict.pop(positional_field))
         if self._var_positional_field:
@@ -277,7 +295,9 @@ class Tool(Generic[AgentDeps]):
 
         return args, args_dict
 
-    def _on_error(self, exc: ValidationError | ModelRetry, call_message: messages.ToolCall) -> messages.RetryPrompt:
+    def _on_error(
+        self, exc: ValidationError | ModelRetry, call_message: _messages.ToolCallPart
+    ) -> _messages.RetryPromptPart:
         self.current_retry += 1
         if self.max_retries is None or self.current_retry > self.max_retries:
             raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {self.max_retries}') from exc
@@ -286,10 +306,10 @@ class Tool(Generic[AgentDeps]):
                 content = exc.errors(include_url=False)
             else:
                 content = exc.message
-            return messages.RetryPrompt(
+            return _messages.RetryPromptPart(
                 tool_name=call_message.tool_name,
                 content=content,
-                tool_id=call_message.tool_id,
+                tool_call_id=call_message.tool_call_id,
             )
 
 
@@ -298,7 +318,7 @@ ObjectJsonSchema: TypeAlias = dict[str, Any]
 
 This type is used to define tools parameters (aka arguments) in [ToolDefinition][pydantic_ai.tools.ToolDefinition].
 
-With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_items=Any`
+With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_parts=Any`
 """