Let message history end on ModelResponse and execute pending tool calls (#2562)

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -23,7 +23,7 @@
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
 from .exceptions import ToolRetryError
 from .output import OutputDataT, OutputSpec
-from .settings import ModelSettings, merge_model_settings
+from .settings import ModelSettings
 from .tools import RunContext, ToolDefinition, ToolKind
 
 if TYPE_CHECKING:
@@ -158,28 +158,7 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
 
     async def run(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> ModelRequestNode[DepsT, NodeRunEndT]:
-        return ModelRequestNode[DepsT, NodeRunEndT](request=await self._get_first_message(ctx))
-
-    async def _get_first_message(
-        self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> _messages.ModelRequest:
-        run_context = build_run_context(ctx)
-        history, next_message = await self._prepare_messages(
-            self.user_prompt, ctx.state.message_history, ctx.deps.get_instructions, run_context
-        )
-        ctx.state.message_history = history
-        run_context.messages = history
-
-        return next_message
-
-    async def _prepare_messages(
-        self,
-        user_prompt: str | Sequence[_messages.UserContent] | None,
-        message_history: list[_messages.ModelMessage] | None,
-        get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]],
-        run_context: RunContext[DepsT],
-    ) -> tuple[list[_messages.ModelMessage], _messages.ModelRequest]:
+    ) -> Union[ModelRequestNode[DepsT, NodeRunEndT], CallToolsNode[DepsT, NodeRunEndT]]:  # noqa UP007
         try:
             ctx_messages = get_captured_run_messages()
         except LookupError:
@@ -191,29 +170,48 @@ async def _prepare_messages(
                 messages = ctx_messages.messages
                 ctx_messages.used = True
 
+        # Add message history to the `capture_run_messages` list, which will be empty at this point
+        messages.extend(ctx.state.message_history)
+        # Use the `capture_run_messages` list as the message history so that new messages are added to it
+        ctx.state.message_history = messages
+
+        run_context = build_run_context(ctx)
+
         parts: list[_messages.ModelRequestPart] = []
-        instructions = await get_instructions(run_context)
-        if message_history:
-            # Shallow copy messages
-            messages.extend(message_history)
+        if messages:
             # Reevaluate any dynamic system prompt parts
             await self._reevaluate_dynamic_prompts(messages, run_context)
         else:
             parts.extend(await self._sys_parts(run_context))
 
-        if user_prompt is not None:
-            parts.append(_messages.UserPromptPart(user_prompt))
-        elif (
-            len(parts) == 0
-            and message_history
-            and (last_message := message_history[-1])
-            and isinstance(last_message, _messages.ModelRequest)
-        ):
-            # Drop last message that came from history and reuse its parts
-            messages.pop()
-            parts.extend(last_message.parts)
+        if messages and (last_message := messages[-1]):
+            if isinstance(last_message, _messages.ModelRequest) and self.user_prompt is None:
+                # Drop last message from history and reuse its parts
+                messages.pop()
+                parts.extend(last_message.parts)
+            elif isinstance(last_message, _messages.ModelResponse):
+                if self.user_prompt is None:
+                    # `CallToolsNode` requires the tool manager to be prepared for the run step
+                    # This will raise errors for any tool name conflicts
+                    ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
+
+                    # Skip ModelRequestNode and go directly to CallToolsNode
+                    return CallToolsNode[DepsT, NodeRunEndT](model_response=last_message)
+                elif any(isinstance(part, _messages.ToolCallPart) for part in last_message.parts):
+                    raise exceptions.UserError(
+                        'Cannot provide a new user prompt when the message history ends with '
+                        'a model response containing unprocessed tool calls. Either process the '
+                        'tool calls first (by calling `iter` with `user_prompt=None`) or append a '
+                        '`ModelRequest` with `ToolResultPart`s.'
+                    )
+
+        if self.user_prompt is not None:
+            parts.append(_messages.UserPromptPart(self.user_prompt))
+
+        instructions = await ctx.deps.get_instructions(run_context)
+        next_message = _messages.ModelRequest(parts, instructions=instructions)
 
-        return messages, _messages.ModelRequest(parts, instructions=instructions)
+        return ModelRequestNode[DepsT, NodeRunEndT](request=next_message)
 
     async def _reevaluate_dynamic_prompts(
         self, messages: list[_messages.ModelMessage], run_context: RunContext[DepsT]
@@ -250,11 +248,6 @@ async def _prepare_request_parameters(
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
 ) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
-    run_context = build_run_context(ctx)
-
-    # This will raise errors for any tool name conflicts
-    ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
-
     output_schema = ctx.deps.output_schema
     output_object = None
     if isinstance(output_schema, _output.NativeOutputSchema):
@@ -357,21 +350,21 @@ async def _prepare_request(
 
         run_context = build_run_context(ctx)
 
-        model_settings = merge_model_settings(ctx.deps.model_settings, None)
+        # This will raise errors for any tool name conflicts
+        ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
+
+        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
 
-        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
-
+        model_settings = ctx.deps.model_settings
         usage = ctx.state.usage
         if ctx.deps.usage_limits.count_tokens_before_request:
             # Copy to avoid modifying the original usage object with the counted usage
             usage = dataclasses.replace(usage)
 
-            counted_usage = await ctx.deps.model.count_tokens(
-                message_history, ctx.deps.model_settings, model_request_parameters
-            )
+            counted_usage = await ctx.deps.model.count_tokens(message_history, model_settings, model_request_parameters)
             usage.incr(counted_usage)
 
         ctx.deps.usage_limits.check_before_request(usage)