pydantic-ai-slim 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

pydantic_ai/agent.py

@@ -22,10 +22,10 @@ from . import (
     result,
     usage as _usage,
 )
-from .result import ResultData
+from .result import ResultDataT
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
-    AgentDeps,
+    AgentDepsT,
     DocstringFormat,
     RunContext,
     Tool,
@@ -51,6 +51,8 @@ else:
 
     logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
+T = TypeVar('T')
+"""An invariant TypeVar."""
 NoneType = type(None)
 EndStrategy = Literal['early', 'exhaustive']
 """The strategy for handling multiple tool calls when a final result is found.
@@ -58,17 +60,17 @@ EndStrategy = Literal['early', 'exhaustive']
 - `'early'`: Stop processing other tool calls once a final result is found
 - `'exhaustive'`: Process all tool calls even after finding a final result
 """
-RunResultData = TypeVar('RunResultData')
+RunResultDataT = TypeVar('RunResultDataT')
 """Type variable for the result data of a run where `result_type` was customized on the run call."""
 
 
 @final
 @dataclasses.dataclass(init=False)
-class Agent(Generic[AgentDeps, ResultData]):
+class Agent(Generic[AgentDepsT, ResultDataT]):
     """Class for defining "agents" - a way to have a specific type of "conversation" with an LLM.
 
-    Agents are generic in the dependency type they take [`AgentDeps`][pydantic_ai.tools.AgentDeps]
-    and the result data type they return, [`ResultData`][pydantic_ai.result.ResultData].
+    Agents are generic in the dependency type they take [`AgentDepsT`][pydantic_ai.tools.AgentDepsT]
+    and the result data type they return, [`ResultDataT`][pydantic_ai.result.ResultDataT].
 
     By default, if neither generic parameter is customised, agents have type `Agent[None, str]`.
 
@@ -104,34 +106,34 @@ class Agent(Generic[AgentDeps, ResultData]):
     """
     _result_tool_name: str = dataclasses.field(repr=False)
     _result_tool_description: str | None = dataclasses.field(repr=False)
-    _result_schema: _result.ResultSchema[ResultData] | None = dataclasses.field(repr=False)
-    _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = dataclasses.field(repr=False)
+    _result_schema: _result.ResultSchema[ResultDataT] | None = dataclasses.field(repr=False)
+    _result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]] = dataclasses.field(repr=False)
     _system_prompts: tuple[str, ...] = dataclasses.field(repr=False)
