Skip to content

Commit 5e551c4

Browse files
Update src/a2a/utils/artifact.py
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 6ba923b commit 5e551c4

6 files changed

Lines changed: 331 additions & 2 deletions

File tree

.pre-commit-config.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ repos:
3131
# ===============================================
3232
# Python Hooks
3333
# ===============================================
34+
3435
# Ruff for linting and formatting
3536
- repo: https://github.com/astral-sh/ruff-pre-commit
3637
rev: v0.12.0

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ dev = [
102102
"types-protobuf",
103103
"types-requests",
104104
"pre-commit",
105+
105106
"trio",
106107
"uvicorn>=0.35.0",
107108
"pytest-timeout>=2.4.0",

src/a2a/utils/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Utility functions for the A2A Python SDK."""
22

33
from a2a.utils.artifact import (
4+
ArtifactStreamer,
45
get_artifact_text,
56
new_artifact,
67
new_data_artifact,
@@ -39,6 +40,7 @@
3940
'DEFAULT_RPC_URL',
4041
'EXTENDED_AGENT_CARD_PATH',
4142
'PREV_AGENT_CARD_WELL_KNOWN_PATH',
43+
'ArtifactStreamer',
4244
'append_artifact_to_task',
4345
'are_modalities_compatible',
4446
'build_text_artifact',

src/a2a/utils/artifact.py

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from typing import Any
66

7-
from a2a.types import Artifact, DataPart, Part, TextPart
7+
from a2a.types import Artifact, DataPart, Part, TaskArtifactUpdateEvent, TextPart
88
from a2a.utils.parts import get_text_parts
99

1010

@@ -86,3 +86,79 @@ def get_artifact_text(artifact: Artifact, delimiter: str = '\n') -> str:
8686
A single string containing all text content, or an empty string if no text parts are found.
8787
"""
8888
return delimiter.join(get_text_parts(artifact.parts))
89+
90+
91+
class ArtifactStreamer:
92+
"""Helper for streaming text into a single artifact across multiple events.
93+
94+
Creates a stable artifact ID on construction so all chunks reference
95+
the same artifact, enabling proper append semantics per the A2A spec.
96+
97+
Example::
98+
99+
streamer = ArtifactStreamer(context_id, task_id, name='response')
100+
101+
async for chunk in llm.stream(prompt):
102+
await event_queue.enqueue_event(streamer.append(chunk))
103+
104+
await event_queue.enqueue_event(streamer.finalize())
105+
106+
Args:
107+
context_id: The context ID associated with the task.
108+
task_id: The task ID associated with the streaming session.
109+
name: A human-readable name for the artifact.
110+
artifact_id: An explicit artifact ID. If omitted a UUID is generated.
111+
"""
112+
113+
def __init__(
114+
self,
115+
context_id: str,
116+
task_id: str,
117+
name: str = 'response',
118+
artifact_id: str | None = None,
119+
) -> None:
120+
self._context_id = context_id
121+
self._task_id = task_id
122+
self._name = name
123+
self._artifact_id = artifact_id or str(uuid.uuid4())
124+
125+
def append(self, text: str) -> TaskArtifactUpdateEvent:
126+
"""Emit a chunk to be appended to the streaming artifact.
127+
128+
Args:
129+
text: The incremental text content for this chunk.
130+
131+
Returns:
132+
A ``TaskArtifactUpdateEvent`` with ``append=True`` and
133+
``last_chunk=False``.
134+
"""
135+
return TaskArtifactUpdateEvent(
136+
context_id=self._context_id,
137+
task_id=self._task_id,
138+
append=True,
139+
last_chunk=False,
140+
artifact=Artifact(
141+
artifact_id=self._artifact_id,
142+
name=self._name,
143+
parts=[Part(root=TextPart(text=text))],
144+
)
145+
)
146+
147+
def finalize(self) -> TaskArtifactUpdateEvent:
148+
"""Signal that the artifact stream is complete.
149+
150+
Returns:
151+
A ``TaskArtifactUpdateEvent`` with ``append=True`` and
152+
``last_chunk=True``.
153+
"""
154+
return TaskArtifactUpdateEvent(
155+
context_id=self._context_id,
156+
task_id=self._task_id,
157+
append=True,
158+
last_chunk=True,
159+
artifact=Artifact(
160+
artifact_id=self._artifact_id,
161+
name=self._name,
162+
parts=[],
163+
)
164+
)

tests/utils/test_artifact.py

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@
77
Artifact,
88
DataPart,
99
Part,
10+
TaskArtifactUpdateEvent,
1011
TextPart,
1112
)
1213
from a2a.utils.artifact import (
14+
ArtifactStreamer,
1315
get_artifact_text,
1416
new_artifact,
1517
new_data_artifact,
@@ -155,5 +157,105 @@ def test_get_artifact_text_empty_parts(self):
155157
assert result == ''
156158

157159

160+
class TestArtifactStreamer(unittest.TestCase):
161+
def setUp(self):
162+
self.context_id = 'ctx-123'
163+
self.task_id = 'task-456'
164+
165+
def test_generates_stable_artifact_id(self):
166+
streamer = ArtifactStreamer(self.context_id, self.task_id)
167+
e1 = streamer.append('hello ')
168+
e2 = streamer.append('world')
169+
self.assertEqual(e1.artifact.artifact_id, e2.artifact.artifact_id)
170+
171+
def test_uses_explicit_artifact_id(self):
172+
streamer = ArtifactStreamer(
173+
self.context_id, self.task_id, artifact_id='my-fixed-id'
174+
)
175+
event = streamer.append('chunk')
176+
self.assertEqual(event.artifact.artifact_id, 'my-fixed-id')
177+
178+
@patch('a2a.utils.artifact.uuid.uuid4')
179+
def test_generated_id_comes_from_uuid4(self, mock_uuid4):
180+
mock_uuid = uuid.UUID('abcdef12-1234-5678-1234-567812345678')
181+
mock_uuid4.return_value = mock_uuid
182+
streamer = ArtifactStreamer(self.context_id, self.task_id)
183+
self.assertEqual(streamer._artifact_id, str(mock_uuid))
184+
185+
def test_default_name_is_response(self):
186+
streamer = ArtifactStreamer(self.context_id, self.task_id)
187+
event = streamer.append('text')
188+
self.assertEqual(event.artifact.name, 'response')
189+
190+
def test_custom_name(self):
191+
streamer = ArtifactStreamer(
192+
self.context_id, self.task_id, name='summary'
193+
)
194+
event = streamer.append('text')
195+
self.assertEqual(event.artifact.name, 'summary')
196+
197+
def test_append_returns_task_artifact_update_event(self):
198+
streamer = ArtifactStreamer(self.context_id, self.task_id)
199+
event = streamer.append('chunk')
200+
self.assertIsInstance(event, TaskArtifactUpdateEvent)
201+
202+
def test_append_sets_correct_context_and_task(self):
203+
streamer = ArtifactStreamer(self.context_id, self.task_id)
204+
event = streamer.append('chunk')
205+
self.assertEqual(event.context_id, self.context_id)
206+
self.assertEqual(event.task_id, self.task_id)
207+
208+
def test_append_sets_append_true_last_chunk_false(self):
209+
streamer = ArtifactStreamer(self.context_id, self.task_id)
210+
event = streamer.append('chunk')
211+
self.assertTrue(event.append)
212+
self.assertFalse(event.last_chunk)
213+
214+
def test_append_creates_single_text_part(self):
215+
streamer = ArtifactStreamer(self.context_id, self.task_id)
216+
event = streamer.append('hello')
217+
self.assertEqual(len(event.artifact.parts), 1)
218+
self.assertIsInstance(event.artifact.parts[0].root, TextPart)
219+
self.assertEqual(event.artifact.parts[0].root.text, 'hello')
220+
221+
def test_finalize_returns_task_artifact_update_event(self):
222+
streamer = ArtifactStreamer(self.context_id, self.task_id)
223+
event = streamer.finalize()
224+
self.assertIsInstance(event, TaskArtifactUpdateEvent)
225+
226+
def test_finalize_sets_append_true_last_chunk_true(self):
227+
streamer = ArtifactStreamer(self.context_id, self.task_id)
228+
event = streamer.finalize()
229+
self.assertTrue(event.append)
230+
self.assertTrue(event.last_chunk)
231+
232+
def test_finalize_has_empty_parts(self):
233+
streamer = ArtifactStreamer(self.context_id, self.task_id)
234+
event = streamer.finalize()
235+
self.assertEqual(event.artifact.parts, [])
236+
237+
def test_finalize_uses_same_artifact_id_as_append(self):
238+
streamer = ArtifactStreamer(self.context_id, self.task_id)
239+
append_event = streamer.append('text')
240+
finalize_event = streamer.finalize()
241+
self.assertEqual(
242+
append_event.artifact.artifact_id,
243+
finalize_event.artifact.artifact_id,
244+
)
245+
246+
def test_multiple_appends_all_share_artifact_id(self):
247+
streamer = ArtifactStreamer(self.context_id, self.task_id)
248+
events = [streamer.append(f'chunk-{i}') for i in range(5)]
249+
ids = {e.artifact.artifact_id for e in events}
250+
self.assertEqual(len(ids), 1)
251+
252+
def test_multiple_appends_carry_distinct_text(self):
253+
streamer = ArtifactStreamer(self.context_id, self.task_id)
254+
texts = ['Hello, ', 'world', '!']
255+
events = [streamer.append(t) for t in texts]
256+
result_texts = [e.artifact.parts[0].root.text for e in events]
257+
self.assertEqual(result_texts, texts)
258+
259+
158260
if __name__ == '__main__':
159261
unittest.main()

0 commit comments

Comments
 (0)