|
3 | 3 | import uuid |
4 | 4 |
|
5 | 5 | 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 |
6 | 10 |
|
7 | 11 | from a2a.types.a2a_pb2 import ( |
8 | 12 | Artifact, |
@@ -57,6 +61,89 @@ def get_message_text(message: Message, delimiter: str = '\n') -> str: |
57 | 61 | return delimiter.join(get_text_parts(message.parts)) |
58 | 62 |
|
59 | 63 |
|
| 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 | + |
60 | 147 | # --- Artifact Helpers --- |
61 | 148 |
|
62 | 149 |
|
@@ -90,6 +177,89 @@ def new_text_artifact( |
90 | 177 | ) |
91 | 178 |
|
92 | 179 |
|
| 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 | + |
93 | 263 | def get_artifact_text(artifact: Artifact, delimiter: str = '\n') -> str: |
94 | 264 | """Extracts and joins all text content from an Artifact's parts.""" |
95 | 265 | return delimiter.join(get_text_parts(artifact.parts)) |
@@ -141,6 +311,62 @@ def new_task( |
141 | 311 | # --- Part Helpers --- |
142 | 312 |
|
143 | 313 |
|
| 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 | + |
144 | 370 | def get_text_parts(parts: Sequence[Part]) -> list[str]: |
145 | 371 | """Extracts text content from all text Parts.""" |
146 | 372 | return [part.text for part in parts if part.HasField('text')] |
|
0 commit comments