-    _function_tools: dict[str, Tool[AgentDeps]] = dataclasses.field(repr=False)
+    _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
-    _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDeps]] = dataclasses.field(repr=False)
-    _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDeps]] = dataclasses.field(
+    _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(repr=False)
+    _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(
         repr=False
     )
-    _deps_type: type[AgentDeps] = dataclasses.field(repr=False)
+    _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
-    _override_deps: _utils.Option[AgentDeps] = dataclasses.field(default=None, repr=False)
+    _override_deps: _utils.Option[AgentDepsT] = dataclasses.field(default=None, repr=False)
     _override_model: _utils.Option[models.Model] = dataclasses.field(default=None, repr=False)
 
     def __init__(
         self,
         model: models.Model | models.KnownModelName | None = None,
         *,
-        result_type: type[ResultData] = str,
+        result_type: type[ResultDataT] = str,
         system_prompt: str | Sequence[str] = (),
-        deps_type: type[AgentDeps] = NoneType,
+        deps_type: type[AgentDepsT] = NoneType,
         name: str | None = None,
         model_settings: ModelSettings | None = None,
         retries: int = 1,
         result_tool_name: str = 'final_result',
         result_tool_description: str | None = None,
         result_retries: int | None = None,
-        tools: Sequence[Tool[AgentDeps] | ToolFuncEither[AgentDeps, ...]] = (),
+        tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
     ):
@@ -200,27 +202,27 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[ResultData]: ...
+    ) -> result.RunResult[ResultDataT]: ...
 
     @overload
     async def run(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     async def run(
         self,
@@ -228,11 +230,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[Any]:
         """Run the agent with a user prompt in async mode.
@@ -338,36 +340,36 @@ class Agent(Generic[AgentDeps, ResultData]):
         *,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[ResultData]: ...
+    ) -> result.RunResult[ResultDataT]: ...
 
     @overload
     def run_sync(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None,
+        result_type: type[RunResultDataT] | None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     def run_sync(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
@@ -428,42 +430,42 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDeps, ResultData]]: ...
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, ResultDataT]]: ...
 
     @overload
     def run_stream(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDeps, RunResultData]]: ...
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunResultDataT]]: ...
 
     @asynccontextmanager
     async def run_stream(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
-        deps: AgentDeps = None,
+        deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AsyncIterator[result.StreamedRunResult[AgentDeps, Any]]:
+    ) -> AsyncIterator[result.StreamedRunResult[AgentDepsT, Any]]:
         """Run the agent with a user prompt in async mode, returning a streamed response.
 
         Example:
@@ -560,10 +562,19 @@ class Agent(Generic[AgentDeps, ResultData]):
                                     parts = await self._process_function_tools(
                                         tool_calls, result_tool_name, run_context, result_schema
                                     )
+                                    if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
+                                        self._incr_result_retry(run_context)
                                     if parts:
                                         messages.append(_messages.ModelRequest(parts))
                                     run_span.set_attribute('all_messages', messages)
 
+                                # The following is not guaranteed to be true, but we consider it a user error if
+                                # there are result validators that might convert the result data from an overridden
+                                # `result_type` to a type that is not valid as such.
+                                result_validators = cast(
+                                    list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators
+                                )
+
                                 yield result.StreamedRunResult(
                                     messages,
                                     new_message_index,
@@ -571,7 +582,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                                     result_stream,
                                     result_schema,
                                     run_context,
-                                    self._result_validators,
+                                    result_validators,
                                     result_tool_name,
                                     on_complete,
                                 )
@@ -597,7 +608,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     def override(
         self,
         *,
-        deps: AgentDeps | _utils.Unset = _utils.UNSET,
+        deps: AgentDepsT | _utils.Unset = _utils.UNSET,
         model: models.Model | models.KnownModelName | _utils.Unset = _utils.UNSET,
     ) -> Iterator[None]:
         """Context manager to temporarily override agent dependencies and model.
@@ -633,13 +644,13 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     @overload
     def system_prompt(
-        self, func: Callable[[RunContext[AgentDeps]], str], /
-    ) -> Callable[[RunContext[AgentDeps]], str]: ...
+        self, func: Callable[[RunContext[AgentDepsT]], str], /
+    ) -> Callable[[RunContext[AgentDepsT]], str]: ...
 
     @overload
     def system_prompt(
-        self, func: Callable[[RunContext[AgentDeps]], Awaitable[str]], /
-    ) -> Callable[[RunContext[AgentDeps]], Awaitable[str]]: ...
+        self, func: Callable[[RunContext[AgentDepsT]], Awaitable[str]], /
+    ) -> Callable[[RunContext[AgentDepsT]], Awaitable[str]]: ...
 
     @overload
     def system_prompt(self, func: Callable[[], str], /) -> Callable[[], str]: ...
@@ -650,17 +661,17 @@ class Agent(Generic[AgentDeps, ResultData]):
     @overload
     def system_prompt(
         self, /, *, dynamic: bool = False
-    ) -> Callable[[_system_prompt.SystemPromptFunc[AgentDeps]], _system_prompt.SystemPromptFunc[AgentDeps]]: ...
+    ) -> Callable[[_system_prompt.SystemPromptFunc[AgentDepsT]], _system_prompt.SystemPromptFunc[AgentDepsT]]: ...
 
     def system_prompt(
         self,
-        func: _system_prompt.SystemPromptFunc[AgentDeps] | None = None,
+        func: _system_prompt.SystemPromptFunc[AgentDepsT] | None = None,
         /,
         *,
         dynamic: bool = False,
     ) -> (
-        Callable[[_system_prompt.SystemPromptFunc[AgentDeps]], _system_prompt.SystemPromptFunc[AgentDeps]]
-        | _system_prompt.SystemPromptFunc[AgentDeps]
+        Callable[[_system_prompt.SystemPromptFunc[AgentDepsT]], _system_prompt.SystemPromptFunc[AgentDepsT]]
+        | _system_prompt.SystemPromptFunc[AgentDepsT]
     ):
         """Decorator to register a system prompt function.
 
@@ -696,9 +707,9 @@ class Agent(Generic[AgentDeps, ResultData]):
         if func is None:
 
             def decorator(
-                func_: _system_prompt.SystemPromptFunc[AgentDeps],
-            ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
-                runner = _system_prompt.SystemPromptRunner(func_, dynamic=dynamic)
+                func_: _system_prompt.SystemPromptFunc[AgentDepsT],
+            ) -> _system_prompt.SystemPromptFunc[AgentDepsT]:
+                runner = _system_prompt.SystemPromptRunner[AgentDepsT](func_, dynamic=dynamic)
                 self._system_prompt_functions.append(runner)
                 if dynamic:
                     self._system_prompt_dynamic_functions[func_.__qualname__] = runner
@@ -712,25 +723,27 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     @overload
     def result_validator(
-        self, func: Callable[[RunContext[AgentDeps], ResultData], ResultData], /
-    ) -> Callable[[RunContext[AgentDeps], ResultData], ResultData]: ...
+        self, func: Callable[[RunContext[AgentDepsT], ResultDataT], ResultDataT], /
+    ) -> Callable[[RunContext[AgentDepsT], ResultDataT], ResultDataT]: ...
 
     @overload
     def result_validator(
-        self, func: Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]], /
-    ) -> Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]]: ...
+        self, func: Callable[[RunContext[AgentDepsT], ResultDataT], Awaitable[ResultDataT]], /
+    ) -> Callable[[RunContext[AgentDepsT], ResultDataT], Awaitable[ResultDataT]]: ...
 
     @overload
