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

pydantic_ai/_pydantic.py

@@ -17,10 +17,10 @@ from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
 
 from ._griffe import doc_descriptions
-from ._utils import ObjectJsonSchema, check_object_json_schema, is_model_like
+from ._utils import check_object_json_schema, is_model_like
 
 if TYPE_CHECKING:
-    pass
+    from .tools import ObjectJsonSchema
 
 
 __all__ = 'function_schema', 'LazyTypeAdapter'
@@ -168,11 +168,13 @@ def takes_ctx(function: Callable[..., Any]) -> bool:
     """
     sig = signature(function)
     try:
-        _, first_param = next(iter(sig.parameters.items()))
+        first_param_name = next(iter(sig.parameters.keys()))
     except StopIteration:
         return False
     else:
-        return first_param.annotation is not sig.empty and _is_call_ctx(first_param.annotation)
+        type_hints = _typing_extra.get_function_type_hints(function)
+        annotation = type_hints[first_param_name]
+        return annotation is not sig.empty and _is_call_ctx(annotation)
 
 
 def _build_schema(

pydantic_ai/_result.py

@@ -14,7 +14,7 @@ from . import _utils, messages
 from .exceptions import ModelRetry
 from .messages import ModelStructuredResponse, ToolCall
 from .result import ResultData
-from .tools import AgentDeps, ResultValidatorFunc, RunContext
+from .tools import AgentDeps, ResultValidatorFunc, RunContext, ToolDefinition
 
 
 @dataclass
@@ -94,10 +94,7 @@ class ResultSchema(Generic[ResultData]):
             allow_text_result = False
 
         def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[ResultData]:
-            return cast(
-                ResultTool[ResultData],
-                ResultTool.build(a, tool_name_, description, multiple),  # pyright: ignore[reportUnknownMemberType]
-            )
+            return cast(ResultTool[ResultData], ResultTool(a, tool_name_, description, multiple))
 
         tools: dict[str, ResultTool[ResultData]] = {}
         if args := get_union_args(response_type):
@@ -121,38 +118,38 @@ class ResultSchema(Generic[ResultData]):
         """Return the names of the tools."""
         return list(self.tools.keys())
 
+    def tool_defs(self) -> list[ToolDefinition]:
+        """Get tool definitions to register with the model."""
+        return [t.tool_def for t in self.tools.values()]
+
 
 DEFAULT_DESCRIPTION = 'The final response which ends this conversation'
 
 
-@dataclass
+@dataclass(init=False)
 class ResultTool(Generic[ResultData]):
-    name: str
-    description: str
+    tool_def: ToolDefinition
     type_adapter: TypeAdapter[Any]
-    json_schema: _utils.ObjectJsonSchema
-    outer_typed_dict_key: str | None
 
-    @classmethod
-    def build(cls, response_type: type[ResultData], name: str, description: str | None, multiple: bool) -> Self | None:
+    def __init__(self, response_type: type[ResultData], name: str, description: str | None, multiple: bool):
         """Build a ResultTool dataclass from a response type."""
         assert response_type is not str, 'ResultTool does not support str as a response type'
 
         if _utils.is_model_like(response_type):
-            type_adapter = TypeAdapter(response_type)
+            self.type_adapter = TypeAdapter(response_type)
             outer_typed_dict_key: str | None = None
             # noinspection PyArgumentList
-            json_schema = _utils.check_object_json_schema(type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
         else:
             response_data_typed_dict = TypedDict('response_data_typed_dict', {'response': response_type})  # noqa
-            type_adapter = TypeAdapter(response_data_typed_dict)
+            self.type_adapter = TypeAdapter(response_data_typed_dict)
             outer_typed_dict_key = 'response'
             # noinspection PyArgumentList
-            json_schema = _utils.check_object_json_schema(type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
             # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
-            json_schema.pop('title')
+            parameters_json_schema.pop('title')
 
-        if json_schema_description := json_schema.pop('description', None):
+        if json_schema_description := parameters_json_schema.pop('description', None):
             if description is None:
                 tool_description = json_schema_description
             else:
@@ -162,11 +159,10 @@ class ResultTool(Generic[ResultData]):
             if multiple:
                 tool_description = f'{union_arg_name(response_type)}: {tool_description}'
 
-        return cls(
+        self.tool_def = ToolDefinition(
             name=name,
             description=tool_description,
-            type_adapter=type_adapter,
-            json_schema=json_schema,
+            parameters_json_schema=parameters_json_schema,
             outer_typed_dict_key=outer_typed_dict_key,
         )
 
@@ -204,7 +200,7 @@ class ResultTool(Generic[ResultData]):
             else:
                 raise
         else:
-            if k := self.outer_typed_dict_key:
+            if k := self.tool_def.outer_typed_dict_key:
                 result = result[k]
             return result
 

pydantic_ai/_system_prompt.py

@@ -21,7 +21,7 @@ class SystemPromptRunner(Generic[AgentDeps]):
 
     async def run(self, deps: AgentDeps) -> str:
         if self._takes_ctx:
-            args = (RunContext(deps, 0, None),)
+            args = (RunContext(deps, 0),)
         else:
             args = ()
 

pydantic_ai/_utils.py

@@ -8,12 +8,15 @@ from dataclasses import dataclass, is_dataclass
 from datetime import datetime, timezone
 from functools import partial
 from types import GenericAlias
-from typing import Any, Callable, Generic, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast, overload
 
 from pydantic import BaseModel
 from pydantic.json_schema import JsonSchemaValue
 from typing_extensions import ParamSpec, TypeAlias, TypeGuard, is_typeddict
 
+if TYPE_CHECKING:
+    from .tools import ObjectJsonSchema
+
 _P = ParamSpec('_P')
 _R = TypeVar('_R')
 
@@ -39,10 +42,6 @@ def is_model_like(type_: Any) -> bool:
     )
 
 
-# With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_items=Any`
-ObjectJsonSchema: TypeAlias = dict[str, Any]
-
-
 def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
     from .exceptions import UserError
 
@@ -127,6 +126,12 @@ class Either(Generic[Left, Right]):
     def whichever(self) -> Left | Right:
         return self._left.value if self._left is not None else self.right
 
+    def __repr__(self):
+        if left := self._left:
+            return f'Either(left={left.value!r})'
+        else:
+            return f'Either(right={self.right!r})'
+
 
 @asynccontextmanager
 async def group_by_temporal(
@@ -218,7 +223,7 @@ async def group_by_temporal(
 
     try:
         yield async_iter_groups()
-    finally:
+    finally:  # pragma: no cover
         # after iteration if a tasks still exists, cancel it, this will only happen if an error occurred
         if task:
             task.cancel('Cancelling due to error in iterator')

pydantic_ai/agent.py

@@ -22,7 +22,17 @@ from . import (
     result,
 )
 from .result import ResultData
-from .tools import AgentDeps, RunContext, Tool, ToolFuncContext, ToolFuncEither, ToolFuncPlain, ToolParams
+from .tools import (
+    AgentDeps,
+    RunContext,
+    Tool,
+    ToolDefinition,
+    ToolFuncContext,
+    ToolFuncEither,
+    ToolFuncPlain,
+    ToolParams,
+    ToolPrepareFunc,
+)
 
 __all__ = ('Agent',)
 
@@ -136,7 +146,10 @@ class Agent(Generic[AgentDeps, ResultData]):
         self._function_tools = {}
         self._default_retries = retries
         for tool in tools:
-            self._register_tool(Tool.infer(tool))
+            if isinstance(tool, Tool):
+                self._register_tool(tool)
+            else:
+                self._register_tool(Tool(tool))
         self._deps_type = deps_type
         self._system_prompt_functions = []
         self._max_result_retries = result_retries if result_retries is not None else retries
@@ -166,7 +179,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
-        model_used, custom_model, agent_model = await self._get_agent_model(model)
+        model_used, mode_selection = await self._get_model(model)
 
         deps = self._get_deps(deps)
 
@@ -174,7 +187,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             '{agent_name} run {prompt=}',
             prompt=user_prompt,
             agent=self,
-            custom_model=custom_model,
+            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
@@ -182,14 +195,17 @@ class Agent(Generic[AgentDeps, ResultData]):
             self.last_run_messages = messages
 
             for tool in self._function_tools.values():
-                tool.reset()
+                tool.current_retry = 0
 
             cost = result.Cost()
 
             run_step = 0
             while True:
                 run_step += 1
-                with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                    agent_model = await self._prepare_model(model_used, deps)
+
+                with _logfire.span('model request', run_step=run_step) as model_req_span:
                     model_response, request_cost = await agent_model.request(messages)
                     model_req_span.set_attribute('response', model_response)
                     model_req_span.set_attribute('cost', request_cost)
@@ -198,12 +214,15 @@ class Agent(Generic[AgentDeps, ResultData]):
                 messages.append(model_response)
                 cost += request_cost
 
-                with _logfire.span('handle model response') as handle_span:
-                    either = await self._handle_model_response(model_response, deps)
+                with _logfire.span('handle model response', run_step=run_step) as handle_span:
+                    final_result, response_messages = await self._handle_model_response(model_response, deps)
 
-                    if isinstance(either, _MarkFinalResult):
-                        # we have a final result, end the conversation
-                        result_data = either.data
+                    # Add all messages to the conversation
+                    messages.extend(response_messages)
+
+                    # Check if we got a final result
+                    if final_result is not None:
+                        result_data = final_result.data
                         run_span.set_attribute('all_messages', messages)
                         run_span.set_attribute('cost', cost)
                         handle_span.set_attribute('result', result_data)
@@ -211,11 +230,9 @@ class Agent(Generic[AgentDeps, ResultData]):
                         return result.RunResult(messages, new_message_index, result_data, cost)
                     else:
                         # continue the conversation
-                        tool_responses = either
-                        handle_span.set_attribute('tool_responses', tool_responses)
-                        response_msgs = ' '.join(m.role for m in tool_responses)
+                        handle_span.set_attribute('tool_responses', response_messages)
+                        response_msgs = ' '.join(r.role for r in response_messages)
                         handle_span.message = f'handle model response -> {response_msgs}'
-                        messages.extend(tool_responses)
 
     def run_sync(
         self,
@@ -273,7 +290,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             # f_back because `asynccontextmanager` adds one frame
             if frame := inspect.currentframe():  # pragma: no branch
                 self._infer_name(frame.f_back)
-        model_used, custom_model, agent_model = await self._get_agent_model(model)
+        model_used, mode_selection = await self._get_model(model)
 
         deps = self._get_deps(deps)
 
@@ -281,7 +298,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             '{agent_name} run stream {prompt=}',
             prompt=user_prompt,
             agent=self,
-            custom_model=custom_model,
+            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
@@ -289,13 +306,17 @@ class Agent(Generic[AgentDeps, ResultData]):
             self.last_run_messages = messages
 
             for tool in self._function_tools.values():
-                tool.reset()
+                tool.current_retry = 0
 
             cost = result.Cost()
 
             run_step = 0
             while True:
                 run_step += 1
+
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                    agent_model = await self._prepare_model(model_used, deps)
+
                 with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
                     async with agent_model.request_stream(messages) as model_response:
                         model_req_span.set_attribute('response_type', model_response.__class__.__name__)
@@ -304,10 +325,16 @@ class Agent(Generic[AgentDeps, ResultData]):
                         model_req_span.__exit__(None, None, None)
 
                         with _logfire.span('handle model response') as handle_span:
-                            either = await self._handle_streamed_model_response(model_response, deps)
+                            final_result, response_messages = await self._handle_streamed_model_response(
+                                model_response, deps
+                            )
+
+                            # Add all messages to the conversation
+                            messages.extend(response_messages)
 
-                            if isinstance(either, _MarkFinalResult):
-                                result_stream = either.data
+                            # Check if we got a final result
+                            if final_result is not None:
+                                result_stream = final_result.data
                                 run_span.set_attribute('all_messages', messages)
                                 handle_span.set_attribute('result_type', result_stream.__class__.__name__)
                                 handle_span.message = 'handle model response -> final result'
@@ -323,11 +350,10 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 )
                                 return
                             else:
-                                tool_responses = either
-                                handle_span.set_attribute('tool_responses', tool_responses)
-                                response_msgs = ' '.join(m.role for m in tool_responses)
+                                # continue the conversation
+                                handle_span.set_attribute('tool_responses', response_messages)
+                                response_msgs = ' '.join(r.role for r in response_messages)
                                 handle_span.message = f'handle model response -> {response_msgs}'
-                                messages.extend(tool_responses)
                                 # the model_response should have been fully streamed by now, we can add it's cost
                                 cost += model_response.cost()
 
@@ -477,7 +503,11 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     @overload
     def tool(
-        self, /, *, retries: int | None = None
+        self,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
 
     def tool(
@@ -486,9 +516,9 @@ class Agent(Generic[AgentDeps, ResultData]):
         /,
         *,
         retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Any:
-        """Decorator to register a tool function which takes
-        [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
+        """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
         Can decorate a sync or async functions.
 
@@ -521,20 +551,23 @@ class Agent(Generic[AgentDeps, ResultData]):
             func: The tool function to register.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
-        """  # noqa: D205
+            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].
+        """
         if func is None:
 
             def tool_decorator(
                 func_: ToolFuncContext[AgentDeps, ToolParams],
             ) -> ToolFuncContext[AgentDeps, ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, True, retries)
+                self._register_function(func_, True, retries, prepare)
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_function(func, True, retries)
+            self._register_function(func, True, retries, prepare)
             return func
 
     @overload
@@ -542,10 +575,21 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     @overload
     def tool_plain(
-        self, /, *, retries: int | None = None
+        self,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
-    def tool_plain(self, func: ToolFuncPlain[ToolParams] | None = None, /, *, retries: int | None = None) -> Any:
+    def tool_plain(
+        self,
+        func: ToolFuncPlain[ToolParams] | None = None,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+    ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
         Can decorate a sync or async functions.
@@ -579,30 +623,38 @@ class Agent(Generic[AgentDeps, ResultData]):
             func: The tool function to register.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
+            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].
         """
         if func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, False, retries)
