Python: add configurable response extraction to agent.as_tool()

python/packages/core/AGENTS.md

@@ -54,6 +54,10 @@ agent_framework/
 - **`SupportsAgentRun`** - Protocol defining the agent interface
 - **`BaseAgent`** - Abstract base class for agents
 - **`Agent`** - Main agent class wrapping a chat client with tools, instructions, and middleware
+- **`BaseAgent.as_tool(return_mode=...)`** - Defaults to `"all"` (the complete response text). `"last"` selects
+  the last assistant message's text without falling back to earlier text when it is empty; a synchronous
+  `AgentResponse -> str` callable supports custom extraction. Selection happens after pending-input checks and
+  does not filter streaming callbacks or change session history.
 - **`RawAgent.open()` / `close()`** (inherited by `Agent`) - Explicitly enter the client's and configured MCP tools'
   async contexts, then release them along with lazily connected MCP tools. `async with agent` delegates to these
   methods; a partially failed `open()` closes resources already entered.

python/packages/core/agent_framework/_agents.py

@@ -667,6 +667,7 @@ def as_tool(
         approval_mode: Literal["always_require", "never_require"] = "never_require",
         stream_callback: Callable[[AgentResponseUpdate], Awaitable[None] | None] | None = None,
         propagate_session: bool = False,
+        return_mode: Literal["all", "last"] | Callable[[AgentResponse], str] = "all",
     ) -> FunctionTool:
         """Create a FunctionTool that wraps this agent.
 
@@ -677,14 +678,20 @@ def as_tool(
             arg_description: The description for the function argument.
                 If None, defaults to "Task for {tool_name}".
             approval_mode: Whether this delegated tool requires approval before execution.
-            stream_callback: Optional callback for streaming responses. If provided, uses run(..., stream=True).
+            stream_callback: Optional observer for streaming child-agent responses. Receives all updates regardless
+                of return_mode.
             propagate_session: If True, the parent agent's session is forwarded
                 to this sub-agent's ``run()`` call. Application-state changes
                 propagate back to the parent, while framework approval continuation
                 state and provider-owned service session state remain isolated.
                 Defaults to False. The sub-agent always receives an AgentSession
                 so session-backed middleware can run. When False, that session is
                 private to this invocation.
+            return_mode: How to extract the completed tool result. "all" (default) returns the response's
+                concatenated text. "last" returns the last assistant message's text, or an empty string if no
+                assistant message exists. An empty or non-text final assistant message does not fall back to
+                earlier text. A synchronous callable receives the complete AgentResponse and must return a str.
+                Pending human-input requests retain their existing behavior and are not passed to the callable.
 
         Returns:
             A FunctionTool that can be used as a tool by other agents.
@@ -718,6 +725,9 @@ def as_tool(
                 # Convert the agent to a tool (independent session)
                 research_tool = agent.as_tool()
 
+                # Return only the final assistant message to the coordinator
+                research_tool = agent.as_tool(return_mode="last")
+
                 # Convert the agent to a tool (shared session with parent)
                 research_tool = agent.as_tool(propagate_session=True)
 
@@ -728,6 +738,9 @@ def as_tool(
         if not isinstance(self, SupportsAgentRun):
             raise TypeError(f"Agent {self.__class__.__name__} must implement SupportsAgentRun to be used as a tool")
 
+        if not callable(return_mode) and return_mode not in ("all", "last"):
+            raise ValueError('return_mode must be "all", "last", or a synchronous callable returning str.')
+
         tool_name = name or _sanitize_agent_name(self.name)
         if tool_name is None:
             raise ValueError("Agent tool name cannot be None. Either provide a name parameter or set the agent's name.")
@@ -869,6 +882,16 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str:
                     raise ToolExecutionException(approval_error)
                 if other_input_requests:
                     raise UserInputRequiredException(contents=other_input_requests)
+                if callable(return_mode):
+                    result = return_mode(final_response)
+                    if not isinstance(result, str):
+                        raise TypeError("return_mode callable must return a str.")
+                    return result
+                if return_mode == "last":
+                    return next(
+                        (message.text for message in reversed(final_response.messages) if message.role == "assistant"),
+                        "",
+                    )
                 # TODO(Copilot): update once #4331 merges
                 return final_response.text
             finally:

python/packages/core/tests/core/test_agents.py

@@ -9,7 +9,7 @@
 import json
 import logging
 from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableSequence, Sequence
-from typing import Any, cast
+from typing import Any, Literal, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 from uuid import uuid4
 
@@ -1950,7 +1950,8 @@ def approve_weather(function_call: Content) -> bool:
     assert result[0].text == "The weather is cloudy."
 
 
-async def test_chat_agent_as_tool_fails_closed_for_unresolved_child_approval() -> None:
+@pytest.mark.parametrize("custom_return_mode", [False, True])
+async def test_chat_agent_as_tool_fails_closed_for_unresolved_child_approval(custom_return_mode: bool) -> None:
     """Test that an unresolved child approval cannot escape into the caller's tool registry."""
     executions = 0
 
