mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: handle streamed A2A update events (#4919)
* Python: handle streamed A2A update events * Python: preserve terminal A2A artifacts during streaming * Python: harden streamed A2A update event handling * Python: simplify streamed A2A update guard --------- Co-authored-by: sztoplover-bit <253473756+sztoplover-bit@users.noreply.github.com> Co-authored-by: Giles Odigwe <79032838+giles17@users.noreply.github.com>
This commit is contained in:
perry
GitHub
@@ -14,8 +14,10 @@ from a2a.types import (
|
||||
FileWithUri,
|
||||
Part,
|
||||
Task,
|
||||
TaskArtifactUpdateEvent,
|
||||
TaskState,
|
||||
TaskStatus,
|
||||
TaskStatusUpdateEvent,
|
||||
TextPart,
|
||||
)
|
||||
from a2a.types import Message as A2AMessage
|
||||
@@ -1189,4 +1191,201 @@ async def test_streaming_working_update_with_empty_parts_is_skipped(
|
||||
assert updates[0].contents[0].text == "Result"
|
||||
|
||||
|
||||
async def test_streaming_artifact_update_event_yields_content(
|
||||
a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
|
||||
) -> None:
|
||||
"""Test that streaming artifact update events yield incremental content."""
|
||||
task = Task(id="task-art", context_id="ctx-art", status=TaskStatus(state=TaskState.working, message=None))
|
||||
artifact = Artifact(
|
||||
artifact_id="artifact-1",
|
||||
parts=[Part(root=TextPart(text="Hello"))],
|
||||
)
|
||||
update_event = TaskArtifactUpdateEvent(task_id="task-art", context_id="ctx-art", artifact=artifact, append=False)
|
||||
mock_a2a_client.responses.append((task, update_event))
|
||||
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
async for update in a2a_agent.run("Hello", stream=True):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 1
|
||||
assert updates[0].text == "Hello"
|
||||
assert updates[0].message_id == "artifact-1"
|
||||
assert updates[0].raw_representation == update_event
|
||||
|
||||
|
||||
async def test_streaming_status_update_event_yields_content(
|
||||
a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
|
||||
) -> None:
|
||||
"""Test that streaming status update events surface message content directly from the update event."""
|
||||
update_event = TaskStatusUpdateEvent(
|
||||
task_id="task-status",
|
||||
context_id="ctx-status",
|
||||
status=TaskStatus(
|
||||
state=TaskState.working,
|
||||
message=A2AMessage(
|
||||
message_id=str(uuid4()),
|
||||
role=A2ARole.agent,
|
||||
parts=[Part(root=TextPart(text="Still working"))],
|
||||
),
|
||||
),
|
||||
final=False,
|
||||
)
|
||||
task = Task(id="task-status", context_id="ctx-status", status=TaskStatus(state=TaskState.working, message=None))
|
||||
mock_a2a_client.responses.append((task, update_event))
|
||||
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
async for update in a2a_agent.run("Hello", stream=True):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 1
|
||||
assert updates[0].text == "Still working"
|
||||
assert updates[0].role == "assistant"
|
||||
assert updates[0].raw_representation == update_event
|
||||
|
||||
|
||||
async def test_streaming_artifact_update_event_does_not_duplicate_terminal_task_artifacts(
|
||||
a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
|
||||
) -> None:
|
||||
"""Test that streamed artifact chunks are not re-emitted from the final terminal task."""
|
||||
working_task = Task(id="task-art-dup", context_id="ctx-art-dup", status=TaskStatus(state=TaskState.working))
|
||||
first_chunk = TaskArtifactUpdateEvent(
|
||||
task_id="task-art-dup",
|
||||
context_id="ctx-art-dup",
|
||||
artifact=Artifact(
|
||||
artifact_id="artifact-dup",
|
||||
parts=[Part(root=TextPart(text="Hello "))],
|
||||
),
|
||||
append=False,
|
||||
)
|
||||
second_chunk = TaskArtifactUpdateEvent(
|
||||
task_id="task-art-dup",
|
||||
context_id="ctx-art-dup",
|
||||
artifact=Artifact(
|
||||
artifact_id="artifact-dup",
|
||||
parts=[Part(root=TextPart(text="world"))],
|
||||
),
|
||||
append=True,
|
||||
)
|
||||
terminal_task = Task(
|
||||
id="task-art-dup",
|
||||
context_id="ctx-art-dup",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
artifacts=[
|
||||
Artifact(
|
||||
artifact_id="artifact-dup",
|
||||
parts=[Part(root=TextPart(text="Hello world"))],
|
||||
)
|
||||
],
|
||||
)
|
||||
terminal_event = TaskStatusUpdateEvent(
|
||||
task_id="task-art-dup",
|
||||
context_id="ctx-art-dup",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
final=True,
|
||||
)
|
||||
|
||||
mock_a2a_client.responses.extend(
|
||||
[
|
||||
(working_task, first_chunk),
|
||||
(working_task, second_chunk),
|
||||
(terminal_task, terminal_event),
|
||||
]
|
||||
)
|
||||
|
||||
stream = a2a_agent.run("Hello", stream=True)
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert [update.text for update in updates] == ["Hello ", "world"]
|
||||
assert response.text == "Hello world"
|
||||
assert len(response.messages) == 1
|
||||
|
||||
|
||||
async def test_streaming_terminal_task_artifacts_are_emitted_when_terminal_event_has_no_content(
|
||||
a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
|
||||
) -> None:
|
||||
"""Test that terminal task artifacts are still emitted when the final status event has no message."""
|
||||
terminal_task = Task(
|
||||
id="task-art-final",
|
||||
context_id="ctx-art-final",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
artifacts=[
|
||||
Artifact(
|
||||
artifact_id="artifact-final",
|
||||
parts=[Part(root=TextPart(text="Final artifact"))],
|
||||
)
|
||||
],
|
||||
)
|
||||
terminal_event = TaskStatusUpdateEvent(
|
||||
task_id="task-art-final",
|
||||
context_id="ctx-art-final",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
final=True,
|
||||
)
|
||||
mock_a2a_client.responses.append((terminal_task, terminal_event))
|
||||
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
async for update in a2a_agent.run("Hello", stream=True):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 1
|
||||
assert updates[0].text == "Final artifact"
|
||||
assert updates[0].message_id == "artifact-final"
|
||||
|
||||
|
||||
async def test_streaming_terminal_task_only_emits_unstreamed_artifacts(
|
||||
a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient
|
||||
) -> None:
|
||||
"""Test that the terminal task only emits artifacts that were not already streamed incrementally."""
|
||||
working_task = Task(id="task-art-mixed", context_id="ctx-art-mixed", status=TaskStatus(state=TaskState.working))
|
||||
streamed_chunk = TaskArtifactUpdateEvent(
|
||||
task_id="task-art-mixed",
|
||||
context_id="ctx-art-mixed",
|
||||
artifact=Artifact(
|
||||
artifact_id="artifact-streamed",
|
||||
parts=[Part(root=TextPart(text="Hello"))],
|
||||
),
|
||||
append=False,
|
||||
)
|
||||
terminal_task = Task(
|
||||
id="task-art-mixed",
|
||||
context_id="ctx-art-mixed",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
artifacts=[
|
||||
Artifact(
|
||||
artifact_id="artifact-streamed",
|
||||
parts=[Part(root=TextPart(text="Hello"))],
|
||||
),
|
||||
Artifact(
|
||||
artifact_id="artifact-final",
|
||||
parts=[Part(root=TextPart(text="Goodbye"))],
|
||||
),
|
||||
],
|
||||
)
|
||||
terminal_event = TaskStatusUpdateEvent(
|
||||
task_id="task-art-mixed",
|
||||
context_id="ctx-art-mixed",
|
||||
status=TaskStatus(state=TaskState.completed, message=None),
|
||||
final=True,
|
||||
)
|
||||
|
||||
mock_a2a_client.responses.extend(
|
||||
[
|
||||
(working_task, streamed_chunk),
|
||||
(terminal_task, terminal_event),
|
||||
]
|
||||
)
|
||||
|
||||
stream = a2a_agent.run("Hello", stream=True)
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert [update.text for update in updates] == ["Hello", "Goodbye"]
|
||||
assert [message.text for message in response.messages] == ["Hello", "Goodbye"]
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
Reference in New Issue
Block a user