pydantic-ai-slim 1.0.5__tar.gz → 1.0.7__tar.gz

{pydantic_ai_slim-1.0.5 → pydantic_ai_slim-1.0.7}/.gitignore

@@ -20,3 +20,4 @@ node_modules/
 .coverage*
 /test_tmp/
 .mcp.json
+.claude/

{pydantic_ai_slim-1.0.5 → pydantic_ai_slim-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 1.0.5
+Version: 1.0.7
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Project-URL: Homepage, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
 Project-URL: Source, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
@@ -33,7 +33,7 @@ Requires-Dist: genai-prices>=0.0.23
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==1.0.5
+Requires-Dist: pydantic-graph==1.0.7
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -51,13 +51,13 @@ Requires-Dist: prompt-toolkit>=3; extra == 'cli'
 Requires-Dist: pyperclip>=1.9.0; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
-Requires-Dist: cohere>=5.17.0; (platform_system != 'Emscripten') and extra == 'cohere'
+Requires-Dist: cohere>=5.18.0; (platform_system != 'Emscripten') and extra == 'cohere'
 Provides-Extra: dbos
 Requires-Dist: dbos>=1.13.0; extra == 'dbos'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==1.0.5; extra == 'evals'
+Requires-Dist: pydantic-evals==1.0.7; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.31.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-1.0.5 → pydantic_ai_slim-1.0.7}/pydantic_ai/_agent_graph.py

@@ -8,7 +8,7 @@ from collections import defaultdict, deque
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
-from dataclasses import field
+from dataclasses import field, replace
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast
 
 from opentelemetry.trace import Tracer
@@ -16,7 +16,7 @@ from typing_extensions import TypeVar, assert_never
 
 from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
 from pydantic_ai._tool_manager import ToolManager
-from pydantic_ai._utils import is_async_callable, run_in_executor
+from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, is_async_callable, run_in_executor
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
 from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
@@ -26,7 +26,9 @@ from .exceptions import ToolRetryError
 from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings
 from .tools import (
+    DeferredToolCallResult,
     DeferredToolResult,
+    DeferredToolResults,
     RunContext,
     ToolApproved,
     ToolDefinition,
@@ -123,7 +125,6 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     builtin_tools: list[AbstractBuiltinTool] = dataclasses.field(repr=False)
     tool_manager: ToolManager[DepsT]
-    tool_call_results: dict[str, DeferredToolResult] | None
 
     tracer: Tracer
     instrumentation_settings: InstrumentationSettings | None
@@ -160,14 +161,18 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
 
     _: dataclasses.KW_ONLY
 
-    instructions: str | None
-    instructions_functions: list[_system_prompt.SystemPromptRunner[DepsT]]
+    deferred_tool_results: DeferredToolResults | None = None
 
-    system_prompts: tuple[str, ...]
-    system_prompt_functions: list[_system_prompt.SystemPromptRunner[DepsT]]
-    system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[DepsT]]
+    instructions: str | None = None
+    instructions_functions: list[_system_prompt.SystemPromptRunner[DepsT]] = dataclasses.field(default_factory=list)
 
