pydantic-ai-slim 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

pydantic_ai/agent.py

@@ -22,10 +22,11 @@ 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,
     ToolDefinition,
@@ -50,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.
@@ -63,11 +66,11 @@ RunResultData = TypeVar('RunResultData')
 
 @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]`.
 
@@ -103,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',
     ):
@@ -199,12 +202,12 @@ 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(
@@ -214,7 +217,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[RunResultData],
         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,
@@ -227,7 +230,7 @@ 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,
@@ -242,9 +245,10 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         agent = Agent('openai:gpt-4o')
 
-        result_sync = agent.run_sync('What is the capital of Italy?')
-        print(result_sync.data)
-        #> Rome
+        async def main():
+            result = await agent.run('What is the capital of France?')
+            print(result.data)
+            #> Paris
         ```
 
         Args:
@@ -336,12 +340,12 @@ 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(
@@ -351,7 +355,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[RunResultData] | 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,
@@ -365,7 +369,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[RunResultData] | 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,
@@ -382,10 +386,9 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         agent = Agent('openai:gpt-4o')
 
-        async def main():
-            result = await agent.run('What is the capital of France?')
-            print(result.data)
-            #> Paris
+        result_sync = agent.run_sync('What is the capital of Italy?')
+        print(result_sync.data)
+        #> Rome
         ```
 
         Args:
@@ -427,12 +430,12 @@ 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(
@@ -442,12 +445,12 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[RunResultData],
         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, RunResultData]]: ...
 
     @asynccontextmanager
     async def run_stream(
@@ -457,12 +460,12 @@ class Agent(Generic[AgentDeps, ResultData]):
         result_type: type[RunResultData] | 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:
@@ -535,7 +538,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                         model_req_span.__exit__(None, None, None)
 
                         with _logfire.span('handle model response') as handle_span:
-                            maybe_final_result = await self._handle_streamed_model_response(
+                            maybe_final_result = await self._handle_streamed_response(
                                 model_response, run_context, result_schema
                             )
 
@@ -559,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, RunResultData]], self._result_validators
+                                )
+
                                 yield result.StreamedRunResult(
                                     messages,
                                     new_message_index,
@@ -570,7 +582,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                                     result_stream,
                                     result_schema,
                                     run_context,
-                                    self._result_validators,
+                                    result_validators,
                                     result_tool_name,
                                     on_complete,
                                 )
@@ -596,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.
@@ -632,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]: ...
@@ -649,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.
 
@@ -695,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
@@ -711,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.
@@ -761,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(
@@ -773,16 +787,20 @@ class Agent(Generic[AgentDeps, ResultData]):
         /,
         *,
         retries: int | None = None,
-        prepare: ToolPrepareFunc[AgentDeps] | None = None,
-    ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+    ) -> 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:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -820,20 +838,23 @@ class Agent(Generic[AgentDeps, ResultData]):
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
         """
         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)
+                self._register_function(func_, True, retries, prepare, docstring_format, require_parameter_descriptions)
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_function(func, True, retries, prepare)
+            self._register_function(func, True, retries, prepare, docstring_format, require_parameter_descriptions)
             return func
 
     @overload
@@ -845,7 +866,9 @@ 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]]: ...
 
     def tool_plain(
@@ -854,7 +877,9 @@ 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:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -892,32 +917,46 @@ class Agent(Generic[AgentDeps, ResultData]):
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
         """
         if func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, False, retries, prepare)
+                self._register_function(
+                    func_, False, retries, prepare, docstring_format, require_parameter_descriptions
+                )
                 return func_
 
             return tool_decorator
         else:
-            self._register_function(func, False, retries, prepare)
+            self._register_function(func, False, retries, prepare, docstring_format, require_parameter_descriptions)
             return func
 
     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(func, takes_ctx=takes_ctx, max_retries=retries_, prepare=prepare)
+        tool = Tool[AgentDepsT](
+            func,
+            takes_ctx=takes_ctx,
+            max_retries=retries_,
+            prepare=prepare,
+            docstring_format=docstring_format,
+            require_parameter_descriptions=require_parameter_descriptions,
+        )
         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
@@ -960,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[RunResultData] | 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)
@@ -979,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.
@@ -1008,7 +1047,10 @@ 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()
@@ -1037,7 +1079,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _handle_model_response(
         self,
         model_response: _messages.ModelResponse,
-        run_context: RunContext[AgentDeps],
+        run_context: RunContext[AgentDepsT],
         result_schema: _result.ResultSchema[RunResultData] | None,
     ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Process a non-streamed response from the model.
@@ -1068,7 +1110,7 @@ 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
+        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultData] | None
     ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Handle a plain text response from the model for non-streaming responses."""
         if self._allow_text_result(result_schema):
@@ -1090,7 +1132,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     async def _handle_structured_response(
         self,
         tool_calls: list[_messages.ToolCallPart],
-        run_context: RunContext[AgentDeps],
+        run_context: RunContext[AgentDepsT],
         result_schema: _result.ResultSchema[RunResultData] | None,
     ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
@@ -1100,14 +1142,13 @@ class Agent(Generic[AgentDeps, ResultData]):
         final_result: _MarkFinalResult[RunResultData] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
-        if result_schema := result_schema:
+        if result_schema is not None:
             if match := result_schema.find_tool(tool_calls):
                 call, result_tool = match
                 try:
                     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)
@@ -1117,13 +1158,16 @@ 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],
+        run_context: RunContext[AgentDepsT],
         result_schema: _result.ResultSchema[RunResultData] | None,
     ) -> list[_messages.ModelRequestPart]:
         """Process function (non-result) tool calls in parallel.
@@ -1170,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:
@@ -1179,99 +1223,85 @@ class Agent(Generic[AgentDeps, ResultData]):
                 parts.extend(task_results)
         return parts
 
-    async def _handle_streamed_model_response(
+    async def _handle_streamed_response(
         self,
-        model_response: models.EitherStreamedResponse,
-        run_context: RunContext[AgentDeps],
+        streamed_response: models.StreamedResponse,
+        run_context: RunContext[AgentDepsT],
         result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> (
-        _MarkFinalResult[models.EitherStreamedResponse]
-        | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]
-    ):
+    ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
         """Process a streamed response from the model.
 
         Returns:
             Either a final result or a tuple of the model response and the tool responses for the next request.
             If a final result is returned, the conversation should end.
         """
-        if isinstance(model_response, models.StreamTextResponse):
-            # plain string response
-            if self._allow_text_result(result_schema):
-                return _MarkFinalResult(model_response, None)
-            else:
-                self._incr_result_retry(run_context)
-                response = _messages.RetryPromptPart(
-                    content='Plain text responses are not permitted, please call one of the functions instead.',
-                )
-                # stream the response, so usage is correct
-                async for _ in model_response:
-                    pass
-
-                text = ''.join(model_response.get(final=True))
-                return _messages.ModelResponse([_messages.TextPart(text)]), [response]
-        elif isinstance(model_response, models.StreamStructuredResponse):
-            if result_schema is not None:
-                # if there's a result schema, iterate over the stream until we find at least one tool
-                # NOTE: this means we ignore any other tools called here
-                structured_msg = model_response.get()
-                while not structured_msg.parts:
-                    try:
-                        await model_response.__anext__()
-                    except StopAsyncIteration:
-                        break
-                    structured_msg = model_response.get()
-
-                if match := result_schema.find_tool(structured_msg.parts):
-                    call, _ = match
-                    return _MarkFinalResult(model_response, call.tool_name)
-
-            # the model is calling a tool function, consume the response to get the next message
-            async for _ in model_response:
-                pass
-            model_response_msg = model_response.get()
-            if not model_response_msg.parts:
-                raise exceptions.UnexpectedModelBehavior('Received empty tool call message')
-
-            # we now run all tool functions in parallel
-            tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
-            parts: list[_messages.ModelRequestPart] = []
-            for item in model_response_msg.parts:
-                if isinstance(item, _messages.ToolCallPart):
-                    call = item
-                    if tool := self._function_tools.get(call.tool_name):
-                        tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name))
-                    else:
-                        parts.append(self._unknown_tool(call.tool_name, run_context, result_schema))
+        received_text = False
+
+        async for maybe_part_event in streamed_response:
+            if isinstance(maybe_part_event, _messages.PartStartEvent):
+                new_part = maybe_part_event.part
+                if isinstance(new_part, _messages.TextPart):
+                    received_text = True
+                    if self._allow_text_result(result_schema):
+                        return _MarkFinalResult(streamed_response, None)
+                elif isinstance(new_part, _messages.ToolCallPart):
+                    if result_schema is not None and (match := result_schema.find_tool([new_part])):
+                        call, _ = match
+                        return _MarkFinalResult(streamed_response, call.tool_name)
+                else:
+                    assert_never(new_part)
 
-            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)
-            return model_response_msg, parts
-        else:
-            assert_never(model_response)
+        tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
+        parts: list[_messages.ModelRequestPart] = []
+        model_response = streamed_response.get()
+        if not model_response.parts:
+            raise exceptions.UnexpectedModelBehavior('Received empty model response')
+        for p in model_response.parts:
+            if isinstance(p, _messages.ToolCallPart):
+                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, 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
+            self._incr_result_retry(run_context)
+            model_response = _messages.RetryPromptPart(
+                content='Plain text responses are not permitted, please call one of the functions instead.',
+            )
+            return streamed_response.get(), [model_response]
+
+        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],
+        run_context: RunContext[AgentDepsT],
         tool_call: _messages.ToolCallPart | None,
     ) -> RunResultData:
         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)
         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:
@@ -1285,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,
     ) -> _messages.RetryPromptPart:
-        self._incr_result_retry(run_context)
         names = list(self._function_tools.keys())
         if result_schema:
             names.extend(result_schema.tool_names())
@@ -1298,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.
@@ -1386,15 +1414,15 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
 
 
 @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."""