chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,721 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user