Cannot handle consecutive model request/model response

[yf-yang]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

If message history contains multiple ModelRequest/ModelResponse of tool return/tool calls, then the framework fails to handle it correctly.

The cause is here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py

Lines 198 to 220 in bee76e6

	if (tool_call_results := ctx.deps.tool_call_results) is not None:
	if messages and (last_message := messages[-1]) and isinstance(last_message, _messages.ModelRequest):
	# If tool call results were provided, that means the previous run ended on deferred tool calls.
	# That run would typically have ended on a `ModelResponse`, but if it had a mix of deferred tool calls and ones that could already be executed,
	# a `ModelRequest` would already have been added to the history with the preliminary results, even if it wouldn't have been sent to the model yet.
	# So now that we have all of the deferred results, we roll back to the last `ModelResponse` and store the contents of the `ModelRequest` on `deferred_tool_results` to be handled by `CallToolsNode`.
	ctx.deps.tool_call_results = self._update_tool_call_results_from_model_request(
	tool_call_results, last_message
	)
	messages.pop()
	if not messages:
	raise exceptions.UserError('Tool call results were provided, but the message history is empty.')
	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):
	call_tools_node = await self._handle_message_history_model_response(ctx, last_message)
	if call_tools_node is not None:
	return call_tools_node

Both blocks are merely one if, instead of a loop, indicating it assumes there should be at most one request/response.

However I don't know what's the proper place to fix it. I personally think there should be a history merge step to merge consecutive requests/responses somewhere.

Example Code

import asyncio

from dotenv import load_dotenv
from pydantic_ai import (
  Agent,
  CallDeferred,
  DeferredToolRequests,
  DeferredToolResults,
)
from pydantic_ai.messages import (
  ModelRequest,
  ModelResponse,
  ToolCallPart,
  ToolReturnPart,
  UserPromptPart,
)

load_dotenv()

agent = Agent("anthropic:claude-sonnet-4-0", output_type=[str, DeferredToolRequests])


@agent.tool_plain
async def sum(a: int, b: int) -> int:
  return a + b


@agent.tool_plain
async def multiply(a: int, b: int) -> int:
  raise CallDeferred


async def main():
  # Successful case
  # messages = [
  #   ModelRequest(parts=[UserPromptPart(content="Calculate 1+2, 3+4, 5*6, 7*8")]),
  #   ModelResponse(
  #     parts=[
  #       ToolCallPart(tool_name="sum", args={"a": 1, "b": 2}, tool_call_id="sum-1"),
  #       ToolCallPart(tool_name="sum", args={"a": 3, "b": 4}, tool_call_id="sum-2"),
  #       ToolCallPart(tool_name="multiply", args={"a": 5, "b": 6}, tool_call_id="multiply-1"),
  #       ToolCallPart(tool_name="multiply", args={"a": 7, "b": 8}, tool_call_id="multiply-2"),
  #     ]
  #   ),
  #   ModelRequest(
  #     parts=[
  #       ToolReturnPart(tool_name="sum", tool_call_id="sum-1", content="3"),
  #       ToolReturnPart(tool_name="sum", tool_call_id="sum-2", content="7"),
  #     ]
  #   ),
  # ]

  # Failed case 1
  messages = [
    ModelRequest(parts=[UserPromptPart(content="Calculate 1+2, 3+4, 5*6, 7*8")]),
    ModelResponse(
      parts=[
        ToolCallPart(tool_name="sum", args={"a": 1, "b": 2}, tool_call_id="sum-1"),
        ToolCallPart(tool_name="sum", args={"a": 3, "b": 4}, tool_call_id="sum-2"),
        ToolCallPart(tool_name="multiply", args={"a": 5, "b": 6}, tool_call_id="multiply-1"),
        ToolCallPart(tool_name="multiply", args={"a": 7, "b": 8}, tool_call_id="multiply-2"),
      ]
    ),
    ModelRequest(
      parts=[
        ToolReturnPart(tool_name="sum", tool_call_id="sum-1", content="3"),
      ]
    ),
    ModelRequest(
      parts=[
        ToolReturnPart(tool_name="sum", tool_call_id="sum-2", content="7"),
      ]
    ),
  ]

  # Failed case 2
  messages = [
    ModelRequest(parts=[UserPromptPart(content="Calculate 1+2, 3+4, 5*6, 7*8")]),
    ModelResponse(parts=[ToolCallPart(tool_name="sum", args={"a": 1, "b": 2}, tool_call_id="sum-1")]),
    ModelResponse(
      parts=[
        ToolCallPart(tool_name="sum", args={"a": 3, "b": 4}, tool_call_id="sum-2"),
      ]
    ),
    ModelResponse(
      parts=[
        ToolCallPart(tool_name="multiply", args={"a": 5, "b": 6}, tool_call_id="multiply-1"),
      ]
    ),
    ModelResponse(
      parts=[
        ToolCallPart(tool_name="multiply", args={"a": 7, "b": 8}, tool_call_id="multiply-2"),
      ]
    ),
    ModelRequest(
      parts=[
        ToolReturnPart(tool_name="sum", tool_call_id="sum-1", content="3"),
        ToolReturnPart(tool_name="sum", tool_call_id="sum-2", content="7"),
      ]
    ),
  ]

  tool_results = DeferredToolResults(calls={"multiply-1": "30", "multiply-2": "56"})

  result = await agent.run(message_history=messages, deferred_tool_results=tool_results)
  print(result.output)
  print(result.all_messages())