-    async def run(
+    system_prompts: tuple[str, ...] = dataclasses.field(default_factory=tuple)
+    system_prompt_functions: list[_system_prompt.SystemPromptRunner[DepsT]] = dataclasses.field(default_factory=list)
+    system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[DepsT]] = dataclasses.field(
+        default_factory=dict
+    )
+
+    async def run(  # noqa: C901
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> ModelRequestNode[DepsT, NodeRunEndT] | CallToolsNode[DepsT, NodeRunEndT]:
         try:
@@ -181,119 +186,127 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
                 messages = ctx_messages.messages
                 ctx_messages.used = True
 
+        message_history = _clean_message_history(ctx.state.message_history)
         # Add message history to the `capture_run_messages` list, which will be empty at this point
-        messages.extend(ctx.state.message_history)
+        messages.extend(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
+        ctx.deps.new_message_index = len(messages)
 
-        run_context = build_run_context(ctx)
-
-        parts: list[_messages.ModelRequestPart] = []
-        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 (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 self.deferred_tool_results is not None:
+            return await self._handle_deferred_tool_results(self.deferred_tool_results, messages, ctx)
 
-            if not messages:
-                raise exceptions.UserError('Tool call results were provided, but the message history is empty.')
+        next_message: _messages.ModelRequest | None = None
 
         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)
+                next_message = _messages.ModelRequest(parts=last_message.parts)
+
+                # Extract `UserPromptPart` content from the popped message and add to `ctx.deps.prompt`
+                user_prompt_parts = [part for part in last_message.parts if isinstance(part, _messages.UserPromptPart)]
+                if user_prompt_parts:
+                    if len(user_prompt_parts) == 1:
+                        ctx.deps.prompt = user_prompt_parts[0].content
+                    else:
+                        combined_content: list[_messages.UserContent] = []
+                        for part in user_prompt_parts:
+                            if isinstance(part.content, str):
+                                combined_content.append(part.content)
+                            else:
+                                combined_content.extend(part.content)
+                        ctx.deps.prompt = combined_content
             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
+                if self.user_prompt is None:
+                    # Skip ModelRequestNode and go directly to CallToolsNode
+                    return CallToolsNode[DepsT, NodeRunEndT](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 contains unprocessed tool calls.'
+                    )
 
-        if self.user_prompt is not None:
-            parts.append(_messages.UserPromptPart(self.user_prompt))
+        # Build the run context after `ctx.deps.prompt` has been updated
+        run_context = build_run_context(ctx)
 
-        instructions = await ctx.deps.get_instructions(run_context)
-        next_message = _messages.ModelRequest(parts, instructions=instructions)
+        parts: list[_messages.ModelRequestPart] = []
+        if messages:
+            await self._reevaluate_dynamic_prompts(messages, run_context)
 
-        return ModelRequestNode[DepsT, NodeRunEndT](request=next_message)
+        if next_message:
+            await self._reevaluate_dynamic_prompts([next_message], run_context)
+        else:
+            parts: list[_messages.ModelRequestPart] = []
+            if not messages:
+                parts.extend(await self._sys_parts(run_context))
 
-    async def _handle_message_history_model_response(
-        self,
-        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
-        message: _messages.ModelResponse,
-    ) -> CallToolsNode[DepsT, NodeRunEndT] | None:
-        unprocessed_tool_calls = any(isinstance(part, _messages.ToolCallPart) for part in message.parts)
-        if unprocessed_tool_calls:
             if self.user_prompt is not None:
-                raise exceptions.UserError(
-                    'Cannot provide a new user prompt when the message history contains unprocessed tool calls.'
-                )
-        else:
-            if ctx.deps.tool_call_results is not None:
-                raise exceptions.UserError(
-                    'Tool call results were provided, but the message history does not contain any unprocessed tool calls.'
-                )
+                parts.append(_messages.UserPromptPart(self.user_prompt))
 
-        if unprocessed_tool_calls or 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
-            run_context = build_run_context(ctx)
-            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=message)
-
-    def _update_tool_call_results_from_model_request(
-        self, tool_call_results: dict[str, DeferredToolResult], message: _messages.ModelRequest
-    ) -> dict[str, DeferredToolResult]:
-        last_tool_return: _messages.ToolReturn | None = None
-        user_content: list[str | _messages.UserContent] = []
-        for part in message.parts:
-            if isinstance(part, _messages.ToolReturnPart):
-                if part.tool_call_id in tool_call_results:
-                    raise exceptions.UserError(
-                        f'Tool call {part.tool_call_id!r} was already executed and its result cannot be overridden.'
-                    )
+            next_message = _messages.ModelRequest(parts=parts)
 
-                last_tool_return = _messages.ToolReturn(return_value=part.content, metadata=part.metadata)
-                tool_call_results[part.tool_call_id] = last_tool_return
-            elif isinstance(part, _messages.RetryPromptPart):
-                if part.tool_call_id in tool_call_results:
-                    raise exceptions.UserError(
-                        f'Tool call {part.tool_call_id!r} was already executed and its result cannot be overridden.'
-                    )
+        next_message.instructions = await ctx.deps.get_instructions(run_context)
 
-                tool_call_results[part.tool_call_id] = part
-            elif isinstance(part, _messages.UserPromptPart):
-                # Tools can return user parts via `ToolReturn.content` or by returning multi-modal content.
-                # These go together with a specific `ToolReturnPart`, but we don't have a way to know which,
-                # so (below) we just add them to the last one, matching the tool-results-before-user-parts order of the request.
-                if isinstance(part.content, str):
-                    user_content.append(part.content)
-                else:
-                    user_content.extend(part.content)
-            else:
-                raise exceptions.UserError(f'Unexpected message part type: {type(part)}')  # pragma: no cover
+        return ModelRequestNode[DepsT, NodeRunEndT](request=next_message)
 
-        if user_content:
-            if last_tool_return is None:
-                raise exceptions.UserError(
-                    'Tool call results were provided, but the last message in the history was a `ModelRequest` with user parts not tied to preliminary tool results.'
-                )
-            assert last_tool_return is not None
-            last_tool_return.content = user_content
+    async def _handle_deferred_tool_results(  # noqa: C901
+        self,
+        deferred_tool_results: DeferredToolResults,
+        messages: list[_messages.ModelMessage],
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+    ) -> CallToolsNode[DepsT, NodeRunEndT]:
+        if not messages:
+            raise exceptions.UserError('Tool call results were provided, but the message history is empty.')
+
+        last_model_request: _messages.ModelRequest | None = None
+        last_model_response: _messages.ModelResponse | None = None
+        for message in reversed(messages):
+            if isinstance(message, _messages.ModelRequest):
+                last_model_request = message
+            elif isinstance(message, _messages.ModelResponse):  # pragma: no branch
+                last_model_response = message
+                break
+
+        if not last_model_response:
+            raise exceptions.UserError(
+                'Tool call results were provided, but the message history does not contain a `ModelResponse`.'
+            )
+        if not any(isinstance(part, _messages.ToolCallPart) for part in last_model_response.parts):
+            raise exceptions.UserError(
+                'Tool call results were provided, but the message history does not contain any unprocessed tool calls.'
+            )
+        if self.user_prompt is not None:
+            raise exceptions.UserError(
+                'Cannot provide a new user prompt when the message history contains unprocessed tool calls.'
+            )
+
+        tool_call_results: dict[str, DeferredToolResult | Literal['skip']] | None = None
+        tool_call_results = {}
+        for tool_call_id, approval in deferred_tool_results.approvals.items():
+            if approval is True:
+                approval = ToolApproved()
+            elif approval is False:
+                approval = ToolDenied()
+            tool_call_results[tool_call_id] = approval
+
+        if calls := deferred_tool_results.calls:
+            call_result_types = get_union_args(DeferredToolCallResult)
+            for tool_call_id, result in calls.items():
+                if not isinstance(result, call_result_types):
+                    result = _messages.ToolReturn(result)
+                tool_call_results[tool_call_id] = result
+
+        if last_model_request:
+            for part in last_model_request.parts:
+                if isinstance(part, _messages.ToolReturnPart | _messages.RetryPromptPart):
+                    if part.tool_call_id in tool_call_results:
+                        raise exceptions.UserError(
+                            f'Tool call {part.tool_call_id!r} was already executed and its result cannot be overridden.'
+                        )
+                    tool_call_results[part.tool_call_id] = 'skip'
 
-        return tool_call_results
+        # Skip ModelRequestNode and go directly to CallToolsNode
+        return CallToolsNode[DepsT, NodeRunEndT](last_model_response, tool_call_results=tool_call_results)
 
     async def _reevaluate_dynamic_prompts(
         self, messages: list[_messages.ModelMessage], run_context: RunContext[DepsT]
@@ -330,6 +343,8 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
                 messages.append(_messages.SystemPromptPart(prompt))
         return messages
 
+    __repr__ = dataclasses_no_defaults_repr
+
 
 async def _prepare_request_parameters(
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
@@ -441,6 +456,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         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)
+        message_history = _clean_message_history(message_history)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
@@ -476,12 +492,15 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
         return self._result
 
+    __repr__ = dataclasses_no_defaults_repr
+
 
 @dataclasses.dataclass
 class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
     """The node that processes a model response, and decides whether to end the run or make a new request."""
 
     model_response: _messages.ModelResponse
+    tool_call_results: dict[str, DeferredToolResult | Literal['skip']] | None = None
 
     _events_iterator: AsyncIterator[_messages.HandleResponseEvent] | None = field(default=None, init=False, repr=False)
     _next_node: ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]] | None = field(
@@ -582,11 +601,20 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
     ) -> AsyncIterator[_messages.HandleResponseEvent]:
         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_parts: list[_messages.ModelRequestPart] = []
         output_final_result: deque[result.FinalResult[NodeRunEndT]] = deque(maxlen=1)
 