+                self._register_function(func_, False, retries, prepare)
                 return func_
 
             return tool_decorator
         else:
-            self._register_function(func, False, retries)
+            self._register_function(func, False, retries, prepare)
             return func
 
     def _register_function(
-        self, func: ToolFuncEither[AgentDeps, ToolParams], takes_ctx: bool, retries: int | None
+        self,
+        func: ToolFuncEither[AgentDeps, ToolParams],
+        takes_ctx: bool,
+        retries: int | None,
+        prepare: ToolPrepareFunc[AgentDeps] | None,
     ) -> 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, max_retries=retries_)
+        tool = Tool(func, takes_ctx=takes_ctx, max_retries=retries_, prepare=prepare)
         self._register_tool(tool)
 
     def _register_tool(self, tool: Tool[AgentDeps]) -> None:
         """Private utility to register a tool instance."""
         if tool.max_retries is None:
+            # noinspection PyTypeChecker
             tool = dataclasses.replace(tool, max_retries=self._default_retries)
 
         if tool.name in self._function_tools:
@@ -613,16 +665,14 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         self._function_tools[tool.name] = tool
 
-    async def _get_agent_model(
-        self, model: models.Model | models.KnownModelName | None
-    ) -> tuple[models.Model, models.Model | None, models.AgentModel]:
+    async def _get_model(self, model: models.Model | models.KnownModelName | None) -> tuple[models.Model, str]:
         """Create a model configured for this agent.
 
         Args:
             model: model to use for this run, required if `model` was not set when creating the agent.
 
         Returns:
-            a tuple of `(model used, custom_model if any, agent_model)`
+            a tuple of `(model used, how the model was selected)`
         """
         model_: models.Model
         if some_model := self._override_model:
@@ -633,19 +683,35 @@ class Agent(Generic[AgentDeps, ResultData]):
                     '(Even when `override(model=...)` is customizing the model that will actually be called)'
                 )
             model_ = some_model.value
-            custom_model = None
+            mode_selection = 'override-model'
         elif model is not None:
-            custom_model = model_ = models.infer_model(model)
+            model_ = models.infer_model(model)
+            mode_selection = 'custom'
         elif self.model is not None:
             # noinspection PyTypeChecker
             model_ = self.model = models.infer_model(self.model)
-            custom_model = None
+            mode_selection = 'from-agent'
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        result_tools = list(self._result_schema.tools.values()) if self._result_schema else None
-        agent_model = await model_.agent_model(self._function_tools, self._allow_text_result, result_tools)
-        return model_, custom_model, agent_model
+        return model_, mode_selection
+
+    async def _prepare_model(self, model: models.Model, deps: AgentDeps) -> models.AgentModel:
+        """Create building tools and create an agent model."""
+        function_tools: list[ToolDefinition] = []
+
+        async def add_tool(tool: Tool[AgentDeps]) -> None:
+            ctx = RunContext(deps, tool.current_retry, tool.name)
+            if tool_def := await tool.prepare_tool_def(ctx):
+                function_tools.append(tool_def)
+
+        await asyncio.gather(*map(add_tool, self._function_tools.values()))
+
+        return await model.agent_model(
+            function_tools=function_tools,
+            allow_text_result=self._allow_text_result,
+            result_tools=self._result_schema.tool_defs() if self._result_schema is not None else [],
+        )
 
     async def _prepare_messages(
         self, deps: AgentDeps, user_prompt: str, message_history: list[_messages.Message] | None
@@ -665,11 +731,11 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     async def _handle_model_response(
         self, model_response: _messages.ModelAnyResponse, deps: AgentDeps
-    ) -> _MarkFinalResult[ResultData] | list[_messages.Message]:
+    ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Process a non-streamed response from the model.
 
         Returns:
-            Return `Either` — left: final result data, right: list of messages to send back to the model.
+            A tuple of `(final_result, messages)`. If `final_result` is not `None`, the conversation should end.
         """
         if model_response.role == 'model-text-response':
             # plain string response
@@ -679,15 +745,15 @@ class Agent(Generic[AgentDeps, ResultData]):
                     result_data = await self._validate_result(result_data_input, deps, None)
                 except _result.ToolRetryError as e:
                     self._incr_result_retry()
-                    return [e.tool_retry]
+                    return None, [e.tool_retry]
                 else:
-                    return _MarkFinalResult(result_data)
+                    return _MarkFinalResult(result_data), []
             else:
                 self._incr_result_retry()
                 response = _messages.RetryPrompt(
                     content='Plain text responses are not permitted, please call one of the functions instead.',
                 )
-                return [response]
+                return None, [response]
         elif model_response.role == 'model-structured-response':
             if self._result_schema is not None:
                 # if there's a result schema, and any of the calls match one of its tools, return the result
@@ -699,9 +765,15 @@ class Agent(Generic[AgentDeps, ResultData]):
                         result_data = await self._validate_result(result_data, deps, call)
                     except _result.ToolRetryError as e:
                         self._incr_result_retry()
-                        return [e.tool_retry]
+                        return None, [e.tool_retry]
                     else:
-                        return _MarkFinalResult(result_data)
+                        # Add a ToolReturn message for the schema tool call
+                        tool_return = _messages.ToolReturn(
+                            tool_name=call.tool_name,
+                            content='Final result processed.',
+                            tool_id=call.tool_id,
+                        )
+                        return _MarkFinalResult(result_data), [tool_return]
 
             if not model_response.calls:
                 raise exceptions.UnexpectedModelBehavior('Received empty tool call message')
@@ -716,26 +788,24 @@ class Agent(Generic[AgentDeps, ResultData]):
                     messages.append(self._unknown_tool(call.tool_name))
 
             with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                messages += await asyncio.gather(*tasks)
-            return messages
+                task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
+                messages.extend(task_results)
+            return None, messages
         else:
             assert_never(model_response)
 
     async def _handle_streamed_model_response(
         self, model_response: models.EitherStreamedResponse, deps: AgentDeps
-    ) -> _MarkFinalResult[models.EitherStreamedResponse] | list[_messages.Message]:
+    ) -> tuple[_MarkFinalResult[models.EitherStreamedResponse] | None, list[_messages.Message]]:
         """Process a streamed response from the model.
 
-        TODO: change the response type to `models.EitherStreamedResponse | list[_messages.Message]` once we drop 3.9
-        (with 3.9 we get `TypeError: Subscripted generics cannot be used with class and instance checks`)
-
         Returns:
-            Return `Either` — left: final result data, right: list of messages to send back to the model.
+            A tuple of (final_result, messages). If final_result is not None, the conversation should end.
         """
         if isinstance(model_response, models.StreamTextResponse):
             # plain string response
             if self._allow_text_result:
-                return _MarkFinalResult(model_response)
+                return _MarkFinalResult(model_response), []
             else:
                 self._incr_result_retry()
                 response = _messages.RetryPrompt(
@@ -745,7 +815,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 async for _ in model_response:
                     pass
 
-                return [response]
+                return None, [response]
         else:
             assert isinstance(model_response, models.StreamStructuredResponse), f'Unexpected response: {model_response}'
             if self._result_schema is not None:
@@ -759,8 +829,14 @@ class Agent(Generic[AgentDeps, ResultData]):
                         break
                     structured_msg = model_response.get()
 
-                if self._result_schema.find_tool(structured_msg):
-                    return _MarkFinalResult(model_response)
+                if match := self._result_schema.find_tool(structured_msg):
+                    call, _ = match
+                    tool_return = _messages.ToolReturn(
+                        tool_name=call.tool_name,
+                        content='Final result processed.',
+                        tool_id=call.tool_id,
+                    )
+                    return _MarkFinalResult(model_response), [tool_return]
 
             # the model is calling a tool function, consume the response to get the next message
             async for _ in model_response:
@@ -779,8 +855,9 @@ class Agent(Generic[AgentDeps, ResultData]):
                     messages.append(self._unknown_tool(call.tool_name))
 
             with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                messages += await asyncio.gather(*tasks)
-            return messages
+                task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
+                messages.extend(task_results)
+            return None, messages
 
     async def _validate_result(
         self, result_data: ResultData, deps: AgentDeps, tool_call: _messages.ToolCall | None
@@ -852,6 +929,8 @@ class _MarkFinalResult(Generic[ResultData]):
     """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.
+
+    It also avoids problems in the case where the result type is itself `None`, but is set.
     """
 
     data: ResultData