@@ -1981,6 +1982,11 @@ def delete_weather_data() -> str:
         tools=[delete_weather_data],
         middleware=[ToolApprovalMiddleware(auto_approval_rules=[lambda function_call: False])],
     )
+    extracted_responses: list[AgentResponse] = []
+
+    def extract_result(response: AgentResponse) -> str:
+        extracted_responses.append(response)
+        return "Approval must not be hidden"
 
     with raises(
         ToolExecutionException,
@@ -1989,9 +1995,12 @@ def delete_weather_data() -> str:
             "ToolApprovalMiddleware.*Use a workflow for interactive, delayed, or durable approval"
         ),
     ):
-        await agent.as_tool().invoke(arguments={"task": "Delete weather data"})
+        await agent.as_tool(return_mode=extract_result if custom_return_mode else "all").invoke(
+            arguments={"task": "Delete weather data"}
+        )
 
     assert executions == 0
+    assert extracted_responses == []
 
 
 @pytest.mark.parametrize("stream", [False, True])
@@ -2195,6 +2204,114 @@ def guarded_write() -> str:
     assert not tool_approval_state.get("pending_approval_requests")
 
 
+@pytest.mark.parametrize("with_stream_callback", [False, True])
+@pytest.mark.parametrize(
+    ("return_mode", "expected"),
+    [
+        ("all", "Searching...Final answer"),
+        ("last", "Final answer"),
+        (
+            lambda response: f"{len(response.messages)} messages: {response.messages[0].text}",
+            "3 messages: Searching...",
+        ),
+    ],
+    ids=["all", "last", "callable"],
+)
+async def test_chat_agent_as_tool_return_mode(
+    return_mode: Literal["all", "last"] | Callable[[AgentResponse], str],
+    expected: str,
+    with_stream_callback: bool,
+) -> None:
+    client = MockBaseChatClient()
+    client.streaming_responses = [
+        [
+            ChatResponseUpdate(role="assistant", message_id="progress", contents=[Content.from_text("Searching...")]),
+            ChatResponseUpdate(role="assistant", message_id="answer", contents=[Content.from_text("Final answer")]),
+            ChatResponseUpdate(
+                role="tool", message_id="result", contents=[Content.from_function_result("call", result="Tool output")]
+            ),
+        ]
+    ]
+    agent = Agent(client=client, name="ChildAgent")
+    updates: list[AgentResponseUpdate] = []
+    agent_tool = agent.as_tool(
+        return_mode=return_mode,
+        stream_callback=updates.append if with_stream_callback else None,
+    )
+
+    result = await agent_tool.invoke(arguments={"task": "Find an answer"})
+
+    assert isinstance(result, list)
+    assert result[0].text == expected
+    assert set(agent_tool.parameters()["properties"]) == {"task"}
+    assert client.call_count == 1
+    if with_stream_callback:
+        assert "".join(update.text for update in updates) == "Searching...Final answer"
+        assert any(content.type == "function_result" for update in updates for content in update.contents)
+    else:
+        assert updates == []
+
+
+@pytest.mark.parametrize(
+    "response_updates",
+    [
+        [],
+        [ChatResponseUpdate(role="tool", contents=[Content.from_function_result("call", result="Tool output")])],
+        [
+            ChatResponseUpdate(role="assistant", message_id="progress", contents=[Content.from_text("Searching...")]),
+            ChatResponseUpdate(role="assistant", message_id="answer", contents=[Content.from_text("")]),
+        ],
+        [
+            ChatResponseUpdate(role="assistant", message_id="progress", contents=[Content.from_text("Searching...")]),
+            ChatResponseUpdate(
+                role="assistant", message_id="answer", contents=[Content.from_text_reasoning(text="Reasoning only")]
+            ),
+        ],
+    ],
+    ids=["empty-response", "no-assistant", "empty-final-assistant", "non-text-final-assistant"],
+)
+async def test_chat_agent_as_tool_return_last_without_final_text(response_updates: list[ChatResponseUpdate]) -> None:
+    client = MockBaseChatClient()
+    client.streaming_responses = [response_updates]
+    agent = Agent(client=client, name="ChildAgent")
+
+    result = await agent.as_tool(return_mode="last").invoke(arguments={"task": "Find an answer"})
+
+    assert isinstance(result, list)
+    assert result[0].text == ""
+
+
+@pytest.mark.parametrize("return_mode", ["unknown", "LAST", None, 0])
+def test_chat_agent_as_tool_rejects_invalid_return_mode(return_mode: Any) -> None:
+    client = MockBaseChatClient()
+    agent = Agent(client=client, name="ChildAgent")
+
+    with raises(ValueError, match="return_mode must be"):
+        agent.as_tool(return_mode=return_mode)
+
+    assert client.call_count == 0
+
+
+async def test_chat_agent_as_tool_rejects_non_string_extraction() -> None:
+    agent = Agent(client=MockBaseChatClient(), name="ChildAgent")
+
+    def extract_result(response: AgentResponse) -> Any:
+        return 42
+
+    with raises(TypeError, match="return_mode callable must return a str"):
+        await agent.as_tool(return_mode=extract_result).invoke(arguments={"task": "Find an answer"})
+
+
+async def test_chat_agent_as_tool_propagates_extraction_error() -> None:
+    agent = Agent(client=MockBaseChatClient(), name="ChildAgent")
+
+    def extract_result(response: AgentResponse) -> str:
+        raise ValueError("Extraction failed")
+
+    with raises(ValueError, match="Extraction failed"):
+        await agent.as_tool(return_mode=extract_result).invoke(arguments={"task": "Find an answer"})
+
+
 async def test_chat_agent_as_tool_with_stream_callback(
     client: SupportsChatGetResponse,
 ) -> None:
