Extracting `new_messages()` in the middle of an `Agent.iter` run to support user-in-the-loop approval of tool calls

[tibbe]

Question

I'm trying to implement an agent which streams outputs to the user and uses tools. Some tools require user approval, which means that we need to

• suspend the agent pending user approval and thus
• serialize all the messages so far so we can wake up the agent once we have the user's approval.

I'm trying to use Agent.iter to achieve this, with a function along the lines of:

# Allows sending the user messages over e.g. a WebSocket:
class OutputStreamBase(ABC):
    def on_tool_call_start(self, event: FunctionToolCallEvent) -> None:
        """Handle the start of a tool call event."""
        ...

    @abstractmethod
    def on_tool_call_end(self, event: FunctionToolResultEvent) -> None:
        """Handle the end of a tool call event."""
        ...

    @abstractmethod
    def on_part_start(self, index: int, part: TextPartDelta) -> None:
        """Handle the start of a part event."""
        ...

    @abstractmethod
    def on_part_delta(self, delta: TextPartDelta) -> None:
        """Handle the delta of a part event."""
        ...

    @abstractmethod
    def on_final_result(self, event: FinalResultEvent) -> None:
        """Handle the final result event."""
        ...

# Allows storing the current state to e.g. a database:
class BaseStore(ABC):
    @abstractmethod
    async def add_messages(self, messages: Iterable[ModelMessage]) -> None:
        ...

    @abstractmethod
    async def get_messages(self) -> list[ModelMessage]:
        ...

# Track the status of the agent (run):
RunStatus = Literal['COMPLETED', 'AWAITING_APPROVAL', 'RUNNING']

class RunResult(BaseModel):
    status: RunStatus
    output: str | None = None

# This is the interesting part:
async def run_agent_graph_step(
    agent: Agent,
    store: BaseStore,
    sink: OutputStreamBase,
    user_prompt: Optional[str] = None,
) -> RunResult:
    messages = await store.get_messages()

    async with agent.iter(user_prompt, message_history=messages) as run:
        async for node in run:
            if agent.is_user_prompt_node(node):
                pass
            elif Agent.is_model_request_node(node):
                async with node.stream(run.ctx) as request_stream:
                    async for event in request_stream:
                        if isinstance(event, PartStartEvent):
                            if isinstance(event.part, TextPartDelta):
                                sink.on_part_start(event.index, event.part)
                        elif isinstance(event, PartDeltaEvent):
                            if isinstance(event.delta, TextPartDelta):
                                sink.on_part_delta(event.delta)
                            elif isinstance(event.delta, ToolCallPartDelta):
                                pass
                        elif isinstance(event, FinalResultEvent):
                            pass
            elif agent.is_call_tools_node(node):
                async with node.stream(run.ctx) as handle_stream:
                    async for tool_event in handle_stream:
                        if isinstance(tool_event, FunctionToolCallEvent):
                            if requires_approval(tool_event):
                              # ----> HERE: we need to get new_messages() and serialize them <----
                              return RunResult(status='AWAITING_APPROVAL')
                            sink.on_tool_call_start(tool_event)
                        elif isinstance(tool_event, FunctionToolResultEvent):
                            sink.on_tool_call_end(tool_event)
            elif Agent.is_end_node(node):
                assert (
                    run.result is not None
                    and run.result.output == node.data.output
                )
                return RunResult(
                    status='COMPLETED',
                    output=node.data.output,
                )

    await store.add_messages(run.result.new_messages())
    return RunResult(status='RUNNING')

Is there a way to get new_messages() (so far) in the middle of an iter run? Is there a better way to model user-in-the-loop pausing/resuming of an agent that need tool call approvals?

Additional Context

No response

[tibbe]

Maybe a different way to phrase this question is: how can I "single step" the LLM, using the following logic:

def handle_request(prompt: str, approved_tool_call_id: str | None = None) -> State:
    messages = load_from_db()
    while True:
        result = agent.step(user_prompt=prompt, messages=messages)
        if instance(result, ToolCall):
            if requires_approval(result) and not approved_tool_call_id == result.tool_call_id:
                save_to_db(result.new_messages())
                return NeedsApproval(tool_call_id=result.tool_call_id)
            else:
                perform_tool_call()
        elif isinstance(result, End):
            save_to_db(result.new_messages())
            return Done()
        elif isinstance(result, Output):
            send_output_to_user(result)
    save_to_db(result.new_messages())
    return Running()