-    def result_validator(self, func: Callable[[ResultData], ResultData], /) -> Callable[[ResultData], ResultData]: ...
+    def result_validator(
+        self, func: Callable[[ResultDataT], ResultDataT], /
+    ) -> Callable[[ResultDataT], ResultDataT]: ...
 
     @overload
     def result_validator(
-        self, func: Callable[[ResultData], Awaitable[ResultData]], /
-    ) -> Callable[[ResultData], Awaitable[ResultData]]: ...
+        self, func: Callable[[ResultDataT], Awaitable[ResultDataT]], /
+    ) -> Callable[[ResultDataT], Awaitable[ResultDataT]]: ...
 
     def result_validator(
-        self, func: _result.ResultValidatorFunc[AgentDeps, ResultData], /
-    ) -> _result.ResultValidatorFunc[AgentDeps, ResultData]:
+        self, func: _result.ResultValidatorFunc[AgentDepsT, ResultDataT], /
+    ) -> _result.ResultValidatorFunc[AgentDepsT, ResultDataT]:
         """Decorator to register a result validator function.
 
         Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
@@ -762,11 +775,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         #> success (no tool calls)
         ```
         """
-        self._result_validators.append(_result.ResultValidator[AgentDeps, Any](func))
+        self._result_validators.append(_result.ResultValidator[AgentDepsT, Any](func))
         return func
 
     @overload
