wehub-resource-sync
722 lines
27 KiB
Python
722 lines
27 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
import asyncio
|
|
import sys
|
|
from collections.abc import AsyncIterable, Awaitable, Callable
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from semantic_kernel.agents.agent import Agent, AgentResponseItem, AgentThread
|
|
from semantic_kernel.agents.orchestration.handoffs import (
|
|
HANDOFF_PLUGIN_NAME,
|
|
HandoffAgentActor,
|
|
HandoffOrchestration,
|
|
OrchestrationHandoffs,
|
|
)
|
|
from semantic_kernel.agents.orchestration.orchestration_base import DefaultTypeAlias
|
|
from semantic_kernel.agents.runtime.in_process.in_process_runtime import InProcessRuntime
|
|
from semantic_kernel.contents.chat_history import ChatHistory
|
|
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
|
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
|
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
|
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
|
from semantic_kernel.contents.utils.author_role import AuthorRole
|
|
from semantic_kernel.kernel import Kernel
|
|
from tests.unit.agents.orchestration.conftest import MockAgent, MockAgentWithException, MockRuntime
|
|
|
|
if sys.version_info >= (3, 12):
|
|
from typing import override # pragma: no cover
|
|
else:
|
|
from typing_extensions import override # pragma: no cover
|
|
|
|
|
|
class MockAgentWithHandoffFunctionCall(Agent):
|
|
"""A mock agent with handoff function call for testing purposes."""
|
|
|
|
target_agent: Agent
|
|
|
|
def __init__(self, target_agent: Agent):
|
|
super().__init__(target_agent=target_agent)
|
|
|
|
@override
|
|
async def get_response(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
kernel: Kernel | None = None,
|
|
**kwargs,
|
|
) -> AgentResponseItem[ChatMessageContent]:
|
|
pass
|
|
|
|
@override
|
|
async def invoke(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
on_intermediate_message: Callable[[ChatMessageContent], Awaitable[None]] | None = None,
|
|
**kwargs,
|
|
) -> AgentResponseItem[ChatMessageContent]:
|
|
pass
|
|
|
|
@override
|
|
async def invoke_stream(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
on_intermediate_message: Callable[[ChatMessageContent], Awaitable[None]] | None = None,
|
|
kernel: Kernel | None = None,
|
|
**kwargs,
|
|
) -> AsyncIterable[AgentResponseItem[StreamingChatMessageContent]]:
|
|
"""Simulate streaming response from the agent."""
|
|
function_call = FunctionCallContent(
|
|
function_name=f"transfer_to_{self.target_agent.name}",
|
|
plugin_name=HANDOFF_PLUGIN_NAME,
|
|
call_id="test_call_id",
|
|
id="test_id",
|
|
)
|
|
|
|
# Simulate some processing time
|
|
await asyncio.sleep(0.1)
|
|
await kernel.invoke_function_call(
|
|
function_call=function_call,
|
|
chat_history=ChatHistory(),
|
|
)
|
|
|
|
# Do not yield any messages, as the agent doesn't yield any tool related messages from the streaming API.
|
|
# Nevertheless, the method needs have a `yield` code path to satisfy the AsyncIterable interface.
|
|
if False:
|
|
yield
|
|
|
|
# Simulate on_intermediate_message callback
|
|
await on_intermediate_message(
|
|
StreamingChatMessageContent(
|
|
role=AuthorRole.ASSISTANT,
|
|
name=self.name,
|
|
choice_index=0,
|
|
items=[function_call],
|
|
)
|
|
)
|
|
await on_intermediate_message(
|
|
StreamingChatMessageContent(
|
|
role=AuthorRole.ASSISTANT,
|
|
name=self.name,
|
|
choice_index=0,
|
|
items=[
|
|
FunctionResultContent(
|
|
function_name=function_call.function_name,
|
|
plugin_name=function_call.plugin_name,
|
|
call_id=function_call.call_id,
|
|
id=function_call.id,
|
|
result=None,
|
|
)
|
|
],
|
|
)
|
|
)
|
|
|
|
|
|
class MockAgentWithCompleteTaskFunctionCall(Agent):
|
|
"""A mock agent with complete_task function call for testing purposes."""
|
|
|
|
@override
|
|
async def get_response(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
kernel: Kernel | None = None,
|
|
**kwargs,
|
|
) -> AgentResponseItem[ChatMessageContent]:
|
|
pass
|
|
|
|
@override
|
|
async def invoke(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
on_intermediate_message: Callable[[ChatMessageContent], Awaitable[None]] | None = None,
|
|
**kwargs,
|
|
) -> AgentResponseItem[ChatMessageContent]:
|
|
pass
|
|
|
|
@override
|
|
async def invoke_stream(
|
|
self,
|
|
messages: str | ChatMessageContent | list[str | ChatMessageContent] | None = None,
|
|
*,
|
|
thread: AgentThread | None = None,
|
|
on_intermediate_message: Callable[[ChatMessageContent], Awaitable[None]] | None = None,
|
|
kernel: Kernel | None = None,
|
|
**kwargs,
|
|
) -> AsyncIterable[AgentResponseItem[StreamingChatMessageContent]]:
|
|
"""Simulate streaming response from the agent."""
|
|
# Simulate some processing time
|
|
await asyncio.sleep(0.1)
|
|
await kernel.invoke_function_call(
|
|
function_call=FunctionCallContent(
|
|
function_name="complete_task",
|
|
plugin_name=HANDOFF_PLUGIN_NAME,
|
|
call_id="test_call_id",
|
|
id="test_id",
|
|
arguments={"task_summary": "test_summary"},
|
|
),
|
|
chat_history=ChatHistory(),
|
|
)
|
|
|
|
# Do not yield any messages, as the agent doesn't yield any tool related messages from the streaming API.
|
|
# Nevertheless, the method needs have a `yield` code path to satisfy the AsyncIterable interface.
|
|
if False:
|
|
yield
|
|
|
|
|
|
# region HandoffOrchestration
|
|
|
|
|
|
def test_init_without_handoffs():
|
|
"""Test the initialization of HandoffOrchestration without handoffs."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
with pytest.raises(ValueError):
|
|
HandoffOrchestration(members=[agent_a, agent_b], handoffs={})
|
|
|
|
|
|
def test_init_with_invalid_handoff():
|
|
"""Test the initialization of HandoffOrchestration with invalid handoff."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
# Invalid handoff agent name
|
|
with pytest.raises(ValueError):
|
|
HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={
|
|
agent_a.name: {agent_b.name: "test", "invalid_agent_name": "test"},
|
|
agent_b.name: {agent_a.name: "test"},
|
|
},
|
|
)
|
|
|
|
# Invalid handoff agent name (not in members)
|
|
with pytest.raises(ValueError):
|
|
HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={
|
|
"invalid_agent_name": {agent_b.name: "test"},
|
|
agent_b.name: {agent_a.name: "test"},
|
|
},
|
|
)
|
|
|
|
# Cannot handoff to self
|
|
with pytest.raises(ValueError):
|
|
HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={
|
|
agent_a.name: {agent_a.name: "test"},
|
|
agent_b.name: {agent_a.name: "test"},
|
|
},
|
|
)
|
|
|
|
|
|
def test_init_with_duplicate_handoffs():
|
|
"""Test the initialization of HandoffOrchestration with duplicate handoffs."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
# Uniqueness guarantee
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={
|
|
agent_a.name: {agent_b.name: "test 1", agent_b.name: "test 2"},
|
|
},
|
|
)
|
|
|
|
assert len(orchestration._handoffs[agent_a.name]) == 1
|
|
|
|
|
|
def test_init_with_dictionary_handoffs():
|
|
"""Test the initialization of HandoffOrchestration with dictionary handoffs."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
orchestration_handoffs = OrchestrationHandoffs(
|
|
{
|
|
agent_a.name: {agent_b.name: "test 1"},
|
|
agent_b.name: {agent_a.name: "test 2"},
|
|
},
|
|
)
|
|
|
|
assert len(orchestration_handoffs) == 2
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name == agent_b.name
|
|
assert handoff_description == "test 1"
|
|
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_b.name].items():
|
|
assert handoff_agent_name == agent_a.name
|
|
assert handoff_description == "test 2"
|
|
|
|
|
|
def test_orchestration_handoff_add():
|
|
"""Test the add method of the OrchestrationHandoffs."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
orchestration_handoffs = OrchestrationHandoffs().add(agent_a, agent_b).add(agent_b, agent_a)
|
|
|
|
assert isinstance(orchestration_handoffs, OrchestrationHandoffs)
|
|
assert len(orchestration_handoffs) == 2
|
|
assert len(orchestration_handoffs[agent_a.name]) == 1
|
|
assert len(orchestration_handoffs[agent_b.name]) == 1
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name == agent_b.name
|
|
assert handoff_description == ""
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_b.name].items():
|
|
assert handoff_agent_name == agent_a.name
|
|
assert handoff_description == ""
|
|
|
|
|
|
def test_orchestration_handoff_add_many():
|
|
"""Test the add_many method of the OrchestrationHandoffs."""
|
|
agent_a = MockAgent(description="agent_a")
|
|
agent_b = MockAgent(description="agent_b")
|
|
agent_c = MockAgent(description="agent_c")
|
|
|
|
# Case 1: Agent instance as source and dictionary as handoffs
|
|
orchestration_handoffs = OrchestrationHandoffs().add_many(
|
|
agent_a,
|
|
{agent_b.name: "test 1", agent_c.name: "test 2"},
|
|
)
|
|
|
|
assert isinstance(orchestration_handoffs, OrchestrationHandoffs)
|
|
assert len(orchestration_handoffs) == 1
|
|
assert len(orchestration_handoffs[agent_a.name]) == 2
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name in [agent_b.name, agent_c.name]
|
|
assert handoff_description in ["test 1", "test 2"]
|
|
|
|
# Case 2: Agent name as source and list of agents as handoffs
|
|
orchestration_handoffs = OrchestrationHandoffs().add_many(
|
|
agent_a.name,
|
|
{agent_b.name: "test 1", agent_c.name: "test 2"},
|
|
)
|
|
|
|
assert isinstance(orchestration_handoffs, OrchestrationHandoffs)
|
|
assert len(orchestration_handoffs) == 1
|
|
assert len(orchestration_handoffs[agent_a.name]) == 2
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name in [agent_b.name, agent_c.name]
|
|
assert handoff_description in ["test 1", "test 2"]
|
|
|
|
# Case 3: Agent instance as source and list of agents as handoffs
|
|
orchestration_handoffs = OrchestrationHandoffs().add_many(agent_a, [agent_b, agent_c])
|
|
assert isinstance(orchestration_handoffs, OrchestrationHandoffs)
|
|
assert len(orchestration_handoffs) == 1
|
|
assert len(orchestration_handoffs[agent_a.name]) == 2
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name in [agent_b.name, agent_c.name]
|
|
assert handoff_description in [agent_b.description, agent_c.description]
|
|
|
|
# Case 4: Agent name as source and list of agent names as handoffs
|
|
orchestration_handoffs = OrchestrationHandoffs().add_many(agent_a.name, [agent_b.name, agent_c.name])
|
|
assert isinstance(orchestration_handoffs, OrchestrationHandoffs)
|
|
assert len(orchestration_handoffs) == 1
|
|
assert len(orchestration_handoffs[agent_a.name]) == 2
|
|
for handoff_agent_name, handoff_description in orchestration_handoffs[agent_a.name].items():
|
|
assert handoff_agent_name in [agent_b.name, agent_c.name]
|
|
assert handoff_description == ""
|
|
|
|
|
|
async def test_prepare():
|
|
"""Test the prepare method of the HandoffOrchestration."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
agent_c = MockAgent()
|
|
|
|
runtime = MockRuntime()
|
|
|
|
package_path = "semantic_kernel.agents.orchestration.handoffs"
|
|
with (
|
|
patch(f"{package_path}.HandoffOrchestration._start"),
|
|
patch(f"{package_path}.HandoffAgentActor.register") as mock_agent_actor_register,
|
|
patch.object(runtime, "add_subscription") as mock_add_subscription,
|
|
):
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b, agent_c],
|
|
handoffs={
|
|
agent_a.name: {agent_b.name: "test"},
|
|
agent_b.name: {agent_c.name: "test"},
|
|
agent_c.name: {agent_a.name: "test"},
|
|
},
|
|
)
|
|
await orchestration.invoke(task="test_message", runtime=runtime)
|
|
|
|
assert mock_agent_actor_register.call_count == 3
|
|
assert mock_add_subscription.call_count == 3
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info < (3, 11),
|
|
reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
|
|
)
|
|
async def test_invoke():
|
|
"""Test the prepare method of the HandoffOrchestration."""
|
|
with (
|
|
patch.object(HandoffAgentActor, "__init__", wraps=HandoffAgentActor.__init__, autospec=True) as mock_init,
|
|
patch.object(MockAgent, "invoke_stream", wraps=MockAgent.invoke_stream, autospec=True) as mock_invoke_stream,
|
|
):
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
agent_c = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b, agent_c],
|
|
handoffs={
|
|
agent_a.name: {agent_b.name: "test", agent_c.name: "test"},
|
|
agent_b.name: {agent_a.name: "test"},
|
|
},
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
|
|
assert mock_init.call_args_list[0][0][3] == {agent_b.name: "test", agent_c.name: "test"}
|
|
assert isinstance(mock_invoke_stream.call_args_list[0][1]["kernel"], Kernel)
|
|
kernel = mock_invoke_stream.call_args_list[0][1]["kernel"]
|
|
assert HANDOFF_PLUGIN_NAME in kernel.plugins
|
|
assert (
|
|
len(kernel.plugins[HANDOFF_PLUGIN_NAME].functions) == 3
|
|
) # two handoff functions + complete task function
|
|
# The kernel in the agent should not be modified
|
|
assert len(agent_a.kernel.plugins) == 0
|
|
assert len(agent_b.kernel.plugins) == 0
|
|
assert len(agent_c.kernel.plugins) == 0
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info < (3, 11),
|
|
reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
|
|
)
|
|
async def test_invoke_with_list():
|
|
"""Test the invoke method of the HandoffOrchestration with a list of messages."""
|
|
with (
|
|
patch.object(MockAgent, "invoke_stream", wraps=MockAgent.invoke_stream, autospec=True) as mock_invoke_stream,
|
|
):
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
messages = [
|
|
ChatMessageContent(role=AuthorRole.USER, content="test_message_1"),
|
|
ChatMessageContent(role=AuthorRole.USER, content="test_message_2"),
|
|
]
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
)
|
|
orchestration_result = await orchestration.invoke(task=messages, runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert mock_invoke_stream.call_count == 1
|
|
# Two messages
|
|
assert len(mock_invoke_stream.call_args_list[0][0][1]) == 2
|
|
# The kernel in the agent should not be modified
|
|
assert len(agent_a.kernel.plugins) == 0
|
|
assert len(agent_b.kernel.plugins) == 0
|
|
|
|
|
|
async def test_invoke_with_response_callback():
|
|
"""Test the invoke method of the HandoffOrchestration with a response callback."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
responses: list[DefaultTypeAlias] = []
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
agent_response_callback=lambda x: responses.append(x),
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert len(responses) == 1
|
|
assert all(isinstance(item, ChatMessageContent) for item in responses)
|
|
assert all(item.content == "mock_response" for item in responses)
|
|
# The kernel in the agent should not be modified
|
|
assert len(agent_a.kernel.plugins) == 0
|
|
assert len(agent_b.kernel.plugins) == 0
|
|
|
|
|
|
async def test_invoke_with_streaming_response_callback():
|
|
"""Test the invoke method of the HandoffOrchestration with a streaming response callback."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
responses: dict[str, list[StreamingChatMessageContent]] = {}
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
streaming_agent_response_callback=lambda x, _: responses.setdefault(x.name, []).append(x),
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert len(responses[agent_a.name]) == 2
|
|
assert all(isinstance(item, StreamingChatMessageContent) for item in responses[agent_a.name])
|
|
agent_a_response = sum(responses[agent_a.name][1:], responses[agent_a.name][0])
|
|
assert agent_a_response.content == "mock_response"
|
|
|
|
# Agent B was not invoked, so it should not have any responses
|
|
assert agent_b.name not in responses or len(responses[agent_b.name]) == 0
|
|
|
|
# The kernel in the agent should not be modified
|
|
assert len(agent_a.kernel.plugins) == 0
|
|
assert len(agent_b.kernel.plugins) == 0
|
|
|
|
|
|
async def test_response_callback_with_handoff_function_call():
|
|
"""Test the response callback of the HandoffOrchestration with a handoff function call."""
|
|
agent_b = MockAgent()
|
|
agent_a = MockAgentWithHandoffFunctionCall(agent_b)
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
responses: list[DefaultTypeAlias] = []
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
agent_response_callback=lambda x: responses.append(x),
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert len(responses) == 3
|
|
|
|
assert responses[0].name == agent_a.name
|
|
assert isinstance(responses[0].items[0], FunctionCallContent)
|
|
assert responses[1].name == agent_a.name
|
|
assert isinstance(responses[1].items[0], FunctionResultContent)
|
|
|
|
assert responses[2].name == agent_b.name
|
|
|
|
|
|
async def test_streaming_response_callback_with_handoff_function_call():
|
|
"""Test the streaming sresponse callback of the HandoffOrchestration with a handoff function call."""
|
|
agent_b = MockAgent()
|
|
agent_a = MockAgentWithHandoffFunctionCall(agent_b)
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
responses: dict[str, list[StreamingChatMessageContent]] = {}
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
streaming_agent_response_callback=lambda x, _: responses.setdefault(x.name, []).append(x),
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert len(responses[agent_a.name]) == 2
|
|
assert len(responses[agent_b.name]) == 2 # 2 chunks
|
|
|
|
assert responses[agent_a.name][0].name == agent_a.name
|
|
assert isinstance(responses[agent_a.name][0].items[0], FunctionCallContent)
|
|
assert responses[agent_a.name][1].name == agent_a.name
|
|
assert isinstance(responses[agent_a.name][1].items[0], FunctionResultContent)
|
|
|
|
assert responses[agent_b.name][0].name == agent_b.name
|
|
assert all([isinstance(response, StreamingChatMessageContent) for response in responses[agent_b.name]])
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info < (3, 11),
|
|
reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
|
|
)
|
|
async def test_invoke_with_human_response_function():
|
|
"""Test the invoke method of the HandoffOrchestration with a human response function."""
|
|
complete_task_agent_instance = MockAgentWithCompleteTaskFunctionCall()
|
|
normal_agent_instance = MockAgent()
|
|
call_sequence = iter([normal_agent_instance.invoke_stream, complete_task_agent_instance.invoke_stream])
|
|
|
|
user_input_count = 0
|
|
|
|
def human_response_function() -> ChatMessageContent:
|
|
# Simulate user input
|
|
nonlocal user_input_count
|
|
user_input_count += 1
|
|
return ChatMessageContent(role=AuthorRole.USER, content="user_input")
|
|
|
|
with (
|
|
patch.object(MockAgent, "invoke_stream") as mock_invoke_stream,
|
|
):
|
|
mock_invoke_stream.side_effect = lambda *args, **kwargs: next(call_sequence)(*args, **kwargs)
|
|
|
|
agent_a = MockAgent(name="agent_a")
|
|
agent_b = MockAgent(name="agent_b")
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
human_response_function=human_response_function,
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert user_input_count == 1
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info < (3, 11),
|
|
reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
|
|
)
|
|
async def test_invoke_with_handoff_function_call():
|
|
"""Test the invoke method of the HandoffOrchestration with a handoff function call."""
|
|
agent_b = MockAgent()
|
|
agent_a = MockAgentWithHandoffFunctionCall(agent_b)
|
|
|
|
with (
|
|
patch.object(
|
|
HandoffAgentActor, "_handoff_to_agent", wraps=HandoffAgentActor._handoff_to_agent, autospec=True
|
|
) as mock_handoff_to_agent,
|
|
):
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
await orchestration_result.get(1.0)
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert mock_handoff_to_agent.call_count == 1
|
|
assert mock_handoff_to_agent.call_args_list[0][0][1] == agent_b.name
|
|
# The kernel in the agent should not be modified
|
|
assert len(agent_a.kernel.plugins) == 0
|
|
assert len(agent_b.kernel.plugins) == 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info < (3, 11),
|
|
reason="Python 3.10 doesn't bound the original function provided to the wraps argument of the patch object.",
|
|
)
|
|
async def test_invoke_cancel_before_completion():
|
|
"""Test the invoke method of the HandoffOrchestration with cancellation before completion."""
|
|
with (
|
|
patch.object(MockAgent, "invoke_stream", wraps=MockAgent.invoke_stream, autospec=True) as mock_invoke_stream,
|
|
):
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
)
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
|
|
# Cancel before first agent completes
|
|
await asyncio.sleep(0.05)
|
|
orchestration_result.cancel()
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
assert mock_invoke_stream.call_count == 1
|
|
|
|
|
|
async def test_invoke_cancel_after_completion():
|
|
"""Test the invoke method of the HandoffOrchestration with cancellation after completion."""
|
|
agent_a = MockAgent()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
)
|
|
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
|
|
# Wait for the orchestration to complete
|
|
await orchestration_result.get(1.0)
|
|
|
|
with pytest.raises(RuntimeError, match="The invocation has already been completed."):
|
|
orchestration_result.cancel()
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
|
|
async def test_invoke_with_agent_raising_exception():
|
|
"""Test the invoke method of the HandoffOrchestration with an agent raising an exception."""
|
|
agent_a = MockAgentWithException()
|
|
agent_b = MockAgent()
|
|
|
|
runtime = InProcessRuntime()
|
|
runtime.start()
|
|
|
|
try:
|
|
orchestration = HandoffOrchestration(
|
|
members=[agent_a, agent_b],
|
|
handoffs={agent_a.name: {agent_b.name: "test"}},
|
|
)
|
|
|
|
orchestration_result = await orchestration.invoke(task="test_message", runtime=runtime)
|
|
|
|
with pytest.raises(RuntimeError, match="Mock agent exception"):
|
|
await orchestration_result.get(1.0)
|
|
assert orchestration_result.exception is not None
|
|
finally:
|
|
await runtime.stop_when_idle()
|
|
|
|
|
|
# endregion
|