[tibbe]

This seems related to #642.

[DouweM]

@tibbe Just after elif agent.is_call_tools_node(node):, before you call async with node.stream(run.ctx) as handle_stream: to actually run the node (which will start executing tools), you can read node.model_response to get the ModelResponse, which contains (among other things) the tool calls:

async def main():
    async with agent.iter('call something please') as run:
        async for node in run:
            if agent.is_call_tools_node(node):
                for part in node.model_response.parts:
                    if isinstance(part, ToolCallPart):
                        return f'tool call requires approval: {part.tool_name}'

                async with node.stream(run.ctx) as handle_stream:
                    async for tool_event in handle_stream:
                        print(tool_event)


result = asyncio.run(main())
print(result)

Let me know if that helps!

[tibbe]

@DouweM that does help. What about the new_messages() part? When can I extract those? I need to extract them before running the tool node (i.e. just as I return f'tool call requires approval: {part.tool_name}') so I can resume later when I get the approval from the user.

From the documentation I gather than run.new_messages() are only available when the run finishes, which I guess it doesn't if we return early?

[restless]

@tibbe You may find #195 (comment) helpful. Not that it's a nice solution.

[tibbe]

@tibbe You may find #195 (comment) helpful. Not that it's a nice solution.

I see. I think at this point it's probably easier to just use pydantic_ai.models than trying to use the agent.

[tibbe]

Also noticed that if you try to "suspend" the tool call (like described above, by gathering the message history and then returning) and then try to run iter(..., message_history=messages) it fails:

Loaded 2 messages from store: [ModelRequest(parts=[SystemPromptPart(content='...', timestamp=datetime.datetime(2025, 6, 17, 8, 47, 41, 815952, tzinfo=datetime.timezone.utc)), UserPromptPart(content='...', timestamp=datetime.datetime(2025, 6, 17, 8, 47, 41, 815959, tzinfo=datetime.timezone.utc))]), ModelResponse(parts=[ToolCallPart(tool_name='get_projects', args='{}', tool_call_id='call_ssE8MNVAdbcij4niI1clZeZZ')], usage=Usage(requests=1, request_tokens=744, response_tokens=467, total_tokens=1211, details={'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 448, 'rejected_prediction_tokens': 0, 'cached_tokens': 0}), model_name='o4-mini', timestamp=datetime.datetime(2025, 6, 17, 8, 47, 42, tzinfo=TzInfo(UTC)))]
Executing node: UserPromptNode(user_prompt='...', instructions=None, instructions_functions=[], system_prompts=(), system_prompt_functions=[SystemPromptRunner(function=<function get_agent.<locals>.system_prompt at 0x11adeb920>, dynamic=False, _takes_ctx=True, _is_async=False)], system_prompt_dynamic_functions={})
Executing node: ModelRequestNode(request=ModelRequest(parts=[UserPromptPart(content='...', timestamp=datetime.datetime(2025, 6, 17, 8, 48, 30, 577726, tzinfo=datetime.timezone.utc))]))

So when I pass message_history iter is trying to executed nodes it already executed on the last invocation, which seems to defeat the purpose of message history.

[DouweM]

@tibbe Right, unfortunately there's currently no way to pick up the agent graph from a CallToolsNode in response to the ModelResponse passed in through message_history. I'm tempted to implement that (I don't think it'd be too hard), but we're actually planning to stop using a graph inside Agent: #1976. It would be useful to get your opinions there on the hooks you'll need and things you liked/disliked about the current graph approach.

[tibbe]

@DouweM since Agent currently can't do what I want I'm currently trying to write my own agent on top of e.g. messages, models, and tools. What I found in doing so is that I have to duplicate a lot of the heavy lifting done inside Agent, such as

• carrying around a RunContext,
• bookkeeping around e.g. Usage,
• initializing tools,
• doing parallel tool calls,
• creating a dynamic system prompt, and
• generally just dealing with the different streaming response parts.

I guess there's two ways to go:

• stick with an inversion-of-control approach, like the current Agent, and try to have enough hooks to allow the user to do what they need to do or
• take a more library approach and provide more building blocks to be able to more easily build your own agent (like I'm currently doing).

If the former just worked I would just use that but there's always a higher likelihood that that inversion-of-control approaches won't be configurable enough.

The latter would benefit from more internal utilities inside Agent becoming public utilities. The jump from Model to Agent is quite large at the moment.

For what it's worth here's a raw, unfinished snapshot of code I was writing for my agent today, to give an idea of what a user might try to do:

import os
import json
from typing import Sequence

from pydantic_ai.direct import model_request_stream
from pydantic_ai import messages
from pydantic_ai import models
from pydantic_ai import usage as _usage
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai import tools as _tools
from opentelemetry import trace

async def run(
    model: models.Model,
    user_prompt: str,
    deps: _tools.AgentDepsT,
    system_prompt: _tools.SystemPromptFunc[_tools.AgentDepsT] | None = None,
    tools: Sequence[_tools.Tool[_tools.AgentDepsT]] = (),
    usage: _usage.Usage | None = None,
) -> messages.ModelResponse:
    ctx = _tools.RunContext(
        deps=deps,
        model=model,
        usage=usage or _usage.Usage(),
        prompt=user_prompt,
        messages=[],
        run_step=0,
    )

    message_history: list[messages.ModelMessage] = [
        messages.ModelRequest(
            parts=[
                messages.SystemPromptPart(
                    content=system_prompt(ctx)
                    if system_prompt is not None
                    else None
                ),
                messages.UserPromptPart(content=user_prompt),
            ]
        ),
    ]
    model_request_parameters = models.ModelRequestParameters(
        function_tools=[await tool.prepare_tool_def(ctx) for tool in tools]
    )

    while True:
        print('Making model request:', message_history)
        async with model_request_stream(
            model,
            message_history,
            model_request_parameters=model_request_parameters,
        ) as stream:
            async for chunk in stream:
                if isinstance(chunk, messages.PartStartEvent):
                    if isinstance(chunk.part, messages.TextPart):
                        print(chunk.part.content, end='', flush=True)
                    elif isinstance(chunk.part, messages.ToolCallPart):
                        print(
                            f'Tool call: {chunk.part.tool_name} with args {chunk.part.args}',
                            flush=True,
                        )
                elif isinstance(chunk, messages.PartDeltaEvent):
                    if isinstance(chunk.delta, messages.TextPartDelta):
                        print(chunk.delta.content_delta, end='', flush=True)
            print()

        model_response = stream.get()
        message_history.append(model_response)
        tool_outputs = []
        for part in model_response.parts:
            if isinstance(part, messages.ToolCallPart):
                print('Tool call response')
                tool = next(
                    (t for t in tools if t.name == part.tool_name), None
                )
                if tool is None:
                    raise ValueError(f'Tool {part.tool_name} not found.')
                tool_response = await tool.run(part, ctx, trace.NoOpTracer())
                if isinstance(tool_response, messages.ToolReturnPart):
                    print(f'Tool response: {tool_response}')
                    tool_outputs.append(tool_response)
        if not tool_outputs:
            print('No tool outputs, ending run.')
            break
        message_history.append(
            messages.ModelRequest(parts=tool_outputs)
        )

# Test it:
def main():
    model = OpenAIModel(
        model_name='gpt-4o',
        provider=OpenAIProvider(api_key=api_key),
    )
    tools = [
        _tools.Tool(database_tools.get_projects),
        _tools.Tool(database_tools.get_tasks),
    ]

    def system_prompt(ctx: _tools.RunContext[database_tools.Dependencies]) -> str:
        return system_prompt_lib.PROMPT_TEMPLATE.format(
            current_date=ctx.deps.today.strftime('%B %d, %Y'),
            user_id=ctx.deps.current_user.UserId,
            user_name=name.full_display_name(
                ctx.deps.current_user.FirstName, ctx.deps.current_user.LastName
            ),
        )

    await run(
        model,
        USER_PROMPT,
        deps=deps,
        system_prompt=system_prompt,
        tools=tools,
    )

[tibbe]

@DouweM to be very concrete, I think that lots of stuff around registering tools, initializing them, etc. doesn't necessarily have anything to do with an agent per se, as a single Model can use tools. Also the tools seem a bit tied to RunContext, which is agent-specific from looking at its fields.

[tibbe]

Actually I wonder if having a pydantic_ai.direct.model_request_iter (similar to Agent.iter) would help. Seems natural and missing giving the other versions in direct.