-    def tool(self, func: ToolFuncContext[AgentDeps, ToolParams], /) -> ToolFuncContext[AgentDeps, ToolParams]: ...
+    def tool(self, func: ToolFuncContext[AgentDepsT, ToolParams], /) -> ToolFuncContext[AgentDepsT, ToolParams]: ...
 
     @overload
     def tool(
@@ -774,18 +787,18 @@ class Agent(Generic[AgentDeps, ResultData]):
         /,
         *,
         retries: int | None = None,
-        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
-    ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
+    ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ...
 
     def tool(
         self,
-        func: ToolFuncContext[AgentDeps, ToolParams] | None = None,
+        func: ToolFuncContext[AgentDepsT, ToolParams] | None = None,
         /,
         *,
         retries: int | None = None,
-        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
     ) -> Any:
@@ -832,8 +845,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         if func is None:
 
             def tool_decorator(
-                func_: ToolFuncContext[AgentDeps, ToolParams],
-            ) -> ToolFuncContext[AgentDeps, ToolParams]:
+                func_: ToolFuncContext[AgentDepsT, ToolParams],
+            ) -> ToolFuncContext[AgentDepsT, ToolParams]:
                 # noinspection PyTypeChecker
                 self._register_function(func_, True, retries, prepare, docstring_format, require_parameter_descriptions)
                 return func_
@@ -853,7 +866,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         /,
         *,
         retries: int | None = None,
-        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
@@ -864,7 +877,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         /,
         *,
         retries: int | None = None,
-        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
     ) -> Any:
@@ -924,16 +937,16 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     def _register_function(
         self,
-        func: ToolFuncEither[AgentDeps, ToolParams],
+        func: ToolFuncEither[AgentDepsT, ToolParams],
         takes_ctx: bool,
         retries: int | None,
-        prepare: ToolPrepareFunc[AgentDeps] | None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None,
         docstring_format: DocstringFormat,
         require_parameter_descriptions: bool,
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
-        tool = Tool(
+        tool = Tool[AgentDepsT](
             func,
             takes_ctx=takes_ctx,
             max_retries=retries_,
@@ -943,7 +956,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         )
         self._register_tool(tool)
 
-    def _register_tool(self, tool: Tool[AgentDeps]) -> None:
+    def _register_tool(self, tool: Tool[AgentDepsT]) -> None:
         """Private utility to register a tool instance."""
         if tool.max_retries is None:
             # noinspection PyTypeChecker
@@ -986,12 +999,12 @@ class Agent(Generic[AgentDeps, ResultData]):
         return model_
 
     async def _prepare_model(
-        self, run_context: RunContext[AgentDeps], result_schema: _result.ResultSchema[RunResultData] | None
+        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
     ) -> models.AgentModel:
         """Build tools and create an agent model."""
         function_tools: list[ToolDefinition] = []
 
-        async def add_tool(tool: Tool[AgentDeps]) -> None:
+        async def add_tool(tool: Tool[AgentDepsT]) -> None:
             ctx = run_context.replace_with(retry=tool.current_retry, tool_name=tool.name)
             if tool_def := await tool.prepare_tool_def(ctx):
                 function_tools.append(tool_def)
@@ -1005,7 +1018,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         )
 
     async def _reevaluate_dynamic_prompts(
-        self, messages: list[_messages.ModelMessage], run_context: RunContext[AgentDeps]
+        self, messages: list[_messages.ModelMessage], run_context: RunContext[AgentDepsT]
     ) -> None:
         """Reevaluate any `SystemPromptPart` with dynamic_ref in the provided messages by running the associated runner function."""
         # Only proceed if there's at least one dynamic runner.
@@ -1022,8 +1035,8 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 )
 
     def _prepare_result_schema(
-        self, result_type: type[RunResultData] | None
-    ) -> _result.ResultSchema[RunResultData] | None:
+        self, result_type: type[RunResultDataT] | None
+    ) -> _result.ResultSchema[RunResultDataT] | None:
         if result_type is not None:
             if self._result_validators:
                 raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