@@ -4750,7 +4867,8 @@ def _as_raw(resp: MagicMock) -> MagicMock:
 # region as_tool user_input_request propagation
 
 
-async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetResponse) -> None:
+@pytest.mark.parametrize("custom_return_mode", [False, True])
+async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetResponse, custom_return_mode: bool) -> None:
     """Test that as_tool raises when the wrapped sub-agent requests user input."""
     from agent_framework.exceptions import UserInputRequiredException
 
@@ -4762,14 +4880,21 @@ async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetRespo
     ]
 
     agent = Agent(client=client, name="OAuthAgent", description="Agent requiring consent")
-    agent_tool = agent.as_tool()
+    extracted_responses: list[AgentResponse] = []
+
+    def extract_result(response: AgentResponse) -> str:
+        extracted_responses.append(response)
+        return "Consent must not be hidden"
+
+    agent_tool = agent.as_tool(return_mode=extract_result if custom_return_mode else "all")
 
     with raises(UserInputRequiredException) as exc_info:
         await agent_tool.invoke(arguments={"task": "Do something"})
 
     assert len(exc_info.value.contents) == 1
     assert exc_info.value.contents[0].type == "oauth_consent_request"
     assert exc_info.value.contents[0].consent_link == "https://login.microsoftonline.com/consent"
+    assert extracted_responses == []
 
 
 # region Per-service-call history persistence scenario matrix