diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 6ffebae8700..d45e821a38a 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2224,6 +2224,11 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse and response.messages[-1].message_id and response.messages[-1].message_id != update.message_id ) + or ( + update.author_name + and response.messages[-1].author_name + and response.messages[-1].author_name != update.author_name + ) or (update.role and response.messages[-1].role != update.role) ): is_new_message = True diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 2c41acd93b4..f9945b42f08 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2000,6 +2000,63 @@ def test_agent_run_response_from_updates(agent_response_update: AgentResponseUpd assert response.text == "Test contentTest content" +@pytest.mark.parametrize( + ("response_type", "update_type"), + [(ChatResponse, ChatResponseUpdate), (AgentResponse, AgentResponseUpdate)], + ids=["chat", "agent"], +) +def test_from_updates_starts_new_message_when_author_changes(response_type: Any, update_type: Any) -> None: + """Test that updates from different authors without message ids are not merged into one message.""" + response = response_type.from_updates([ + update_type(contents=[Content.from_text("Draft ")], role="assistant", author_name="writer"), + update_type(contents=[Content.from_text("text.")], role="assistant"), + update_type(contents=[Content.from_text("Looks good.")], role="assistant", author_name="reviewer"), + ]) + + assert [(message.author_name, message.text) for message in response.messages] == [ + ("writer", "Draft text."), + ("reviewer", "Looks good."), + ] + + +def test_from_updates_keeps_one_message_when_author_is_first_set() -> None: + """Test that a message without an author takes the author of a later update instead of splitting.""" + response = AgentResponse.from_updates([ + AgentResponseUpdate(contents=[Content.from_text("Hello ")], role="assistant"), + AgentResponseUpdate(contents=[Content.from_text("world")], role="assistant", author_name="writer"), + ]) + + assert [(message.author_name, message.text) for message in response.messages] == [("writer", "Hello world")] + + +def test_from_updates_starts_new_message_for_interleaved_authors() -> None: + """Test that interleaved updates from two authors keep each run of updates with its author.""" + response = AgentResponse.from_updates([ + AgentResponseUpdate(contents=[Content.from_text("A1")], role="assistant", author_name="alpha"), + AgentResponseUpdate(contents=[Content.from_text("B1")], role="assistant", author_name="beta"), + AgentResponseUpdate(contents=[Content.from_text("A2")], role="assistant", author_name="alpha"), + ]) + + assert [(message.author_name, message.text) for message in response.messages] == [ + ("alpha", "A1"), + ("beta", "B1"), + ("alpha", "A2"), + ] + + +def test_from_updates_starts_new_message_for_authors_sharing_a_message_id() -> None: + """Test that a different author starts a new message even when the message id is the same.""" + response = AgentResponse.from_updates([ + AgentResponseUpdate(contents=[Content.from_text("Hi")], role="assistant", author_name="alpha", message_id="m1"), + AgentResponseUpdate(contents=[Content.from_text("Yo")], role="assistant", author_name="beta", message_id="m1"), + ]) + + assert [(message.author_name, message.message_id, message.text) for message in response.messages] == [ + ("alpha", "m1", "Hi"), + ("beta", "m1", "Yo"), + ] + + def test_agent_run_response_from_updates_uses_last_non_none_agent_id() -> None: response = AgentResponse.from_updates([ AgentResponseUpdate(agent_id="first-agent"), diff --git a/python/packages/orchestrations/tests/test_sequential.py b/python/packages/orchestrations/tests/test_sequential.py index 162b2e8971b..4e4006456ed 100644 --- a/python/packages/orchestrations/tests/test_sequential.py +++ b/python/packages/orchestrations/tests/test_sequential.py @@ -193,6 +193,25 @@ async def test_sequential_as_agent_returns_only_last_agent_response() -> None: assert "hello as_agent" not in combined +async def test_sequential_as_agent_streaming_keeps_one_message_per_agent() -> None: + """Streaming `workflow.as_agent()` with `output_from="all"` keeps each agent's reply as its own message, + the same messages a non-streaming run returns.""" + + def build_agent() -> Any: + participants = [_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2")] + return SequentialBuilder(participants=participants, output_from="all").build().as_agent() + + response = await build_agent().run("hello") + stream = build_agent().run("hello", stream=True) + async for _ in stream: + pass + streamed = await stream.get_final_response() + + assert [m.text for m in response.messages] == ["A1 reply", "A2 reply"] + assert [m.text for m in streamed.messages] == ["A1 reply", "A2 reply"] + assert [m.author_name for m in streamed.messages] == ["A1", "A2"] + + async def test_sequential_with_custom_executor_summary() -> None: """A custom-executor terminator yields its own AgentResponse — that becomes the workflow output.