Refactor duplicated code

Author: Danipulok

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -826,50 +826,27 @@ async def process_tool_calls(  # noqa: C901
 
     # First, we handle output tool calls
     for call in tool_calls_by_kind['output']:
-        if final_result:
-            if final_result.tool_call_id == call.tool_call_id:
-                part = _messages.ToolReturnPart(
-                    tool_name=call.tool_name,
-                    content='Final result processed.',
-                    tool_call_id=call.tool_call_id,
-                )
-                output_parts.append(part)
-            # With early strategy, execute only the first output tool
-            elif ctx.deps.end_strategy == 'early':
-                yield _messages.FunctionToolCallEvent(call)
-                part = _messages.ToolReturnPart(
-                    tool_name=call.tool_name,
-                    content='Output tool not used - a final result was already processed.',
-                    tool_call_id=call.tool_call_id,
-                )
-                yield _messages.FunctionToolResultEvent(part)
-                output_parts.append(part)
-            # With exhaustive strategy, execute all output tools
-            elif ctx.deps.end_strategy == 'exhaustive':
-                try:
-                    await tool_manager.handle_call(call)
-                except exceptions.UnexpectedModelBehavior as e:
-                    ctx.state.increment_retries(
-                        ctx.deps.max_result_retries, error=e, model_settings=ctx.deps.model_settings
-                    )
-                    raise e  # pragma: lax no cover
-                except ToolRetryError as e:
-                    ctx.state.increment_retries(
-                        ctx.deps.max_result_retries, error=e, model_settings=ctx.deps.model_settings
-                    )
-                    yield _messages.FunctionToolCallEvent(call)
-                    output_parts.append(e.tool_retry)
-                    yield _messages.FunctionToolResultEvent(e.tool_retry)
-                else:
-                    part = _messages.ToolReturnPart(
-                        tool_name=call.tool_name,
-                        content='Final result processed.',
-                        tool_call_id=call.tool_call_id,
-                    )
-                    output_parts.append(part)
-            else:
-                assert_never(ctx.deps.end_strategy)
-        else:
+        # In case we got two tool calls with the same ID
+        if final_result and final_result.tool_call_id == call.tool_call_id:
+            part = _messages.ToolReturnPart(
+                tool_name=call.tool_name,
+                content='Final result processed.',
+                tool_call_id=call.tool_call_id,
+            )
+            output_parts.append(part)
+        # Early strategy is chosen and final result is already set
+        elif ctx.deps.end_strategy == 'early' and final_result:
+            yield _messages.FunctionToolCallEvent(call)
+            part = _messages.ToolReturnPart(
+                tool_name=call.tool_name,
+                content='Output tool not used - a final result was already processed.',
+                tool_call_id=call.tool_call_id,
+            )
+            yield _messages.FunctionToolResultEvent(part)
+            output_parts.append(part)
+        # Early strategy is chosen and final result is not yet set
+        # Or exhaustive strategy is chosen
+        elif (ctx.deps.end_strategy == 'early' and not final_result) or ctx.deps.end_strategy == 'exhaustive':
             try:
                 result_data = await tool_manager.handle_call(call)
             except exceptions.UnexpectedModelBehavior as e:
@@ -891,7 +868,13 @@ async def process_tool_calls(  # noqa: C901
                     tool_call_id=call.tool_call_id,
                 )
                 output_parts.append(part)
-                final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
+
+                # Even in exhaustive mode, use the first output tool's result as the final result
+                if not final_result:
+                    final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
+        # Unknown strategy or invalid state
+        else:
+            assert_never(ctx.deps.end_strategy)
 
     # Then, we handle function tool calls
     calls_to_run: list[_messages.ToolCallPart] = []