diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 5f356cfb094..1d87c683949 100644 --- a/python/packages/core/AGENTS.md +++ b/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. diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index cca47bd978f..aced39739bd 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/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,7 +678,8 @@ 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 @@ -685,6 +687,11 @@ def as_tool( 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: diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 096005930f8..3ee6053d59b 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/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,7 +4880,13 @@ 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"}) @@ -4770,6 +4894,7 @@ async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetRespo 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