mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Fix WorkflowAgent event handling and kwargs forwarding (#2946)
* Fix kwargs propagation through workflow.as_agent() * Fix WorkflowAgent to respect AgentExecutor output_response setting
This commit is contained in:
committed by
GitHub
Unverified
parent
a841bdd1cc
commit
b0a7a1fcb8
@@ -1,11 +1,13 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
AgentRunUpdateEvent,
|
||||
@@ -422,6 +424,48 @@ class TestWorkflowAgent:
|
||||
assert isinstance(updates[2].raw_representation, CustomData)
|
||||
assert updates[2].raw_representation.value == 42
|
||||
|
||||
async def test_workflow_as_agent_yield_output_with_list_of_chat_messages(self) -> None:
|
||||
"""Test that yield_output with list[ChatMessage] extracts contents from all messages.
|
||||
|
||||
Note: TextContent items are coalesced by _finalize_response, so multiple text contents
|
||||
become a single merged TextContent in the final response.
|
||||
"""
|
||||
|
||||
@executor
|
||||
async def list_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
|
||||
# Yield a list of ChatMessages (as SequentialBuilder does)
|
||||
msg_list = [
|
||||
ChatMessage(role=Role.USER, contents=[TextContent(text="first message")]),
|
||||
ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="second message")]),
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[TextContent(text="third"), TextContent(text="fourth")],
|
||||
),
|
||||
]
|
||||
await ctx.yield_output(msg_list)
|
||||
|
||||
workflow = WorkflowBuilder().set_start_executor(list_yielding_executor).build()
|
||||
agent = workflow.as_agent("list-msg-agent")
|
||||
|
||||
# Verify streaming returns the update with all 4 contents before coalescing
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in agent.run_stream("test"):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 1
|
||||
assert len(updates[0].contents) == 4
|
||||
texts = [c.text for c in updates[0].contents if isinstance(c, TextContent)]
|
||||
assert texts == ["first message", "second message", "third", "fourth"]
|
||||
|
||||
# Verify run() coalesces text contents (expected behavior)
|
||||
result = await agent.run("test")
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
# TextContent items are coalesced into one
|
||||
assert len(result.messages[0].contents) == 1
|
||||
assert result.messages[0].text == "first messagesecond messagethirdfourth"
|
||||
|
||||
async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
|
||||
"""Test that conversation history from thread is included when running WorkflowAgent.
|
||||
|
||||
@@ -521,6 +565,142 @@ class TestWorkflowAgent:
|
||||
checkpoints = await checkpoint_storage.list_checkpoints(workflow.id)
|
||||
assert len(checkpoints) > 0, "Checkpoints should have been created when checkpoint_storage is provided"
|
||||
|
||||
async def test_agent_executor_output_response_false_filters_streaming_events(self):
|
||||
"""Test that AgentExecutor with output_response=False does not surface streaming events."""
|
||||
|
||||
class MockAgent(AgentProtocol):
|
||||
"""Mock agent for testing."""
|
||||
|
||||
def __init__(self, name: str, response_text: str) -> None:
|
||||
self._name = name
|
||||
self._response_text = response_text
|
||||
self._description: str | None = None
|
||||
|
||||
@property
|
||||
def name(self) -> str | None:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def description(self) -> str | None:
|
||||
return self._description
|
||||
|
||||
def get_new_thread(self) -> AgentThread:
|
||||
return AgentThread()
|
||||
|
||||
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentRunResponse:
|
||||
return AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text=self._response_text)],
|
||||
text=self._response_text,
|
||||
)
|
||||
|
||||
async def run_stream(
|
||||
self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any
|
||||
) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
for word in self._response_text.split():
|
||||
yield AgentRunResponseUpdate(
|
||||
contents=[TextContent(text=word + " ")],
|
||||
role=Role.ASSISTANT,
|
||||
author_name=self._name,
|
||||
)
|
||||
|
||||
@executor
|
||||
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
|
||||
from agent_framework import AgentExecutorRequest
|
||||
|
||||
await ctx.yield_output("Start output")
|
||||
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
|
||||
|
||||
# Build workflow: start -> agent1 (no output) -> agent2 (output_response=True)
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.register_executor(lambda: start_executor, "start")
|
||||
.register_agent(lambda: MockAgent("agent1", "Agent1 output - should NOT appear"), "agent1")
|
||||
.register_agent(
|
||||
lambda: MockAgent("agent2", "Agent2 output - SHOULD appear"), "agent2", output_response=True
|
||||
)
|
||||
.set_start_executor("start")
|
||||
.add_edge("start", "agent1")
|
||||
.add_edge("agent1", "agent2")
|
||||
.build()
|
||||
)
|
||||
|
||||
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
|
||||
result = await agent.run("Test input")
|
||||
|
||||
# Collect all message texts
|
||||
texts = [msg.text for msg in result.messages if msg.text]
|
||||
|
||||
# Start output should appear (from yield_output)
|
||||
assert any("Start output" in t for t in texts), "Start output should appear"
|
||||
|
||||
# Agent1 output should NOT appear (output_response=False)
|
||||
assert not any("Agent1" in t for t in texts), "Agent1 output should NOT appear"
|
||||
|
||||
# Agent2 output should appear (output_response=True)
|
||||
assert any("Agent2" in t for t in texts), "Agent2 output should appear"
|
||||
|
||||
async def test_agent_executor_output_response_no_duplicate_from_workflow_output_event(self):
|
||||
"""Test that AgentExecutor with output_response=True does not duplicate content."""
|
||||
|
||||
class MockAgent(AgentProtocol):
|
||||
"""Mock agent for testing."""
|
||||
|
||||
def __init__(self, name: str, response_text: str) -> None:
|
||||
self._name = name
|
||||
self._response_text = response_text
|
||||
self._description: str | None = None
|
||||
|
||||
@property
|
||||
def name(self) -> str | None:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def description(self) -> str | None:
|
||||
return self._description
|
||||
|
||||
def get_new_thread(self) -> AgentThread:
|
||||
return AgentThread()
|
||||
|
||||
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentRunResponse:
|
||||
return AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text=self._response_text)],
|
||||
text=self._response_text,
|
||||
)
|
||||
|
||||
async def run_stream(
|
||||
self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any
|
||||
) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
yield AgentRunResponseUpdate(
|
||||
contents=[TextContent(text=self._response_text)],
|
||||
role=Role.ASSISTANT,
|
||||
author_name=self._name,
|
||||
)
|
||||
|
||||
@executor
|
||||
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
|
||||
from agent_framework import AgentExecutorRequest
|
||||
|
||||
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
|
||||
|
||||
# Build workflow with single agent that has output_response=True
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.register_executor(lambda: start_executor, "start")
|
||||
.register_agent(lambda: MockAgent("agent", "Unique response text"), "agent", output_response=True)
|
||||
.set_start_executor("start")
|
||||
.add_edge("start", "agent")
|
||||
.build()
|
||||
)
|
||||
|
||||
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
|
||||
result = await agent.run("Test input")
|
||||
|
||||
# Count occurrences of the unique response text
|
||||
unique_text_count = sum(1 for msg in result.messages if msg.text and "Unique response text" in msg.text)
|
||||
|
||||
# Should appear exactly once (not duplicated from both streaming and WorkflowOutputEvent)
|
||||
assert unique_text_count == 1, f"Response should appear exactly once, but appeared {unique_text_count} times"
|
||||
|
||||
|
||||
class TestWorkflowAgentMergeUpdates:
|
||||
"""Test cases specifically for the WorkflowAgent.merge_updates static method."""
|
||||
|
||||
@@ -492,6 +492,117 @@ async def test_magentic_kwargs_stored_in_shared_state() -> None:
|
||||
# endregion
|
||||
|
||||
|
||||
# region WorkflowAgent (as_agent) kwargs Tests
|
||||
|
||||
|
||||
async def test_workflow_as_agent_run_propagates_kwargs_to_underlying_agent() -> None:
|
||||
"""Test that kwargs passed to workflow_agent.run() flow through to the underlying agents."""
|
||||
agent = _KwargsCapturingAgent(name="inner_agent")
|
||||
workflow = SequentialBuilder().participants([agent]).build()
|
||||
workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
|
||||
|
||||
custom_data = {"endpoint": "https://api.example.com", "version": "v1"}
|
||||
user_token = {"user_name": "alice", "access_level": "admin"}
|
||||
|
||||
_ = await workflow_agent.run(
|
||||
"test message",
|
||||
custom_data=custom_data,
|
||||
user_token=user_token,
|
||||
)
|
||||
|
||||
# Verify inner agent received kwargs
|
||||
assert len(agent.captured_kwargs) >= 1, "Inner agent should have been invoked at least once"
|
||||
received = agent.captured_kwargs[0]
|
||||
assert "custom_data" in received, "Inner agent should receive custom_data kwarg"
|
||||
assert "user_token" in received, "Inner agent should receive user_token kwarg"
|
||||
assert received["custom_data"] == custom_data
|
||||
assert received["user_token"] == user_token
|
||||
|
||||
|
||||
async def test_workflow_as_agent_run_stream_propagates_kwargs_to_underlying_agent() -> None:
|
||||
"""Test that kwargs passed to workflow_agent.run_stream() flow through to the underlying agents."""
|
||||
agent = _KwargsCapturingAgent(name="inner_agent")
|
||||
workflow = SequentialBuilder().participants([agent]).build()
|
||||
workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
|
||||
|
||||
custom_data = {"session_id": "xyz123"}
|
||||
api_token = "secret-token"
|
||||
|
||||
async for _ in workflow_agent.run_stream(
|
||||
"test message",
|
||||
custom_data=custom_data,
|
||||
api_token=api_token,
|
||||
):
|
||||
pass
|
||||
|
||||
# Verify inner agent received kwargs
|
||||
assert len(agent.captured_kwargs) >= 1, "Inner agent should have been invoked at least once"
|
||||
received = agent.captured_kwargs[0]
|
||||
assert "custom_data" in received, "Inner agent should receive custom_data kwarg"
|
||||
assert "api_token" in received, "Inner agent should receive api_token kwarg"
|
||||
assert received["custom_data"] == custom_data
|
||||
assert received["api_token"] == api_token
|
||||
|
||||
|
||||
async def test_workflow_as_agent_propagates_kwargs_to_multiple_agents() -> None:
|
||||
"""Test that kwargs flow to all agents when using workflow.as_agent()."""
|
||||
agent1 = _KwargsCapturingAgent(name="agent1")
|
||||
agent2 = _KwargsCapturingAgent(name="agent2")
|
||||
workflow = SequentialBuilder().participants([agent1, agent2]).build()
|
||||
workflow_agent = workflow.as_agent(name="MultiAgentWorkflow")
|
||||
|
||||
custom_data = {"batch_id": "batch-001"}
|
||||
|
||||
_ = await workflow_agent.run("test message", custom_data=custom_data)
|
||||
|
||||
# Both agents should have received kwargs
|
||||
assert len(agent1.captured_kwargs) >= 1, "First agent should be invoked"
|
||||
assert len(agent2.captured_kwargs) >= 1, "Second agent should be invoked"
|
||||
assert agent1.captured_kwargs[0].get("custom_data") == custom_data
|
||||
assert agent2.captured_kwargs[0].get("custom_data") == custom_data
|
||||
|
||||
|
||||
async def test_workflow_as_agent_kwargs_with_none_values() -> None:
|
||||
"""Test that kwargs with None values are passed through correctly via as_agent()."""
|
||||
agent = _KwargsCapturingAgent(name="none_test_agent")
|
||||
workflow = SequentialBuilder().participants([agent]).build()
|
||||
workflow_agent = workflow.as_agent(name="NoneTestWorkflow")
|
||||
|
||||
_ = await workflow_agent.run("test", optional_param=None, other_param="value")
|
||||
|
||||
assert len(agent.captured_kwargs) >= 1
|
||||
received = agent.captured_kwargs[0]
|
||||
assert "optional_param" in received
|
||||
assert received["optional_param"] is None
|
||||
assert received["other_param"] == "value"
|
||||
|
||||
|
||||
async def test_workflow_as_agent_kwargs_with_complex_nested_data() -> None:
|
||||
"""Test that complex nested data structures flow through correctly via as_agent()."""
|
||||
agent = _KwargsCapturingAgent(name="nested_agent")
|
||||
workflow = SequentialBuilder().participants([agent]).build()
|
||||
workflow_agent = workflow.as_agent(name="NestedDataWorkflow")
|
||||
|
||||
complex_data = {
|
||||
"level1": {
|
||||
"level2": {
|
||||
"level3": ["a", "b", "c"],
|
||||
"number": 42,
|
||||
},
|
||||
"list": [1, 2, {"nested": True}],
|
||||
},
|
||||
}
|
||||
|
||||
_ = await workflow_agent.run("test", complex_data=complex_data)
|
||||
|
||||
assert len(agent.captured_kwargs) >= 1
|
||||
received = agent.captured_kwargs[0]
|
||||
assert received.get("complex_data") == complex_data
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region SubWorkflow (WorkflowExecutor) Tests
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user