Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/_agent_graph.py (Python)

Code before:
            )

    async def _handle_text_response(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        validation_ctx: Any | Callable[[RunContext[DepsT]], Any],
        text: str,
        text_processor: _output.BaseOutputProcessor[NodeRunEndT],
    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
        run_context = build_run_context(ctx)
        validation_context = build_validation_context(validation_ctx, run_context)

        result_data = await text_processor.process(text, run_context, validation_context)

        for validator in ctx.deps.output_validators:
            result_data = await validator.validate(result_data, run_context)
        return self._handle_final_result(ctx, result.FinalResult(result_data), [])

    async def _handle_image_response(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        image: _messages.BinaryImage,
    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
        result_data = cast(NodeRunEndT, image)
        return self._handle_final_result(ctx, result.FinalResult(result_data), [])

    def _handle_final_result(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        final_result: result.FinalResult[NodeRunEndT],
        tool_responses: list[_messages.ModelRequestPart],
    ) -> End[result.FinalResult[NodeRunEndT]]:
        messages = ctx.state.message_history

        # For backwards compatibility, append a new ModelRequest using the tool returns and retries
        if tool_responses:
            messages.append(_messages.ModelRequest(parts=tool_responses, run_id=ctx.state.run_id))

        return End(final_result)

    __repr__ = dataclasses_no_defaults_repr


@dataclasses.dataclass
class SetFinalResult(AgentNode[DepsT, NodeRunEndT]):
    """A node that immediately ends the graph run after a streaming response produced a final result."""

    final_result: result.FinalResult[NodeRunEndT]

    async def run(
        self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]

Review comment (refactor):
Ideally, output validators would have access to the validation context as well, as they could call `model_validate` themselves. They already have access to `RunContext`, so maybe it'd make sense to store the validation context on there? As the validation context callable itself needs `RunContext`, building run context could look like `ctx = <run context>; validation_ctx = callable(ctx); ctx = replace(ctx, validation_ctx = validation_ctx)`

I think that'd be a refactor worth exploring that could allow us to drop a lot of the new arguments.

Code after:
            output_parts=output_parts,
            output_final_result=output_final_result,
        ):
            yield event

        if output_final_result:
            final_result = output_final_result[0]
            self._next_node = self._handle_final_result(ctx, final_result, output_parts)
        else:
            instructions = await ctx.deps.get_instructions(run_context)
            self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
                _messages.ModelRequest(parts=output_parts, instructions=instructions)
            )

    async def _handle_text_response(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        text: str,
        text_processor: _output.BaseOutputProcessor[NodeRunEndT],
    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
        run_context = build_run_context(ctx)

        result_data = await text_processor.process(text, run_context=run_context)

        for validator in ctx.deps.output_validators:
            result_data = await validator.validate(result_data, run_context)
        return self._handle_final_result(ctx, result.FinalResult(result_data), [])

    async def _handle_image_response(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        image: _messages.BinaryImage,
    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
        result_data = cast(NodeRunEndT, image)
        return self._handle_final_result(ctx, result.FinalResult(result_data), [])

    def _handle_final_result(
        self,
        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
        final_result: result.FinalResult[NodeRunEndT],
        tool_responses: list[_messages.ModelRequestPart],
    ) -> End[result.FinalResult[NodeRunEndT]]:
        messages = ctx.state.message_history

        # For backwards compatibility, append a new ModelRequest using the tool returns and retries
        if tool_responses:
            messages.append(_messages.ModelRequest(parts=tool_responses, run_id=ctx.state.run_id))

        return End(final_result)

    __repr__ = dataclasses_no_defaults_repr