@@ -1034,10 +1047,13 @@ class Agent(Generic[AgentDeps, ResultData]):
             return self._result_schema  # pyright: ignore[reportReturnType]
 
     async def _prepare_messages(
-        self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[AgentDeps]
+        self,
+        user_prompt: str,
+        message_history: list[_messages.ModelMessage] | None,
+        run_context: RunContext[AgentDepsT],
     ) -> list[_messages.ModelMessage]:
         try:
-            ctx_messages = _messages_ctx_var.get()
+            ctx_messages = get_captured_run_messages()
         except LookupError:
             messages: list[_messages.ModelMessage] = []
         else:
@@ -1063,9 +1079,9 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _handle_model_response(
         self,
         model_response: _messages.ModelResponse,
-        run_context: RunContext[AgentDeps],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        run_context: RunContext[AgentDepsT],
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Process a non-streamed response from the model.
 
         Returns:
@@ -1094,11 +1110,11 @@ class Agent(Generic[AgentDeps, ResultData]):
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
     async def _handle_text_response(
-        self, text: str, run_context: RunContext[AgentDeps], result_schema: _result.ResultSchema[RunResultData] | None
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a plain text response from the model for non-streaming responses."""
         if self._allow_text_result(result_schema):
-            result_data_input = cast(RunResultData, text)
+            result_data_input = cast(RunResultDataT, text)
             try:
                 result_data = await self._validate_result(result_data_input, run_context, None)
             except _result.ToolRetryError as e:
@@ -1116,14 +1132,14 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _handle_structured_response(
         self,
         tool_calls: list[_messages.ToolCallPart],
-        run_context: RunContext[AgentDeps],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        run_context: RunContext[AgentDepsT],
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
 
         # first look for the result tool call
-        final_result: _MarkFinalResult[RunResultData] | None = None
+        final_result: _MarkFinalResult[RunResultDataT] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
         if result_schema is not None:
@@ -1133,7 +1149,6 @@ class Agent(Generic[AgentDeps, ResultData]):
                     result_data = result_tool.validate(call)
                     result_data = await self._validate_result(result_data, run_context, call)
                 except _result.ToolRetryError as e:
-                    self._incr_result_retry(run_context)
                     parts.append(e.tool_retry)
                 else:
                     final_result = _MarkFinalResult(result_data, call.tool_name)
@@ -1143,14 +1158,17 @@ class Agent(Generic[AgentDeps, ResultData]):
             tool_calls, final_result and final_result.tool_name, run_context, result_schema
         )
 
+        if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
+            self._incr_result_retry(run_context)
+
         return final_result, parts
 
     async def _process_function_tools(
         self,
         tool_calls: list[_messages.ToolCallPart],
         result_tool_name: str | None,
-        run_context: RunContext[AgentDeps],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        run_context: RunContext[AgentDepsT],
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> list[_messages.ModelRequestPart]:
         """Process function (non-result) tool calls in parallel.
 
@@ -1196,7 +1214,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                         )
                     )
             else:
-                parts.append(self._unknown_tool(call.tool_name, run_context, result_schema))
+                parts.append(self._unknown_tool(call.tool_name, result_schema))
 
         # Run all tool tasks in parallel
         if tasks:
@@ -1208,8 +1226,8 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _handle_streamed_response(
         self,
         streamed_response: models.StreamedResponse,
-        run_context: RunContext[AgentDeps],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        run_context: RunContext[AgentDepsT],
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
         """Process a streamed response from the model.
 
@@ -1243,7 +1261,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 if tool := self._function_tools.get(p.tool_name):
                     tasks.append(asyncio.create_task(tool.run(p, run_context), name=p.tool_name))
                 else:
-                    parts.append(self._unknown_tool(p.tool_name, run_context, result_schema))
+                    parts.append(self._unknown_tool(p.tool_name, result_schema))
 
         if received_text and not tasks and not parts:
             # Can only get here if self._allow_text_result returns `False` for the provided result_schema
@@ -1256,30 +1274,34 @@ class Agent(Generic[AgentDeps, ResultData]):
         with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
             task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
             parts.extend(task_results)
+
+        if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
+            self._incr_result_retry(run_context)
+
         return model_response, parts
 
     async def _validate_result(
         self,
-        result_data: RunResultData,
-        run_context: RunContext[AgentDeps],
+        result_data: RunResultDataT,
+        run_context: RunContext[AgentDepsT],
         tool_call: _messages.ToolCallPart | None,
-    ) -> RunResultData:
+    ) -> RunResultDataT:
         if self._result_validators:
-            agent_result_data = cast(ResultData, result_data)
+            agent_result_data = cast(ResultDataT, result_data)
             for validator in self._result_validators:
                 agent_result_data = await validator.validate(agent_result_data, tool_call, run_context)
-            return cast(RunResultData, agent_result_data)
+            return cast(RunResultDataT, agent_result_data)
         else:
             return result_data
 
-    def _incr_result_retry(self, run_context: RunContext[AgentDeps]) -> None:
+    def _incr_result_retry(self, run_context: RunContext[AgentDepsT]) -> None:
         run_context.retry += 1
         if run_context.retry > self._max_result_retries:
             raise exceptions.UnexpectedModelBehavior(
                 f'Exceeded maximum retries ({self._max_result_retries}) for result validation'
             )
 
-    async def _sys_parts(self, run_context: RunContext[AgentDeps]) -> list[_messages.ModelRequestPart]:
+    async def _sys_parts(self, run_context: RunContext[AgentDepsT]) -> list[_messages.ModelRequestPart]:
         """Build the initial messages for the conversation."""
         messages: list[_messages.ModelRequestPart] = [_messages.SystemPromptPart(p) for p in self._system_prompts]
         for sys_prompt_runner in self._system_prompt_functions:
@@ -1293,10 +1315,8 @@ class Agent(Generic[AgentDeps, ResultData]):
     def _unknown_tool(
         self,
         tool_name: str,
-        run_context: RunContext[AgentDeps],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _messages.RetryPromptPart:
-        self._incr_result_retry(run_context)
         names = list(self._function_tools.keys())
         if result_schema:
             names.extend(result_schema.tool_names())
@@ -1306,7 +1326,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             msg = 'No tools available.'
         return _messages.RetryPromptPart(content=f'Unknown tool name: {tool_name!r}. {msg}')
 
-    def _get_deps(self, deps: AgentDeps) -> AgentDeps:
+    def _get_deps(self: Agent[T, Any], deps: T) -> T:
         """Get deps for a run.
 
         If we've overridden deps via `_override_deps`, use that, otherwise use the deps passed to the call.
@@ -1338,7 +1358,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                             return
 
     @staticmethod
-    def _allow_text_result(result_schema: _result.ResultSchema[RunResultData] | None) -> bool:
+    def _allow_text_result(result_schema: _result.ResultSchema[RunResultDataT] | None) -> bool:
         return result_schema is None or result_schema.allow_text_result
 
     @property
@@ -1393,16 +1413,20 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
             _messages_ctx_var.reset(token)
 
 
+def get_captured_run_messages() -> _RunMessages:
+    return _messages_ctx_var.get()
+
+
 @dataclasses.dataclass
-class _MarkFinalResult(Generic[ResultData]):
+class _MarkFinalResult(Generic[ResultDataT]):
     """Marker class to indicate that the result is the final result.
 
-    This allows us to use `isinstance`, which wouldn't be possible if we were returning `ResultData` directly.
+    This allows us to use `isinstance`, which wouldn't be possible if we were returning `ResultDataT` directly.
 
     It also avoids problems in the case where the result type is itself `None`, but is set.
     """
 
-    data: ResultData
+    data: ResultDataT
     """The final result data."""
     tool_name: str | None
     """Name of the final result tool, None if the result is a string."""