Unwrap Annotated in get_union_args util

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -864,7 +864,7 @@ async def _call_tools(
         output_deferred_calls[deferred_calls_by_index[k]].append(tool_calls[k])
 
 
-async def _call_tool(  # noqa: C901
+async def _call_tool(
     tool_manager: ToolManager[DepsT],
     tool_call: _messages.ToolCallPart,
     tool_call_result: DeferredToolResult | None,
@@ -893,10 +893,6 @@ async def _call_tool(  # noqa: C901
             tool_call_result.tool_name = tool_call.tool_name
             tool_call_result.tool_call_id = tool_call.tool_call_id
             raise ToolRetryError(tool_call_result)
-        elif isinstance(tool_call_result, _messages.ToolReturnPart):
-            tool_call_result.tool_name = tool_call.tool_name
-            tool_call_result.tool_call_id = tool_call.tool_call_id
-            return tool_call_result, None
         else:
             tool_result = tool_call_result
     except ToolRetryError as e:

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -459,14 +459,23 @@ def strip_markdown_fences(text: str) -> str:
     return text
 
 
+def _unwrap_annotated(tp: Any) -> Any:
+    origin = get_origin(tp)
+    while typing_objects.is_annotated(origin):
+        tp = tp.__origin__
+        origin = get_origin(tp)
+    return tp
+
+
 def get_union_args(tp: Any) -> tuple[Any, ...]:
     """Extract the arguments of a Union type if `tp` is a union, otherwise return an empty tuple."""
     if typing_objects.is_typealiastype(tp):
         tp = tp.__value__
 
+    tp = _unwrap_annotated(tp)
     origin = get_origin(tp)
     if is_union_origin(origin):
-        return get_args(tp)
+        return tuple(_unwrap_annotated(arg) for arg in get_args(tp))
     else:
         return ()
 

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -628,10 +628,12 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
                     approval = ToolDenied()
                 tool_call_results[tool_call_id] = approval
 
-            for tool_call_id, result in deferred_tool_results.calls.items():
-                if not isinstance(result, _utils.get_union_args(DeferredToolCallResult)):
-                    result = _messages.ToolReturn(result)
-                tool_call_results[tool_call_id] = result
+            if calls := deferred_tool_results.calls:
+                call_result_types = _utils.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
 
         graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunOutputDataT](
             user_deps=deps,