Skip to content

Commit e2a1f26

Browse files
feat(helpers): add new_data_part, new_raw_part, and new_url_part helpers
Part.data is google.protobuf.Value in the v1.0 spec, which requires a non-obvious ParseDict dance to construct from a plain dict. Add new_data_part() to hide that complexity. Also add new_raw_part() and new_url_part() to cover the remaining non-text Part variants that had no helpers, keeping the API consistent with new_text_message() et al.
1 parent 24db37e commit e2a1f26

2 files changed

Lines changed: 406 additions & 0 deletions

File tree

src/a2a/helpers/proto_helpers.py

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@
33
import uuid
44

55
from collections.abc import Sequence
6+
from typing import Any
7+
8+
from google.protobuf import struct_pb2
9+
from google.protobuf.json_format import ParseDict
610

711
from a2a.types.a2a_pb2 import (
812
Artifact,
@@ -57,6 +61,89 @@ def get_message_text(message: Message, delimiter: str = '\n') -> str:
5761
return delimiter.join(get_text_parts(message.parts))
5862

5963

64+
def new_data_message(
65+
data: Any,
66+
role: Role = Role.ROLE_AGENT,
67+
context_id: str | None = None,
68+
task_id: str | None = None,
69+
) -> Message:
70+
"""Creates a new message containing a single data Part.
71+
72+
Args:
73+
data: JSON-serializable data to embed (dict, list, str, etc.).
74+
role: The role of the message sender (default: ROLE_AGENT).
75+
context_id: Optional context ID.
76+
task_id: Optional task ID.
77+
78+
Returns:
79+
A Message with a single data Part.
80+
"""
81+
return new_message(
82+
parts=[new_data_part(data)],
83+
role=role,
84+
context_id=context_id,
85+
task_id=task_id,
86+
)
87+
88+
89+
def new_raw_message( # noqa: PLR0913
90+
raw: bytes,
91+
media_type: str | None = None,
92+
filename: str | None = None,
93+
role: Role = Role.ROLE_AGENT,
94+
context_id: str | None = None,
95+
task_id: str | None = None,
96+
) -> Message:
97+
"""Creates a new message containing a single raw bytes Part.
98+
99+
Args:
100+
raw: The raw bytes content.
101+
media_type: Optional MIME type (e.g. 'image/png').
102+
filename: Optional filename.
103+
role: The role of the message sender (default: ROLE_AGENT).
104+
context_id: Optional context ID.
105+
task_id: Optional task ID.
106+
107+
Returns:
108+
A Message with a single raw Part.
109+
"""
110+
return new_message(
111+
parts=[new_raw_part(raw, media_type=media_type, filename=filename)],
112+
role=role,
113+
context_id=context_id,
114+
task_id=task_id,
115+
)
116+
117+
118+
def new_url_message( # noqa: PLR0913
119+
url: str,
120+
media_type: str | None = None,
121+
filename: str | None = None,
122+
role: Role = Role.ROLE_AGENT,
123+
context_id: str | None = None,
124+
task_id: str | None = None,
125+
) -> Message:
126+
"""Creates a new message containing a single URL Part.
127+
128+
Args:
129+
url: The URL pointing to the file content.
130+
media_type: Optional MIME type (e.g. 'image/png').
131+
filename: Optional filename.
132+
role: The role of the message sender (default: ROLE_AGENT).
133+
context_id: Optional context ID.
134+
task_id: Optional task ID.
135+
136+
Returns:
137+
A Message with a single URL Part.
138+
"""
139+
return new_message(
140+
parts=[new_url_part(url, media_type=media_type, filename=filename)],
141+
role=role,
142+
context_id=context_id,
143+
task_id=task_id,
144+
)
145+
146+
60147
# --- Artifact Helpers ---
61148

62149

@@ -90,6 +177,89 @@ def new_text_artifact(
90177
)
91178

92179

180+
def new_data_artifact(
181+
name: str,
182+
data: Any,
183+
description: str | None = None,
184+
artifact_id: str | None = None,
185+
) -> Artifact:
186+
"""Creates a new Artifact object containing only a single data Part.
187+
188+
Args:
189+
name: The name of the artifact.
190+
data: JSON-serializable data to embed (dict, list, str, etc.).
191+
description: Optional description.
192+
artifact_id: Optional artifact ID (auto-generated if not provided).
193+
194+
Returns:
195+
An Artifact with a single data Part.
196+
"""
197+
return new_artifact(
198+
[new_data_part(data)],
199+
name,
200+
description,
201+
artifact_id=artifact_id,
202+
)
203+
204+
205+
def new_raw_artifact( # noqa: PLR0913
206+
name: str,
207+
raw: bytes,
208+
media_type: str | None = None,
209+
filename: str | None = None,
210+
description: str | None = None,
211+
artifact_id: str | None = None,
212+
) -> Artifact:
213+
"""Creates a new Artifact object containing only a single raw bytes Part.
214+
215+
Args:
216+
name: The name of the artifact.
217+
raw: The raw bytes content.
218+
media_type: Optional MIME type (e.g. 'image/png').
219+
filename: Optional filename.
220+
description: Optional description.
221+
artifact_id: Optional artifact ID (auto-generated if not provided).
222+
223+
Returns:
224+
An Artifact with a single raw Part.
225+
"""
226+
return new_artifact(
227+
[new_raw_part(raw, media_type=media_type, filename=filename)],
228+
name,
229+
description,
230+
artifact_id=artifact_id,
231+
)
232+
233+
234+
def new_url_artifact( # noqa: PLR0913
235+
name: str,
236+
url: str,
237+
media_type: str | None = None,
238+
filename: str | None = None,
239+
description: str | None = None,
240+
artifact_id: str | None = None,
241+
) -> Artifact:
242+
"""Creates a new Artifact object containing only a single URL Part.
243+
244+
Args:
245+
name: The name of the artifact.
246+
url: The URL pointing to the file content.
247+
media_type: Optional MIME type (e.g. 'image/png').
248+
filename: Optional filename.
249+
description: Optional description.
250+
artifact_id: Optional artifact ID (auto-generated if not provided).
251+
252+
Returns:
253+
An Artifact with a single URL Part.
254+
"""
255+
return new_artifact(
256+
[new_url_part(url, media_type=media_type, filename=filename)],
257+
name,
258+
description,
259+
artifact_id=artifact_id,
260+
)
261+
262+
93263
def get_artifact_text(artifact: Artifact, delimiter: str = '\n') -> str:
94264
"""Extracts and joins all text content from an Artifact's parts."""
95265
return delimiter.join(get_text_parts(artifact.parts))
@@ -141,6 +311,62 @@ def new_task(
141311
# --- Part Helpers ---
142312

143313

314+
def new_data_part(data: Any) -> Part:
315+
"""Creates a Part with structured data (google.protobuf.Value).
316+
317+
Args:
318+
data: JSON-serializable data to embed (dict, list, str, etc.).
319+
320+
Returns:
321+
A Part with the data field set.
322+
"""
323+
return Part(data=ParseDict(data, struct_pb2.Value()))
324+
325+
326+
def new_raw_part(
327+
raw: bytes,
328+
media_type: str | None = None,
329+
filename: str | None = None,
330+
) -> Part:
331+
"""Creates a Part with raw bytes content.
332+
333+
Args:
334+
raw: The raw bytes content.
335+
media_type: Optional MIME type (e.g. 'image/png').
336+
filename: Optional filename.
337+
338+
Returns:
339+
A Part with the raw field set.
340+
"""
341+
return Part(
342+
raw=raw,
343+
media_type=media_type or '',
344+
filename=filename or '',
345+
)
346+
347+
348+
def new_url_part(
349+
url: str,
350+
media_type: str | None = None,
351+
filename: str | None = None,
352+
) -> Part:
353+
"""Creates a Part with a URL pointing to file content.
354+
355+
Args:
356+
url: The URL to the file content.
357+
media_type: Optional MIME type (e.g. 'image/png').
358+
filename: Optional filename.
359+
360+
Returns:
361+
A Part with the url field set.
362+
"""
363+
return Part(
364+
url=url,
365+
media_type=media_type or '',
366+
filename=filename or '',
367+
)
368+
369+
144370
def get_text_parts(parts: Sequence[Part]) -> list[str]:
145371
"""Extracts text content from all text Parts."""
146372
return [part.text for part in parts if part.HasField('text')]

0 commit comments

Comments
 (0)