if __name__ == "__main__":
  asyncio.run(main())

Python, Pydantic AI & LLM client version

1.0.3

[DouweM]

@yf-yang Under what circumstances are you getting a message history with a number of back-to-back responses or requests, as opposed to the typical request/response/request/response/etc? You're right that that isn't currently supported, but that's because I didn't think this pattern existed in the wild. So it'd be good to understand when it does.

[yf-yang]

@DouweM Inspired by ag-ui, I create a wrapper around "parts" with additional data and ser/deser them to/from db. Therefore a message history is now saved as a list of parts. When passing to a new agent call, it is wrapper again with ModelRequest/ModelResponse via a map function.

Actually I can handle it from my side, but I think the framework should be able to handle it.

There are some small dx benefits, although that's not too much. Let me illustrate a few scenarios:

1. Frontend, in a frontend UI the messages are clearly organized by a list of different part types, not model request/model response.
2. Double texting. Since some model disallow a bare thinking assistant message as the last message, I can insert a modelresponse/textpart with zero length content, without needing to check if last message is an assistant messae

[yf-yang]

pydantic-ai/pydantic_ai_slim/pydantic_ai/ag_ui.py

Lines 572 to 620 in 619d2aa

	def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
	"""Convert a AG-UI history to a Pydantic AI one."""
	result: list[ModelMessage] = []
	tool_calls: dict[str, str] = {} # Tool call ID to tool name mapping.
	for msg in messages:
	if isinstance(msg, UserMessage):
	result.append(ModelRequest(parts=[UserPromptPart(content=msg.content)]))
	elif isinstance(msg, AssistantMessage):
	if msg.content:
	result.append(ModelResponse(parts=[TextPart(content=msg.content)]))
	if msg.tool_calls:
	for tool_call in msg.tool_calls:
	tool_calls[tool_call.id] = tool_call.function.name
	result.append(
	ModelResponse(
	parts=[
	ToolCallPart(
	tool_name=tool_call.function.name,
	tool_call_id=tool_call.id,
	args=tool_call.function.arguments,
	)
	for tool_call in msg.tool_calls
	]
	)
	)
	elif isinstance(msg, SystemMessage):
	result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
	elif isinstance(msg, ToolMessage):
	tool_name = tool_calls.get(msg.tool_call_id)
	if tool_name is None: # pragma: no cover
	raise _ToolCallNotFoundError(tool_call_id=msg.tool_call_id)
	result.append(
	ModelRequest(
	parts=[
	ToolReturnPart(
	tool_name=tool_name,
	content=msg.content,
	tool_call_id=msg.tool_call_id,
	)
	]
	)
	)
	elif isinstance(msg, DeveloperMessage): # pragma: no branch
	result.append(ModelRequest(parts=[SystemPromptPart(content=msg.content)]))
	return result

