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:
@@ -365,6 +365,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
)
|
||||
|
||||
all_updates: list[AgentResponseUpdate] = []
|
||||
streamed_artifact_ids_by_task: dict[str, set[str]] = {}
|
||||
async for item in a2a_stream:
|
||||
if isinstance(item, A2AMessage):
|
||||
# Process A2A Message
|
||||
@@ -378,12 +379,21 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
all_updates.append(update)
|
||||
yield update
|
||||
elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], Task):
|
||||
task, _update_event = item
|
||||
for update in self._updates_from_task(
|
||||
task, update_event = item
|
||||
updates = self._updates_from_task(
|
||||
task,
|
||||
update_event=update_event,
|
||||
background=background,
|
||||
emit_intermediate=emit_intermediate,
|
||||
streamed_artifact_ids=streamed_artifact_ids_by_task.get(task.id),
|
||||
)
|
||||
if isinstance(update_event, TaskArtifactUpdateEvent) and any(
|
||||
update.raw_representation is update_event for update in updates
|
||||
):
|
||||
streamed_artifact_ids_by_task.setdefault(task.id, set()).add(update_event.artifact.artifact_id)
|
||||
if task.status.state in TERMINAL_TASK_STATES:
|
||||
streamed_artifact_ids_by_task.pop(task.id, None)
|
||||
for update in updates:
|
||||
all_updates.append(update)
|
||||
yield update
|
||||
else:
|
||||
@@ -403,8 +413,10 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
self,
|
||||
task: Task,
|
||||
*,
|
||||
update_event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None,
|
||||
background: bool = False,
|
||||
emit_intermediate: bool = False,
|
||||
streamed_artifact_ids: set[str] | None = None,
|
||||
) -> list[AgentResponseUpdate]:
|
||||
"""Convert an A2A Task into AgentResponseUpdate(s).
|
||||
|
||||
@@ -418,8 +430,21 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
"""
|
||||
status = task.status
|
||||
|
||||
if (
|
||||
emit_intermediate
|
||||
and update_event is not None
|
||||
and (event_updates := self._updates_from_task_update_event(update_event))
|
||||
):
|
||||
return event_updates
|
||||
|
||||
if status.state in TERMINAL_TASK_STATES:
|
||||
task_messages = self._parse_messages_from_task(task)
|
||||
if task.artifacts is not None and streamed_artifact_ids:
|
||||
task_messages = [
|
||||
message
|
||||
for message in task_messages
|
||||
if getattr(message.raw_representation, "artifact_id", None) not in streamed_artifact_ids
|
||||
]
|
||||
if task_messages:
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
@@ -431,6 +456,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
)
|
||||
for message in task_messages
|
||||
]
|
||||
if task.artifacts is not None:
|
||||
return []
|
||||
return [AgentResponseUpdate(contents=[], role="assistant", response_id=task.id, raw_representation=task)]
|
||||
|
||||
if background and status.state in IN_PROGRESS_TASK_STATES:
|
||||
@@ -467,6 +494,44 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
|
||||
|
||||
return []
|
||||
|
||||
def _updates_from_task_update_event(
|
||||
self, update_event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent
|
||||
) -> list[AgentResponseUpdate]:
|
||||
"""Convert A2A task update events into streaming AgentResponseUpdates."""
|
||||
if isinstance(update_event, TaskArtifactUpdateEvent):
|
||||
contents = self._parse_contents_from_a2a(update_event.artifact.parts)
|
||||
if not contents:
|
||||
return []
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=contents,
|
||||
role="assistant",
|
||||
response_id=update_event.task_id,
|
||||
message_id=update_event.artifact.artifact_id,
|
||||
raw_representation=update_event,
|
||||
)
|
||||
]
|
||||
|
||||
if not isinstance(update_event, TaskStatusUpdateEvent):
|
||||
return []
|
||||
|
||||
message = update_event.status.message
|
||||
if message is None or not message.parts:
|
||||
return []
|
||||
|
||||
contents = self._parse_contents_from_a2a(message.parts)
|
||||
if not contents:
|
||||
return []
|
||||
|
||||
return [
|
||||
AgentResponseUpdate(
|
||||
contents=contents,
|
||||
role="assistant" if message.role == A2ARole.agent else "user",
|
||||
response_id=update_event.task_id,
|
||||
raw_representation=update_event,
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _build_continuation_token(task: Task) -> A2AContinuationToken | None:
|
||||
"""Build an A2AContinuationToken from an A2A Task if it is still in progress."""
|
||||
|
||||
Reference in New Issue
Block a user