-        async for event in process_function_tools(
-            ctx.deps.tool_manager, tool_calls, None, ctx, output_parts, output_final_result
+        async for event in process_tool_calls(
+            tool_manager=ctx.deps.tool_manager,
+            tool_calls=tool_calls,
+            tool_call_results=self.tool_call_results,
+            final_result=None,
+            ctx=ctx,
+            output_parts=output_parts,
+            output_final_result=output_final_result,
         ):
             yield event
 
@@ -639,6 +667,8 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         else:
             return self._handle_final_result(ctx, result.FinalResult(result_data), [])
 
+    __repr__ = dataclasses_no_defaults_repr
+
 
 def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]]) -> RunContext[DepsT]:
     """Build a `RunContext` object from the current agent graph run context."""
@@ -652,13 +682,14 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         trace_include_content=ctx.deps.instrumentation_settings is not None
         and ctx.deps.instrumentation_settings.include_content,
         run_step=ctx.state.run_step,
-        tool_call_approved=ctx.state.run_step == 0 and ctx.deps.tool_call_results is not None,
+        tool_call_approved=ctx.state.run_step == 0,
     )
 
 
-async def process_function_tools(  # noqa: C901
+async def process_tool_calls(  # noqa: C901
     tool_manager: ToolManager[DepsT],
     tool_calls: list[_messages.ToolCallPart],
+    tool_call_results: dict[str, DeferredToolResult | Literal['skip']] | None,
     final_result: result.FinalResult[NodeRunEndT] | None,
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
     output_parts: list[_messages.ModelRequestPart],
@@ -739,14 +770,13 @@ async def process_function_tools(  # noqa: C901
         ctx.state.increment_retries(ctx.deps.max_result_retries)
         calls_to_run.extend(tool_calls_by_kind['unknown'])
 
-    deferred_tool_results: dict[str, DeferredToolResult] = {}
-    if build_run_context(ctx).tool_call_approved and ctx.deps.tool_call_results is not None:
-        deferred_tool_results = ctx.deps.tool_call_results
+    calls_to_run_results: dict[str, DeferredToolResult] = {}
+    if tool_call_results is not None:
         # Deferred tool calls are "run" as well, by reading their value from the tool call results
         calls_to_run.extend(tool_calls_by_kind['external'])
         calls_to_run.extend(tool_calls_by_kind['unapproved'])
 
-        result_tool_call_ids = set(deferred_tool_results.keys())
+        result_tool_call_ids = set(tool_call_results.keys())
         tool_call_ids_to_run = {call.tool_call_id for call in calls_to_run}
         if tool_call_ids_to_run != result_tool_call_ids:
             raise exceptions.UserError(
@@ -754,24 +784,29 @@ async def process_function_tools(  # noqa: C901
                 f'Expected: {tool_call_ids_to_run}, got: {result_tool_call_ids}'
             )
 
+        # Filter out calls that were already executed before and should now be skipped
+        calls_to_run_results = {call_id: result for call_id, result in tool_call_results.items() if result != 'skip'}
+        calls_to_run = [call for call in calls_to_run if call.tool_call_id in calls_to_run_results]
+
     deferred_calls: dict[Literal['external', 'unapproved'], list[_messages.ToolCallPart]] = defaultdict(list)
 
     if calls_to_run:
         async for event in _call_tools(
-            tool_manager,
-            calls_to_run,
-            deferred_tool_results,
-            ctx.deps.tracer,
-            ctx.deps.usage_limits,
-            output_parts,
-            deferred_calls,
+            tool_manager=tool_manager,
+            tool_calls=calls_to_run,
+            tool_call_results=calls_to_run_results,
+            tracer=ctx.deps.tracer,
+            usage_limits=ctx.deps.usage_limits,
+            output_parts=output_parts,
+            output_deferred_calls=deferred_calls,
         ):
             yield event
 
     # Finally, we handle deferred tool calls (unless they were already included in the run because results were provided)
-    if not deferred_tool_results:
+    if tool_call_results is None:
+        calls = [*tool_calls_by_kind['external'], *tool_calls_by_kind['unapproved']]
         if final_result:
-            for call in [*tool_calls_by_kind['external'], *tool_calls_by_kind['unapproved']]:
+            for call in calls:
                 output_parts.append(
                     _messages.ToolReturnPart(
                         tool_name=call.tool_name,
@@ -779,13 +814,11 @@ async def process_function_tools(  # noqa: C901
                         tool_call_id=call.tool_call_id,
                     )
                 )
-        else:
-            for call in tool_calls_by_kind['external']:
-                deferred_calls['external'].append(call)
-                yield _messages.FunctionToolCallEvent(call)
+        elif calls:
+            deferred_calls['external'].extend(tool_calls_by_kind['external'])
+            deferred_calls['unapproved'].extend(tool_calls_by_kind['unapproved'])
 
-            for call in tool_calls_by_kind['unapproved']:
-                deferred_calls['unapproved'].append(call)
+            for call in calls:
                 yield _messages.FunctionToolCallEvent(call)
 
     if not final_result and deferred_calls:
@@ -807,7 +840,7 @@ async def process_function_tools(  # noqa: C901
 async def _call_tools(
     tool_manager: ToolManager[DepsT],
     tool_calls: list[_messages.ToolCallPart],
-    deferred_tool_results: dict[str, DeferredToolResult],
+    tool_call_results: dict[str, DeferredToolResult],
     tracer: Tracer,
     usage_limits: _usage.UsageLimits | None,
     output_parts: list[_messages.ModelRequestPart],
@@ -853,7 +886,7 @@ async def _call_tools(
         if tool_manager.should_call_sequentially(tool_calls):
             for index, call in enumerate(tool_calls):
                 if event := await handle_call_or_result(
-                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id), usage_limits),
                     index,
                 ):
                     yield event
@@ -861,7 +894,7 @@ async def _call_tools(
         else:
             tasks = [
                 asyncio.create_task(
-                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id), usage_limits),
                     name=call.tool_name,
                 )
                 for call in tool_calls
@@ -1079,3 +1112,51 @@ async def _process_message_history(
     # Replaces the message history in the state with the processed messages
     state.message_history = messages
     return messages
+
+
+def _clean_message_history(messages: list[_messages.ModelMessage]) -> list[_messages.ModelMessage]:
+    """Clean the message history by merging consecutive messages of the same type."""
+    clean_messages: list[_messages.ModelMessage] = []
+    for message in messages:
+        last_message = clean_messages[-1] if len(clean_messages) > 0 else None
+
+        if isinstance(message, _messages.ModelRequest):
+            if (
+                last_message
+                and isinstance(last_message, _messages.ModelRequest)
+                # Requests can only be merged if they have the same instructions
+                and (
+                    not last_message.instructions
+                    or not message.instructions
+                    or last_message.instructions == message.instructions
+                )
+            ):
+                parts = [*last_message.parts, *message.parts]
+                parts.sort(
+                    # Tool return parts always need to be at the start
+                    key=lambda x: 0 if isinstance(x, _messages.ToolReturnPart | _messages.RetryPromptPart) else 1
+                )
+                merged_message = _messages.ModelRequest(
+                    parts=parts,
+                    instructions=last_message.instructions or message.instructions,
+                )
+                clean_messages[-1] = merged_message
+            else:
+                clean_messages.append(message)
+        elif isinstance(message, _messages.ModelResponse):  # pragma: no branch
+            if (
+                last_message
+                and isinstance(last_message, _messages.ModelResponse)
+                # Responses can only be merged if they didn't really come from an API
+                and last_message.provider_response_id is None
+                and last_message.provider_name is None
+                and last_message.model_name is None
+                and message.provider_response_id is None
+                and message.provider_name is None
+                and message.model_name is None
+            ):
+                merged_message = replace(last_message, parts=[*last_message.parts, *message.parts])
+                clean_messages[-1] = merged_message
+            else:
+                clean_messages.append(message)
+    return clean_messages

{pydantic_ai_slim-1.0.5 → pydantic_ai_slim-1.0.7}/pydantic_ai/ag_ui.py

@@ -23,6 +23,7 @@ from typing import (
 )
 
 from pydantic import BaseModel, ValidationError
+from typing_extensions import assert_never
 
 from . import _utils
 from ._agent_graph import CallToolsNode, ModelRequestNode
@@ -32,7 +33,9 @@ from .messages import (
     FunctionToolResultEvent,
     ModelMessage,
     ModelRequest,
+    ModelRequestPart,
     ModelResponse,
+    ModelResponsePart,
     ModelResponseStreamEvent,
     PartDeltaEvent,
     PartStartEvent,
@@ -573,49 +576,57 @@ 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.
+    request_parts: list[ModelRequestPart] | None = None
+    response_parts: list[ModelResponsePart] | None = None
     for msg in messages:
-        if isinstance(msg, UserMessage):
-            result.append(ModelRequest(parts=[UserPromptPart(content=msg.content)]))
+        if isinstance(msg, UserMessage | SystemMessage | DeveloperMessage | ToolMessage):
+            if request_parts is None:
+                request_parts = []
+                result.append(ModelRequest(parts=request_parts))
+                response_parts = None
+
+            if isinstance(msg, UserMessage):
+                request_parts.append(UserPromptPart(content=msg.content))
+            elif isinstance(msg, SystemMessage | DeveloperMessage):
+                request_parts.append(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)
+
+                request_parts.append(
+                    ToolReturnPart(
+                        tool_name=tool_name,
+                        content=msg.content,
+                        tool_call_id=msg.tool_call_id,
+                    )
+                )
+            else:
+                assert_never(msg)
+
         elif isinstance(msg, AssistantMessage):
+            if response_parts is None:
+                response_parts = []
+                result.append(ModelResponse(parts=response_parts))
+                request_parts = None
+
             if msg.content:
-                result.append(ModelResponse(parts=[TextPart(content=msg.content)]))
+                response_parts.append(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
-                        ]
+                response_parts.extend(
+                    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)]))
+        else:
+            assert_never(msg)
 
     return result