diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index 8b66cb4b16..304fdfea82 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/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, diff --git a/python/packages/main/tests/main/test_agents.py b/python/packages/main/tests/main/test_agents.py index 9737810b60..f41aeee9e1 100644 --- a/python/packages/main/tests/main/test_agents.py +++ b/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