From 03e6b3bcc1cd27dba3aa0c53956194871a295a88 Mon Sep 17 00:00:00 2001 From: Holt Skinner Date: Wed, 16 Jul 2025 17:22:14 +0100 Subject: [PATCH 1/3] feat: Add `get_data_parts()` and `get_file_parts()` helper methods --- src/a2a/utils/__init__.py | 4 ++++ src/a2a/utils/message.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/a2a/utils/__init__.py b/src/a2a/utils/__init__.py index c2cc97e0b..06ac11236 100644 --- a/src/a2a/utils/__init__.py +++ b/src/a2a/utils/__init__.py @@ -17,6 +17,8 @@ create_task_obj, ) from a2a.utils.message import ( + get_data_parts, + get_file_parts, get_message_text, get_text_parts, new_agent_parts_message, @@ -37,6 +39,8 @@ 'build_text_artifact', 'completed_task', 'create_task_obj', + 'get_data_parts', + 'get_file_parts', 'get_message_text', 'get_text_parts', 'new_agent_parts_message', diff --git a/src/a2a/utils/message.py b/src/a2a/utils/message.py index 0e08597a5..708d85c21 100644 --- a/src/a2a/utils/message.py +++ b/src/a2a/utils/message.py @@ -2,7 +2,13 @@ import uuid +from typing import Any + from a2a.types import ( + DataPart, + FilePart, + FileWithBytes, + FileWithUri, Message, Part, Role, @@ -70,6 +76,30 @@ def get_text_parts(parts: list[Part]) -> list[str]: return [part.root.text for part in parts if isinstance(part.root, TextPart)] +def get_data_parts(parts: list[Part]) -> list[dict[str, Any]]: + """Extracts dictionary data from all DataPart objects in a list of Parts. + + Args: + parts: A list of `Part` objects. + + Returns: + A list of dictionaries containing the data from any `DataPart` objects found. + """ + return [part.root.data for part in parts if isinstance(part.root, DataPart)] + + +def get_file_parts(parts: list[Part]) -> list[FileWithBytes | FileWithUri]: + """Extracts file data from all FilePart objects in a list of Parts. + + Args: + parts: A list of `Part` objects. + + Returns: + A list of `FileWithBytes` or `FileWithUri` objects containing the file data from any `FilePart` objects found. + """ + return [part.root.file for part in parts if isinstance(part.root, FilePart)] + + def get_message_text(message: Message, delimiter: str = '\n') -> str: """Extracts and joins all text content from a Message's parts. From fd132ae49b5d00bae3b6c58b5a79c8380584071c Mon Sep 17 00:00:00 2001 From: Holt Skinner Date: Wed, 16 Jul 2025 17:22:19 +0100 Subject: [PATCH 2/3] Add Tests --- tests/utils/test_message.py | 133 ++++++++++++++++++++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/tests/utils/test_message.py b/tests/utils/test_message.py index 55a81ec94..aba8d85b0 100644 --- a/tests/utils/test_message.py +++ b/tests/utils/test_message.py @@ -4,12 +4,17 @@ from a2a.types import ( DataPart, + FilePart, + FileWithBytes, + FileWithUri, Message, Part, Role, TextPart, ) from a2a.utils.message import ( + get_data_parts, + get_file_parts, get_message_text, get_text_parts, new_agent_parts_message, @@ -178,6 +183,134 @@ def test_get_text_parts_empty_list(self): assert result == [] +class TestGetDataParts: + def test_get_data_parts_single_data_part(self): + # Setup + parts = [Part(root=DataPart(data={'key': 'value'}))] + + # Exercise + result = get_data_parts(parts) + + # Verify + assert result == [{'key': 'value'}] + + def test_get_data_parts_multiple_data_parts(self): + # Setup + parts = [ + Part(root=DataPart(data={'key1': 'value1'})), + Part(root=DataPart(data={'key2': 'value2'})), + ] + + # Exercise + result = get_data_parts(parts) + + # Verify + assert result == [{'key1': 'value1'}, {'key2': 'value2'}] + + def test_get_data_parts_mixed_parts(self): + # Setup + parts = [ + Part(root=TextPart(text='some text')), + Part(root=DataPart(data={'key1': 'value1'})), + Part(root=DataPart(data={'key2': 'value2'})), + ] + + # Exercise + result = get_data_parts(parts) + + # Verify + assert result == [{'key1': 'value1'}, {'key2': 'value2'}] + + def test_get_data_parts_no_data_parts(self): + # Setup + parts = [ + Part(root=TextPart(text='some text')), + ] + + # Exercise + result = get_data_parts(parts) + + # Verify + assert result == [] + + def test_get_data_parts_empty_list(self): + # Setup + parts = [] + + # Exercise + result = get_data_parts(parts) + + # Verify + assert result == [] + + +class TestGetFileParts: + def test_get_file_parts_single_file_part(self): + # Setup + file_with_uri = FileWithUri(uri='file://path/to/file', mimeType='text/plain') + parts = [Part(root=FilePart(file=file_with_uri))] + + # Exercise + result = get_file_parts(parts) + + # Verify + assert result == [file_with_uri] + + def test_get_file_parts_multiple_file_parts(self): + # Setup + file_with_uri1 = FileWithUri(uri='file://path/to/file1', mimeType='text/plain') + file_with_bytes = FileWithBytes( + bytes='ZmlsZSBjb250ZW50', mimeType='application/octet-stream' # 'file content' + ) + parts = [ + Part(root=FilePart(file=file_with_uri1)), + Part(root=FilePart(file=file_with_bytes)), + ] + + # Exercise + result = get_file_parts(parts) + + # Verify + assert result == [file_with_uri1, file_with_bytes] + + def test_get_file_parts_mixed_parts(self): + # Setup + file_with_uri = FileWithUri(uri='file://path/to/file', mimeType='text/plain') + parts = [ + Part(root=TextPart(text='some text')), + Part(root=FilePart(file=file_with_uri)), + ] + + # Exercise + result = get_file_parts(parts) + + # Verify + assert result == [file_with_uri] + + def test_get_file_parts_no_file_parts(self): + # Setup + parts = [ + Part(root=TextPart(text='some text')), + Part(root=DataPart(data={'key': 'value'})), + ] + + # Exercise + result = get_file_parts(parts) + + # Verify + assert result == [] + + def test_get_file_parts_empty_list(self): + # Setup + parts = [] + + # Exercise + result = get_file_parts(parts) + + # Verify + assert result == [] + + class TestGetMessageText: def test_get_message_text_single_part(self): # Setup From e341394f2c35385ba8d9e81e2146f3359cbf34be Mon Sep 17 00:00:00 2001 From: Holt Skinner Date: Wed, 16 Jul 2025 17:22:58 +0100 Subject: [PATCH 3/3] Formatting --- tests/utils/test_message.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/utils/test_message.py b/tests/utils/test_message.py index aba8d85b0..86358b7f5 100644 --- a/tests/utils/test_message.py +++ b/tests/utils/test_message.py @@ -247,7 +247,9 @@ def test_get_data_parts_empty_list(self): class TestGetFileParts: def test_get_file_parts_single_file_part(self): # Setup - file_with_uri = FileWithUri(uri='file://path/to/file', mimeType='text/plain') + file_with_uri = FileWithUri( + uri='file://path/to/file', mimeType='text/plain' + ) parts = [Part(root=FilePart(file=file_with_uri))] # Exercise @@ -258,9 +260,12 @@ def test_get_file_parts_single_file_part(self): def test_get_file_parts_multiple_file_parts(self): # Setup - file_with_uri1 = FileWithUri(uri='file://path/to/file1', mimeType='text/plain') + file_with_uri1 = FileWithUri( + uri='file://path/to/file1', mimeType='text/plain' + ) file_with_bytes = FileWithBytes( - bytes='ZmlsZSBjb250ZW50', mimeType='application/octet-stream' # 'file content' + bytes='ZmlsZSBjb250ZW50', + mimeType='application/octet-stream', # 'file content' ) parts = [ Part(root=FilePart(file=file_with_uri1)), @@ -275,7 +280,9 @@ def test_get_file_parts_multiple_file_parts(self): def test_get_file_parts_mixed_parts(self): # Setup - file_with_uri = FileWithUri(uri='file://path/to/file', mimeType='text/plain') + file_with_uri = FileWithUri( + uri='file://path/to/file', mimeType='text/plain' + ) parts = [ Part(root=TextPart(text='some text')), Part(root=FilePart(file=file_with_uri)),