Current ag-ui implementation also takes a hybrid approach. Tool calls are merged, tool results are not, tool calls and texts are split, thinking is not supported. Anyway I like ag-ui's concept but it's implementation is not mature enough.

[DouweM]

@yf-yang Thanks, that makes sense.

I personally think there should be a history merge step to merge consecutive requests/responses somewhere.

I agree. The easiest place to do that would be in Agent.iter, where we currently pass message_history and new_message_index into the graph, since both would need to be in sync. (We should probably also be updating new_message_index in that "rewind" logic in UserPromptNode.run...)

We'd need to take special care of ModelResponse.provider_response_id -- I think we should only merge responses that have matching IDs (either both None or both the same string), as these IDs are meaningful for example in #2756.

Would you be up for submitting a PR?

[yf-yang]

I can, would you mind explaining this one in detail? I know what new_message_index does but I would like to know how you would like to implement that.

We should probably also be updating new_message_index in that "rewind" logic in UserPromptNode.run...

[DouweM]

@yf-yang The expectation is that if you're keeping track of the message history in a variable messages, you can do a result = agent.run(message_history=messages) call and then update the history for the next call either by setting messages = result.all_messages() or messages += result.new_messages().

new_messages() is all_messages()[new_message_index:], so we need to make sure that new_message_index points at the right place, so that even if we pop things off messages inside the UserPromptNode, the new messages the user gets can still be appended onto the messages they already had to get an equivalent history to all_messages().

That's currently not working correctly in this test for deferred_tool_calls with a message history ending in ModelRequest that already had 1 tool call result:

pydantic-ai/tests/test_agent.py

Lines 4553 to 4701 in 8d4889c

	async def test_hitl_tool_approval():
	def model_function(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
	if len(messages) == 1:
	return ModelResponse(
	parts=[
	ToolCallPart(
	tool_name='create_file',
	args={'path': 'new_file.py', 'content': 'print("Hello, world!")'},
	tool_call_id='create_file',
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'ok_to_delete.py'}, tool_call_id='ok_to_delete'
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'never_delete.py'}, tool_call_id='never_delete'
	),
	]
	)
	else:
	return ModelResponse(parts=[TextPart('Done!')])
	model = FunctionModel(model_function)
	agent = Agent(model, output_type=[str, DeferredToolRequests])
	@agent.tool_plain(requires_approval=True)
	def delete_file(path: str) -> str:
	return f'File {path!r} deleted'
	@agent.tool_plain
	def create_file(path: str, content: str) -> str:
	return f'File {path!r} created with content: {content}'
	result = await agent.run('Create new_file.py and delete ok_to_delete.py and never_delete.py')
	messages = result.all_messages()
	assert messages == snapshot(
	[
	ModelRequest(
	parts=[
	UserPromptPart(
	content='Create new_file.py and delete ok_to_delete.py and never_delete.py',
	timestamp=IsDatetime(),
	)
	]
	),
	ModelResponse(
	parts=[
	ToolCallPart(
	tool_name='create_file',
	args={'path': 'new_file.py', 'content': 'print("Hello, world!")'},
	tool_call_id='create_file',
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'ok_to_delete.py'}, tool_call_id='ok_to_delete'
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'never_delete.py'}, tool_call_id='never_delete'
	),
	],
	usage=RequestUsage(input_tokens=60, output_tokens=23),
	model_name='function:model_function:',
	timestamp=IsDatetime(),
	),
	ModelRequest(
	parts=[
	ToolReturnPart(
	tool_name='create_file',
	content='File \'new_file.py\' created with content: print("Hello, world!")',
	tool_call_id='create_file',
	timestamp=IsDatetime(),
	)
	]
	),
	]
	)
	assert result.output == snapshot(
	DeferredToolRequests(
	approvals=[
	ToolCallPart(tool_name='delete_file', args={'path': 'ok_to_delete.py'}, tool_call_id='ok_to_delete'),
	ToolCallPart(tool_name='delete_file', args={'path': 'never_delete.py'}, tool_call_id='never_delete'),
	]
	)
	)
	result = await agent.run(
	message_history=messages,
	deferred_tool_results=DeferredToolResults(
	approvals={'ok_to_delete': True, 'never_delete': ToolDenied('File cannot be deleted')},
	),
	)
	assert result.all_messages() == snapshot(
	[
	ModelRequest(
	parts=[
	UserPromptPart(
	content='Create new_file.py and delete ok_to_delete.py and never_delete.py',
	timestamp=IsDatetime(),
	)
	]
	),
	ModelResponse(
	parts=[
	ToolCallPart(
	tool_name='create_file',
	args={'path': 'new_file.py', 'content': 'print("Hello, world!")'},
	tool_call_id='create_file',
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'ok_to_delete.py'}, tool_call_id='ok_to_delete'
	),
	ToolCallPart(
	tool_name='delete_file', args={'path': 'never_delete.py'}, tool_call_id='never_delete'
	),
	],
	usage=RequestUsage(input_tokens=60, output_tokens=23),
	model_name='function:model_function:',
	timestamp=IsDatetime(),
	),
	ModelRequest(
	parts=[
	ToolReturnPart(
	tool_name='create_file',
	content='File \'new_file.py\' created with content: print("Hello, world!")',
	tool_call_id='create_file',
	timestamp=IsDatetime(),
	),
	ToolReturnPart(
	tool_name='delete_file',
	content="File 'ok_to_delete.py' deleted",
	tool_call_id='ok_to_delete',
	timestamp=IsDatetime(),
	),
	ToolReturnPart(
	tool_name='delete_file',
	content='File cannot be deleted',
	tool_call_id='never_delete',
	timestamp=IsDatetime(),
	),
	]
	),
	ModelResponse(
	parts=[TextPart(content='Done!')],
	usage=RequestUsage(input_tokens=78, output_tokens=24),
	model_name='function:model_function:',
	timestamp=IsDatetime(),
	),
	]
	)
	assert result.output == snapshot('Done!')

