pydantic-ai-slim 0.0.21__py3-none-any.whl → 0.0.22__py3-none-any.whl

pydantic_ai/agent.py

@@ -5,14 +5,17 @@ import dataclasses
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
-from contextvars import ContextVar
 from types import FrameType
-from typing import Any, Callable, Generic, Literal, cast, final, overload
+from typing import Any, Callable, Generic, cast, final, overload
 
 import logfire_api
-from typing_extensions import TypeVar, assert_never, deprecated
+from typing_extensions import TypeVar, deprecated
+
+from pydantic_graph import Graph, GraphRunContext, HistoryStep
+from pydantic_graph.nodes import End
 
 from . import (
+    _agent_graph,
     _result,
     _system_prompt,
     _utils,
@@ -22,6 +25,7 @@ from . import (
     result,
     usage as _usage,
 )
+from ._agent_graph import EndStrategy, capture_run_messages  # imported for re-export
 from .result import ResultDataT
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -29,7 +33,6 @@ from .tools import (
     DocstringFormat,
     RunContext,
     Tool,
-    ToolDefinition,
     ToolFuncContext,
     ToolFuncEither,
     ToolFuncPlain,
@@ -52,14 +55,7 @@ else:
     logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
-"""An invariant TypeVar."""
 NoneType = type(None)
-EndStrategy = Literal['early', 'exhaustive']
-"""The strategy for handling multiple tool calls when a final result is found.
-
-- `'early'`: Stop processing other tool calls once a final result is found
-- `'exhaustive'`: Process all tool calls even after finding a final result
-"""
 RunResultDataT = TypeVar('RunResultDataT')
 """Type variable for the result data of a run where `result_type` was customized on the run call."""
 
@@ -104,18 +100,24 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     Note, if `model_settings` is provided by `run`, `run_sync`, or `run_stream`, those settings will
     be merged with this value, with the runtime argument taking priority.
     """
+
+    result_type: type[ResultDataT] = dataclasses.field(repr=False)
+    """
+    The type of the result data, used to validate the result data, defaults to `str`.
+    """
+
+    _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _result_tool_name: str = dataclasses.field(repr=False)
     _result_tool_description: str | None = dataclasses.field(repr=False)
     _result_schema: _result.ResultSchema[ResultDataT] | None = dataclasses.field(repr=False)
     _result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]] = dataclasses.field(repr=False)
     _system_prompts: tuple[str, ...] = dataclasses.field(repr=False)
-    _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
-    _default_retries: int = dataclasses.field(repr=False)
     _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(repr=False)
     _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(
         repr=False
     )
-    _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
+    _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
+    _default_retries: int = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
     _override_deps: _utils.Option[AgentDepsT] = dataclasses.field(default=None, repr=False)
     _override_model: _utils.Option[models.Model] = dataclasses.field(default=None, repr=False)
@@ -174,25 +176,30 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self.end_strategy = end_strategy
         self.name = name
         self.model_settings = model_settings
+        self.result_type = result_type
+
+        self._deps_type = deps_type
+
         self._result_tool_name = result_tool_name
         self._result_tool_description = result_tool_description
-        self._result_schema = _result.ResultSchema[result_type].build(
+        self._result_schema: _result.ResultSchema[ResultDataT] | None = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
+        self._result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]] = []
 
         self._system_prompts = (system_prompt,) if isinstance(system_prompt, str) else tuple(system_prompt)
-        self._function_tools = {}
+        self._system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = []
+        self._system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = {}
+
+        self._function_tools: dict[str, Tool[AgentDepsT]] = {}
+
         self._default_retries = retries
+        self._max_result_retries = result_retries if result_retries is not None else retries
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
             else:
                 self._register_tool(Tool(tool))
-        self._deps_type = deps_type
-        self._system_prompt_functions = []
-        self._system_prompt_dynamic_functions = {}
-        self._max_result_retries = result_retries if result_retries is not None else retries
-        self._result_validators = []
 
     @overload
     async def run(
@@ -272,66 +279,80 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
-        result_schema = self._prepare_result_schema(result_type)
+        result_schema: _result.ResultSchema[RunResultDataT] | None = self._prepare_result_schema(result_type)
+
+        # Build the graph
+        graph = self._build_graph(result_type)
+
+        # Build the initial state
+        state = _agent_graph.GraphAgentState(
+            message_history=message_history[:] if message_history else [],
+            usage=usage or _usage.Usage(),
+            retries=0,
+            run_step=0,
+        )
+
+        # We consider it a user error if a user tries to restrict the result type while having a result validator that
+        # may change the result type from the restricted type to something else. Therefore, we consider the following
+        # typecast reasonable, even though it is possible to violate it with otherwise-type-checked code.
+        result_validators = cast(list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators)
+
+        # TODO: Instead of this, copy the function tools to ensure they don't share current_retry state between agent
+        #  runs. Requires some changes to `Tool` to make them copyable though.
+        for v in self._function_tools.values():
+            v.current_retry = 0
+
+        model_settings = merge_model_settings(self.model_settings, model_settings)
+        usage_limits = usage_limits or _usage.UsageLimits()
 
         with _logfire.span(
             '{agent_name} run {prompt=}',
             prompt=user_prompt,
             agent=self,
-            model_name=model_used.name(),
+            model_name=model_used.name() if model_used else 'no-model',
             agent_name=self.name or 'agent',
         ) as run_span:
-            run_context = RunContext(deps, model_used, usage or _usage.Usage(), user_prompt)
-            messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            run_context.messages = messages
-
-            for tool in self._function_tools.values():
-                tool.current_retry = 0
-
-            model_settings = merge_model_settings(self.model_settings, model_settings)
-            usage_limits = usage_limits or _usage.UsageLimits()
-
-            while True:
-                usage_limits.check_before_request(run_context.usage)
-
-                run_context.run_step += 1
-                with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
-                    agent_model = await self._prepare_model(run_context, result_schema)
-
-                with _logfire.span('model request', run_step=run_context.run_step) as model_req_span:
-                    model_response, request_usage = await agent_model.request(messages, model_settings)
-                    model_req_span.set_attribute('response', model_response)
-                    model_req_span.set_attribute('usage', request_usage)
+            # Build the deps object for the graph
+            graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunResultDataT](
+                user_deps=deps,
+                prompt=user_prompt,
+                new_message_index=new_message_index,
+                model=model_used,
+                model_settings=model_settings,
+                usage_limits=usage_limits,
+                max_result_retries=self._max_result_retries,
+                end_strategy=self.end_strategy,
+                result_schema=result_schema,
+                result_tools=self._result_schema.tool_defs() if self._result_schema else [],
+                result_validators=result_validators,
+                function_tools=self._function_tools,
+                run_span=run_span,
+            )
 
-                messages.append(model_response)
-                run_context.usage.incr(request_usage, requests=1)
-                usage_limits.check_tokens(run_context.usage)
+            start_node = _agent_graph.UserPromptNode[AgentDepsT](
+                user_prompt=user_prompt,
+                system_prompts=self._system_prompts,
+                system_prompt_functions=self._system_prompt_functions,
+                system_prompt_dynamic_functions=self._system_prompt_dynamic_functions,
+            )
 
-                with _logfire.span('handle model response', run_step=run_context.run_step) as handle_span:
-                    final_result, tool_responses = await self._handle_model_response(
-                        model_response, run_context, result_schema
-                    )
+            # Actually run
+            end_result, _ = await graph.run(
+                start_node,
+                state=state,
+                deps=graph_deps,
+                infer_name=False,
+            )
 
-                    if tool_responses:
-                        # Add parts to the conversation as a new message
-                        messages.append(_messages.ModelRequest(tool_responses))
-
-                    # Check if we got a final result
-                    if final_result is not None:
-                        result_data = final_result.data
-                        result_tool_name = final_result.tool_name
-                        run_span.set_attribute('all_messages', messages)
-                        run_span.set_attribute('usage', run_context.usage)
-                        handle_span.set_attribute('result', result_data)
-                        handle_span.message = 'handle model response -> final result'
-                        return result.RunResult(
-                            messages, new_message_index, result_data, result_tool_name, run_context.usage
-                        )
-                    else:
-                        # continue the conversation
-                        handle_span.set_attribute('tool_responses', tool_responses)
-                        tool_responses_str = ' '.join(r.part_kind for r in tool_responses)
-                        handle_span.message = f'handle model response -> {tool_responses_str}'
+        # Build final run result
+        # We don't do any advanced checking if the data is actually from a final result or not
+        return result.RunResult(
+            state.message_history,
+            new_message_index,
+            end_result.data,
+            end_result.tool_name,
+            state.usage,
+        )
 
     @overload
     def run_sync(
@@ -503,7 +524,31 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
-        result_schema = self._prepare_result_schema(result_type)
+        result_schema: _result.ResultSchema[RunResultDataT] | None = self._prepare_result_schema(result_type)
+
+        # Build the graph
+        graph = self._build_stream_graph(result_type)
+
+        # Build the initial state
+        graph_state = _agent_graph.GraphAgentState(
+            message_history=message_history[:] if message_history else [],
+            usage=usage or _usage.Usage(),
+            retries=0,
+            run_step=0,
+        )
+
+        # We consider it a user error if a user tries to restrict the result type while having a result validator that
+        # may change the result type from the restricted type to something else. Therefore, we consider the following
+        # typecast reasonable, even though it is possible to violate it with otherwise-type-checked code.
+        result_validators = cast(list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators)
+
+        # TODO: Instead of this, copy the function tools to ensure they don't share current_retry state between agent
+        #  runs. Requires some changes to `Tool` to make them copyable though.
+        for v in self._function_tools.values():
+            v.current_retry = 0
+
+        model_settings = merge_model_settings(self.model_settings, model_settings)
+        usage_limits = usage_limits or _usage.UsageLimits()
 
         with _logfire.span(
             '{agent_name} run stream {prompt=}',
@@ -512,97 +557,53 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
-            run_context = RunContext(deps, model_used, usage or _usage.Usage(), user_prompt)
-            messages = await self._prepare_messages(user_prompt, message_history, run_context)
-            run_context.messages = messages
-
-            for tool in self._function_tools.values():
-                tool.current_retry = 0
+            # Build the deps object for the graph
+            graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunResultDataT](
+                user_deps=deps,
+                prompt=user_prompt,
+                new_message_index=new_message_index,
+                model=model_used,
+                model_settings=model_settings,
+                usage_limits=usage_limits,
+                max_result_retries=self._max_result_retries,
+                end_strategy=self.end_strategy,
+                result_schema=result_schema,
+                result_tools=self._result_schema.tool_defs() if self._result_schema else [],
+                result_validators=result_validators,
+                function_tools=self._function_tools,
+                run_span=run_span,
+            )
 
-            model_settings = merge_model_settings(self.model_settings, model_settings)
-            usage_limits = usage_limits or _usage.UsageLimits()
+            start_node = _agent_graph.StreamUserPromptNode[AgentDepsT](
+                user_prompt=user_prompt,
+                system_prompts=self._system_prompts,
+                system_prompt_functions=self._system_prompt_functions,
+                system_prompt_dynamic_functions=self._system_prompt_dynamic_functions,
+            )
 
+            # Actually run
+            node = start_node
+            history: list[HistoryStep[_agent_graph.GraphAgentState, RunResultDataT]] = []
             while True:
-                run_context.run_step += 1
-                usage_limits.check_before_request(run_context.usage)
-
-                with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
-                    agent_model = await self._prepare_model(run_context, result_schema)
-
-                with _logfire.span('model request {run_step=}', run_step=run_context.run_step) as model_req_span:
-                    async with agent_model.request_stream(messages, model_settings) as model_response:
-                        run_context.usage.requests += 1
-                        model_req_span.set_attribute('response_type', model_response.__class__.__name__)
-                        # We want to end the "model request" span here, but we can't exit the context manager
-                        # in the traditional way
-                        model_req_span.__exit__(None, None, None)
-
-                        with _logfire.span('handle model response') as handle_span:
-                            maybe_final_result = await self._handle_streamed_response(
-                                model_response, run_context, result_schema
-                            )
-
-                            # Check if we got a final result
-                            if isinstance(maybe_final_result, _MarkFinalResult):
-                                result_stream = maybe_final_result.data
-                                result_tool_name = maybe_final_result.tool_name
-                                handle_span.message = 'handle model response -> final result'
-
-                                async def on_complete():
-                                    """Called when the stream has completed.
-
-                                    The model response will have been added to messages by now
-                                    by `StreamedRunResult._marked_completed`.
-                                    """
-                                    last_message = messages[-1]
-                                    assert isinstance(last_message, _messages.ModelResponse)
-                                    tool_calls = [
-                                        part for part in last_message.parts if isinstance(part, _messages.ToolCallPart)
-                                    ]
-                                    parts = await self._process_function_tools(
-                                        tool_calls, result_tool_name, run_context, result_schema
-                                    )
-                                    if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
-                                        self._incr_result_retry(run_context)
-                                    if parts:
-                                        messages.append(_messages.ModelRequest(parts))
-                                    run_span.set_attribute('all_messages', messages)
-
-                                # The following is not guaranteed to be true, but we consider it a user error if
-                                # there are result validators that might convert the result data from an overridden
-                                # `result_type` to a type that is not valid as such.
-                                result_validators = cast(
-                                    list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators
-                                )
-
-                                yield result.StreamedRunResult(
-                                    messages,
-                                    new_message_index,
-                                    usage_limits,
-                                    result_stream,
-                                    result_schema,
-                                    run_context,
-                                    result_validators,
-                                    result_tool_name,
-                                    on_complete,
-                                )
-                                return
-                            else:
-                                # continue the conversation
-                                model_response_msg, tool_responses = maybe_final_result
-                                # if we got a model response add that to messages
-                                messages.append(model_response_msg)
-                                if tool_responses:
-                                    # if we got one or more tool response parts, add a model request message
-                                    messages.append(_messages.ModelRequest(tool_responses))
-
-                                handle_span.set_attribute('tool_responses', tool_responses)
-                                tool_responses_str = ' '.join(r.part_kind for r in tool_responses)
-                                handle_span.message = f'handle model response -> {tool_responses_str}'
-                                # the model_response should have been fully streamed by now, we can add its usage
-                                model_response_usage = model_response.usage()
-                                run_context.usage.incr(model_response_usage)
-                                usage_limits.check_tokens(run_context.usage)
+                if isinstance(node, _agent_graph.StreamModelRequestNode):
+                    node = cast(
+                        _agent_graph.StreamModelRequestNode[
+                            AgentDepsT, result.StreamedRunResult[AgentDepsT, RunResultDataT]
+                        ],
+                        node,
+                    )
+                    async with node.run_to_result(GraphRunContext(graph_state, graph_deps)) as r:
+                        if isinstance(r, End):
+                            yield r.data
+                            break
+                assert not isinstance(node, End)  # the previous line should be hit first
+                node = await graph.next(
+                    node,
+                    history,
+                    state=graph_state,
+                    deps=graph_deps,
+                    infer_name=False,
+                )
 
     @contextmanager
     def override(
@@ -718,7 +719,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             return decorator
         else:
             assert not dynamic, "dynamic can't be True in this case"
-            self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func, dynamic=dynamic))
+            self._system_prompt_functions.append(_system_prompt.SystemPromptRunner[AgentDepsT](func, dynamic=dynamic))
             return func
 
     @overload
@@ -998,335 +999,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         return model_
 
-    async def _prepare_model(
-        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
-    ) -> models.AgentModel:
-        """Build tools and create an agent model."""
-        function_tools: list[ToolDefinition] = []
-
-        async def add_tool(tool: Tool[AgentDepsT]) -> None:
-            ctx = run_context.replace_with(retry=tool.current_retry, tool_name=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 run_context.model.agent_model(
-            function_tools=function_tools,
-            allow_text_result=self._allow_text_result(result_schema),
-            result_tools=result_schema.tool_defs() if result_schema is not None else [],
-        )
-
-    async def _reevaluate_dynamic_prompts(
-        self, messages: list[_messages.ModelMessage], run_context: RunContext[AgentDepsT]
-    ) -> None:
-        """Reevaluate any `SystemPromptPart` with dynamic_ref in the provided messages by running the associated runner function."""
-        # Only proceed if there's at least one dynamic runner.
-        if self._system_prompt_dynamic_functions:
-            for msg in messages:
-                if isinstance(msg, _messages.ModelRequest):
-                    for i, part in enumerate(msg.parts):
-                        if isinstance(part, _messages.SystemPromptPart) and part.dynamic_ref:
-                            # Look up the runner by its ref
-                            if runner := self._system_prompt_dynamic_functions.get(part.dynamic_ref):
-                                updated_part_content = await runner.run(run_context)
-                                msg.parts[i] = _messages.SystemPromptPart(
-                                    updated_part_content, dynamic_ref=part.dynamic_ref
-                                )
-
-    def _prepare_result_schema(
-        self, result_type: type[RunResultDataT] | None
-    ) -> _result.ResultSchema[RunResultDataT] | None:
-        if result_type is not None:
-            if self._result_validators:
-                raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
-            return _result.ResultSchema[result_type].build(
-                result_type, self._result_tool_name, self._result_tool_description
-            )
-        else:
-            return self._result_schema  # pyright: ignore[reportReturnType]
-
-    async def _prepare_messages(
-        self,
-        user_prompt: str,
-        message_history: list[_messages.ModelMessage] | None,
-        run_context: RunContext[AgentDepsT],
-    ) -> list[_messages.ModelMessage]:
-        try:
-            ctx_messages = get_captured_run_messages()
-        except LookupError:
-            messages: list[_messages.ModelMessage] = []
-        else:
-            if ctx_messages.used:
-                messages = []
-            else:
-                messages = ctx_messages.messages
-                ctx_messages.used = True
-
-        if message_history:
-            # Shallow copy messages
-            messages.extend(message_history)
-            # Reevaluate any dynamic system prompt parts
-            await self._reevaluate_dynamic_prompts(messages, run_context)
-            messages.append(_messages.ModelRequest([_messages.UserPromptPart(user_prompt)]))
-        else:
-            parts = await self._sys_parts(run_context)
-            parts.append(_messages.UserPromptPart(user_prompt))
-            messages.append(_messages.ModelRequest(parts))
-
-        return messages
-
-    async def _handle_model_response(
-        self,
-        model_response: _messages.ModelResponse,
-        run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultDataT] | None,
-    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
-        """Process a non-streamed response from the model.
-
-        Returns:
-            A tuple of `(final_result, request parts)`. If `final_result` is not `None`, the conversation should end.
-        """
-        texts: list[str] = []
-        tool_calls: list[_messages.ToolCallPart] = []
-        for part in model_response.parts:
-            if isinstance(part, _messages.TextPart):
-                # ignore empty content for text parts, see #437
-                if part.content:
-                    texts.append(part.content)
-            else:
-                tool_calls.append(part)
-
-        # At the moment, we prioritize at least executing tool calls if they are present.
-        # In the future, we'd consider making this configurable at the agent or run level.
-        # This accounts for cases like anthropic returns that might contain a text response
-        # and a tool call response, where the text response just indicates the tool call will happen.
-        if tool_calls:
-            return await self._handle_structured_response(tool_calls, run_context, result_schema)
-        elif texts:
-            text = '\n\n'.join(texts)
-            return await self._handle_text_response(text, run_context, result_schema)
-        else:
-            raise exceptions.UnexpectedModelBehavior('Received empty model response')
-
-    async def _handle_text_response(
-        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
-    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
-        """Handle a plain text response from the model for non-streaming responses."""
-        if self._allow_text_result(result_schema):
-            result_data_input = cast(RunResultDataT, text)
-            try:
-                result_data = await self._validate_result(result_data_input, run_context, None)
-            except _result.ToolRetryError as e:
-                self._incr_result_retry(run_context)
-                return None, [e.tool_retry]
-            else:
-                return _MarkFinalResult(result_data, None), []
-        else:
-            self._incr_result_retry(run_context)
-            response = _messages.RetryPromptPart(
-                content='Plain text responses are not permitted, please call one of the functions instead.',
-            )
-            return None, [response]
-
-    async def _handle_structured_response(
-        self,
-        tool_calls: list[_messages.ToolCallPart],
-        run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultDataT] | None,
-    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
-        """Handle a structured response containing tool calls from the model for non-streaming responses."""
-        assert tool_calls, 'Expected at least one tool call'
-
-        # first look for the result tool call
-        final_result: _MarkFinalResult[RunResultDataT] | None = None
-
-        parts: list[_messages.ModelRequestPart] = []
-        if result_schema is not None:
-            if match := result_schema.find_tool(tool_calls):
-                call, result_tool = match
-                try:
-                    result_data = result_tool.validate(call)
-                    result_data = await self._validate_result(result_data, run_context, call)
-                except _result.ToolRetryError as e:
-                    parts.append(e.tool_retry)
-                else:
-                    final_result = _MarkFinalResult(result_data, call.tool_name)
-
-        # Then build the other request parts based on end strategy
-        parts += await self._process_function_tools(
-            tool_calls, final_result and final_result.tool_name, run_context, result_schema
-        )
-
-        if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
-            self._incr_result_retry(run_context)
-
-        return final_result, parts
-
-    async def _process_function_tools(
-        self,
-        tool_calls: list[_messages.ToolCallPart],
-        result_tool_name: str | None,
-        run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultDataT] | None,
-    ) -> list[_messages.ModelRequestPart]:
-        """Process function (non-result) tool calls in parallel.
-
-        Also add stub return parts for any other tools that need it.
-        """
-        parts: list[_messages.ModelRequestPart] = []
-        tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
-
-        stub_function_tools = bool(result_tool_name) and self.end_strategy == 'early'
-
-        # we rely on the fact that if we found a result, it's the first result tool in the last
-        found_used_result_tool = False
-        for call in tool_calls:
-            if call.tool_name == result_tool_name and not found_used_result_tool:
-                found_used_result_tool = True
-                parts.append(
-                    _messages.ToolReturnPart(
-                        tool_name=call.tool_name,
-                        content='Final result processed.',
-                        tool_call_id=call.tool_call_id,
-                    )
-                )
-            elif tool := self._function_tools.get(call.tool_name):
-                if stub_function_tools:
-                    parts.append(
-                        _messages.ToolReturnPart(
-                            tool_name=call.tool_name,
-                            content='Tool not executed - a final result was already processed.',
-                            tool_call_id=call.tool_call_id,
-                        )
-                    )
-                else:
-                    tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name))
-            elif result_schema is not None and call.tool_name in result_schema.tools:
-                # if tool_name is in _result_schema, it means we found a result tool but an error occurred in
-                # validation, we don't add another part here
-                if result_tool_name is not None:
-                    parts.append(
-                        _messages.ToolReturnPart(
-                            tool_name=call.tool_name,
-                            content='Result tool not used - a final result was already processed.',
-                            tool_call_id=call.tool_call_id,
-                        )
-                    )
-            else:
-                parts.append(self._unknown_tool(call.tool_name, result_schema))
-
-        # Run all tool tasks in parallel
-        if tasks:
-            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
-                parts.extend(task_results)
-        return parts
-
-    async def _handle_streamed_response(
-        self,
-        streamed_response: models.StreamedResponse,
-        run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultDataT] | None,
-    ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
-        """Process a streamed response from the model.
-
-        Returns:
-            Either a final result or a tuple of the model response and the tool responses for the next request.
-            If a final result is returned, the conversation should end.
-        """
-        received_text = False
-
-        async for maybe_part_event in streamed_response:
-            if isinstance(maybe_part_event, _messages.PartStartEvent):
-                new_part = maybe_part_event.part
-                if isinstance(new_part, _messages.TextPart):
-                    received_text = True
-                    if self._allow_text_result(result_schema):
-                        return _MarkFinalResult(streamed_response, None)
-                elif isinstance(new_part, _messages.ToolCallPart):
-                    if result_schema is not None and (match := result_schema.find_tool([new_part])):
-                        call, _ = match
-                        return _MarkFinalResult(streamed_response, call.tool_name)
-                else:
-                    assert_never(new_part)
-
-        tasks: list[asyncio.Task[_messages.ModelRequestPart]] = []
-        parts: list[_messages.ModelRequestPart] = []
-        model_response = streamed_response.get()
-        if not model_response.parts:
-            raise exceptions.UnexpectedModelBehavior('Received empty model response')
-        for p in model_response.parts:
-            if isinstance(p, _messages.ToolCallPart):
-                if tool := self._function_tools.get(p.tool_name):
-                    tasks.append(asyncio.create_task(tool.run(p, run_context), name=p.tool_name))
-                else:
-                    parts.append(self._unknown_tool(p.tool_name, result_schema))
-
-        if received_text and not tasks and not parts:
-            # Can only get here if self._allow_text_result returns `False` for the provided result_schema
-            self._incr_result_retry(run_context)
-            model_response = _messages.RetryPromptPart(
-                content='Plain text responses are not permitted, please call one of the functions instead.',
-            )
-            return streamed_response.get(), [model_response]
-
-        with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-            task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
-            parts.extend(task_results)
-
-        if any(isinstance(part, _messages.RetryPromptPart) for part in parts):
-            self._incr_result_retry(run_context)
-
-        return model_response, parts
-
-    async def _validate_result(
-        self,
-        result_data: RunResultDataT,
-        run_context: RunContext[AgentDepsT],
-        tool_call: _messages.ToolCallPart | None,
-    ) -> RunResultDataT:
-        if self._result_validators:
-            agent_result_data = cast(ResultDataT, result_data)
-            for validator in self._result_validators:
-                agent_result_data = await validator.validate(agent_result_data, tool_call, run_context)
-            return cast(RunResultDataT, agent_result_data)
-        else:
-            return result_data
-
-    def _incr_result_retry(self, run_context: RunContext[AgentDepsT]) -> None:
-        run_context.retry += 1
-        if run_context.retry > self._max_result_retries:
-            raise exceptions.UnexpectedModelBehavior(
-                f'Exceeded maximum retries ({self._max_result_retries}) for result validation'
-            )
-
-    async def _sys_parts(self, run_context: RunContext[AgentDepsT]) -> list[_messages.ModelRequestPart]:
-        """Build the initial messages for the conversation."""
-        messages: list[_messages.ModelRequestPart] = [_messages.SystemPromptPart(p) for p in self._system_prompts]
-        for sys_prompt_runner in self._system_prompt_functions:
-            prompt = await sys_prompt_runner.run(run_context)
-            if sys_prompt_runner.dynamic:
-                messages.append(_messages.SystemPromptPart(prompt, dynamic_ref=sys_prompt_runner.function.__qualname__))
-            else:
-                messages.append(_messages.SystemPromptPart(prompt))
-        return messages
-
-    def _unknown_tool(
-        self,
-        tool_name: str,
-        result_schema: _result.ResultSchema[RunResultDataT] | None,
-    ) -> _messages.RetryPromptPart:
-        names = list(self._function_tools.keys())
-        if result_schema:
-            names.extend(result_schema.tool_names())
-        if names:
-            msg = f'Available tools: {", ".join(names)}'
-        else:
-            msg = 'No tools available.'
-        return _messages.RetryPromptPart(content=f'Unknown tool name: {tool_name!r}. {msg}')
-
-    def _get_deps(self: Agent[T, Any], deps: T) -> T:
+    def _get_deps(self: Agent[T, ResultDataT], deps: T) -> T:
         """Get deps for a run.
 
         If we've overridden deps via `_override_deps`, use that, otherwise use the deps passed to the call.
@@ -1357,10 +1030,6 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                             self.name = name
                             return
 
-    @staticmethod
-    def _allow_text_result(result_schema: _result.ResultSchema[RunResultDataT] | None) -> bool:
-        return result_schema is None or result_schema.allow_text_result
-
     @property
     @deprecated(
         'The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.', category=None
@@ -1368,65 +1037,24 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     def last_run_messages(self) -> list[_messages.ModelMessage]:
         raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
 
+    def _build_graph(
+        self, result_type: type[RunResultDataT] | None
+    ) -> Graph[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any], Any]:
+        return _agent_graph.build_agent_graph(self.name, self._deps_type, result_type or self.result_type)
 
-@dataclasses.dataclass
-class _RunMessages:
-    messages: list[_messages.ModelMessage]
-    used: bool = False
-
-
-_messages_ctx_var: ContextVar[_RunMessages] = ContextVar('var')
-
-
-@contextmanager
-def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
-    """Context manager to access the messages used in a [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] call.
-
-    Useful when a run may raise an exception, see [model errors](../agents.md#model-errors) for more information.
-
-    Examples:
-    ```python
-    from pydantic_ai import Agent, capture_run_messages
-
-    agent = Agent('test')
-
-    with capture_run_messages() as messages:
-        try:
-            result = agent.run_sync('foobar')
-        except Exception:
-            print(messages)
-            raise
-    ```
-
-    !!! note
-        If you call `run`, `run_sync`, or `run_stream` more than once within a single `capture_run_messages` context,
-        `messages` will represent the messages exchanged during the first call only.
-    """
-    try:
-        yield _messages_ctx_var.get().messages
-    except LookupError:
-        messages: list[_messages.ModelMessage] = []
-        token = _messages_ctx_var.set(_RunMessages(messages))
-        try:
-            yield messages
-        finally:
-            _messages_ctx_var.reset(token)
-
-
-def get_captured_run_messages() -> _RunMessages:
-    return _messages_ctx_var.get()
-
-
-@dataclasses.dataclass
-class _MarkFinalResult(Generic[ResultDataT]):
-    """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 `ResultDataT` directly.
-
-    It also avoids problems in the case where the result type is itself `None`, but is set.
-    """
+    def _build_stream_graph(
+        self, result_type: type[RunResultDataT] | None
+    ) -> Graph[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any], Any]:
+        return _agent_graph.build_agent_stream_graph(self.name, self._deps_type, result_type or self.result_type)
 
-    data: ResultDataT
-    """The final result data."""
-    tool_name: str | None
-    """Name of the final result tool, None if the result is a string."""
+    def _prepare_result_schema(
+        self, result_type: type[RunResultDataT] | None
+    ) -> _result.ResultSchema[RunResultDataT] | None:
+        if result_type is not None:
+            if self._result_validators:
+                raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
+            return _result.ResultSchema[result_type].build(
+                result_type, self._result_tool_name, self._result_tool_description
+            )
+        else:
+            return self._result_schema  # pyright: ignore[reportReturnType]