Python: Introduce as_tool() for BaseAgent (#684)

* introduce as_tool for BaseAgent

* fix types

* fix tests

* add async callback support

* address comments

* Update python/packages/main/agent_framework/_agents.py

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>

* address comments

---------

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>

python/packages/main/agent_framework/_agents.py

@@ -1,12 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import inspect
 import sys
-from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
+from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
 from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
 from uuid import uuid4
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr, create_model
 
 from ._clients import BaseChatClient, ChatClientProtocol
 from ._logging import get_logger
@@ -15,7 +16,7 @@ from ._memory import AggregateContextProvider, Context, ContextProvider
 from ._middleware import Middleware, use_agent_middleware
 from ._pydantic import AFBaseModel
 from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
-from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
+from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, AIFunction, ToolProtocol
 from ._types import (
     AgentRunResponse,
     AgentRunResponseUpdate,
@@ -173,6 +174,70 @@ class BaseAgent(AFBaseModel):
         await deserialize_thread_state(thread, serialized_thread, **kwargs)
         return thread
 
+    def as_tool(
+        self,
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        arg_name: str = "task",
+        arg_description: str | None = None,
+        stream_callback: Callable[[AgentRunResponseUpdate], None]
+        | Callable[[AgentRunResponseUpdate], Awaitable[None]]
+        | None = None,
+    ) -> AIFunction[BaseModel, str]:
+        """Create an AIFunction tool that wraps this agent.
+
+        Args:
+            name: The name for the tool. If None, uses the agent's name.
+            description: The description for the tool. If None, uses the agent's description or empty string.
+            arg_name: The name of the function argument (default: "task").
+            arg_description: The description for the function argument.
+                If None, defaults to "Input for {self.display_name}".
+            stream_callback: Optional callback for streaming responses. If provided, uses run_stream.
+
+        Returns:
+            An AIFunction that can be used as a tool by other agents.
+        """
+        # Verify that self implements AgentProtocol
+        if not isinstance(self, AgentProtocol):
+            raise TypeError(f"Agent {self.__class__.__name__} must implement AgentProtocol to be used as a tool")
+
+        tool_name = name or 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.")
+        tool_description = description or self.description or ""
+        argument_description = arg_description or f"Task for {tool_name}"
+
+        # Create dynamic input model with the specified argument name
+        field_info = Field(..., description=argument_description)
+        input_model = create_model(f"{name or self.name or 'agent'}_task", **{arg_name: (str, field_info)})  # type: ignore[call-overload]
+
+        # Check if callback is async once, outside the wrapper
+        is_async_callback = stream_callback is not None and inspect.iscoroutinefunction(stream_callback)
+
+        async def agent_wrapper(**kwargs: Any) -> str:
+            """Wrapper function that calls the agent."""
+            # Extract the input from kwargs using the specified arg_name
+            input_text = kwargs.get(arg_name, "")
+
+            if stream_callback is None:
+                # Use non-streaming mode
+                return (await self.run(input_text)).text
+
+            # Use streaming mode - accumulate updates and create final response
+            response_updates: list[AgentRunResponseUpdate] = []
+            async for update in self.run_stream(input_text):
+                response_updates.append(update)
+                if is_async_callback:
+                    await stream_callback(update)  # type: ignore[misc]
+                else:
+                    stream_callback(update)
+
+            # Create final text from accumulated updates
+            return AgentRunResponse.from_agent_run_response_updates(response_updates).text
+
+        return AIFunction(name=tool_name, description=tool_description, func=agent_wrapper, input_model=input_model)
+
     def _normalize_messages(
         self,
         messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,

python/packages/main/tests/main/test_agents.py

@@ -394,3 +394,135 @@ async def test_chat_agent_context_providers_with_thread_service_id(chat_client_b
     # messages_adding should be called with the service thread ID from response
     assert mock_provider.messages_adding_called
     assert mock_provider.messages_adding_thread_id == "service-thread-123"  # Updated thread ID from response
+
+
+# Tests for as_tool method
+async def test_chat_agent_as_tool_basic(chat_client: ChatClientProtocol) -> None:
+    """Test basic as_tool functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Test agent for as_tool")
+
+    tool = agent.as_tool()
+
+    assert tool.name == "TestAgent"
+    assert tool.description == "Test agent for as_tool"
+    assert hasattr(tool, "func")
+    assert hasattr(tool, "input_model")
+
+
+async def test_chat_agent_as_tool_custom_parameters(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with custom parameters."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Original description")
+
+    tool = agent.as_tool(
+        name="CustomTool",
+        description="Custom description",
+        arg_name="query",
+        arg_description="Custom input description",
+    )
+
+    assert tool.name == "CustomTool"
+    assert tool.description == "Custom description"
+
+    # Check that the input model has the custom field name
+    schema = tool.input_model.model_json_schema()
+    assert "query" in schema["properties"]
+    assert schema["properties"]["query"]["description"] == "Custom input description"
+
+
+async def test_chat_agent_as_tool_defaults(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with default parameters."""
+    agent = ChatAgent(
+        chat_client=chat_client,
+        name="TestAgent",
+        # No description provided
+    )
+
+    tool = agent.as_tool()
+
+    assert tool.name == "TestAgent"
+    assert tool.description == ""  # Should default to empty string
+
+    # Check default input field
+    schema = tool.input_model.model_json_schema()
+    assert "task" in schema["properties"]
+    assert "Task for TestAgent" in schema["properties"]["task"]["description"]
+
+
+async def test_chat_agent_as_tool_no_name(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool when agent has no name (should raise ValueError)."""
+    agent = ChatAgent(chat_client=chat_client)  # No name provided
+
+    # Should raise ValueError since agent has no name
+    with raises(ValueError, match="Agent tool name cannot be None"):
+        agent.as_tool()
+
+
+async def test_chat_agent_as_tool_function_execution(chat_client: ChatClientProtocol) -> None:
+    """Test that the generated AIFunction can be executed."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Test agent")
+
+    tool = agent.as_tool()
+
+    # Test function execution
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should return the agent's response text
+    assert isinstance(result, str)
+    assert result == "test response"  # From mock chat client
+
+
+async def test_chat_agent_as_tool_with_stream_callback(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with stream callback functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="StreamingAgent")
+
+    # Collect streaming updates
+    collected_updates: list[AgentRunResponseUpdate] = []
+
+    def stream_callback(update: AgentRunResponseUpdate) -> None:
+        collected_updates.append(update)
+
+    tool = agent.as_tool(stream_callback=stream_callback)
+
+    # Execute the tool
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should have collected streaming updates
+    assert len(collected_updates) > 0
+    assert isinstance(result, str)
+    # Result should be concatenation of all streaming updates
+    expected_text = "".join(update.text for update in collected_updates)
+    assert result == expected_text
+
+
+async def test_chat_agent_as_tool_with_custom_arg_name(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with custom argument name."""
+    agent = ChatAgent(chat_client=chat_client, name="CustomArgAgent")
+
+    tool = agent.as_tool(arg_name="prompt", arg_description="Custom prompt input")
+
+    # Test that the custom argument name works
+    result = await tool.invoke(arguments=tool.input_model(prompt="Test prompt"))
+    assert result == "test response"
+
+
+async def test_chat_agent_as_tool_with_async_stream_callback(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with async stream callback functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="AsyncStreamingAgent")
+
+    # Collect streaming updates using an async callback
+    collected_updates: list[AgentRunResponseUpdate] = []
+
+    async def async_stream_callback(update: AgentRunResponseUpdate) -> None:
+        collected_updates.append(update)
+
+    tool = agent.as_tool(stream_callback=async_stream_callback)
+
+    # Execute the tool
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should have collected streaming updates
+    assert len(collected_updates) > 0
+    assert isinstance(result, str)
+    # Result should be concatenation of all streaming updates
+    expected_text = "".join(update.text for update in collected_updates)
+    assert result == expected_text