If I call result.new_messages() there, it returns just the final ModelResponse, meaning that the 2 new ToolReturnParts we added to the last ModelRequest in history are not included, and would be missing from their message history if they used messages += result.new_messages().

We could do new_message_index -= 1 when we do messages.pop(), but then their message history would end up with a duplicated ToolReturnPart, because the existing history ends in this:

ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='create_file',
                        content='File \'new_file.py\' created with content: print("Hello, world!")',
                        tool_call_id='create_file',
                        timestamp=IsDatetime(),
                    )
                ]
            ),

And the new one starts with this:

ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='create_file',
                        content='File \'new_file.py\' created with content: print("Hello, world!")',
                        tool_call_id='create_file',
                        timestamp=IsDatetime(),
                    ),
                    ToolReturnPart(
                        tool_name='delete_file',
                        content="File 'ok_to_delete.py' deleted",
                        tool_call_id='ok_to_delete',
                        timestamp=IsDatetime(),
                    ),
                    ToolReturnPart(
                        tool_name='delete_file',
                        content='File cannot be deleted',
                        tool_call_id='never_delete',
                        timestamp=IsDatetime(),
                    ),
                ]
            ),

I think we have 2 options to fix this:

1. Decrement the new_message_index, accept that there could be multiple ToolReturnParts for the same tool_call_id in the message history, and in your new merge logic, filter those out (assuming their content is the same) so that the merged history only has one.
2. Modify the deferred_tool_results logic to not replace the incomplete ModelRequest (with 1 return part) with a complete one (with 3 return parts), but to instead add a new ModelRequest with the 2 new return parts, so that new_messages() will contain that + the final ModelResponse, and there will be no duplication of ToolCallParts. There will be 2 consecutive ModelRequests, but your merging logic will deal with that.

Solution 2 is the cleanest, but may be tricky to implement. I'll have a look at that, but I may end up asking you to implement solution 1 in in your PR instead.