Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/a2a/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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',
Expand Down
30 changes: 30 additions & 0 deletions src/a2a/utils/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,13 @@

import uuid

from typing import Any

from a2a.types import (
DataPart,
FilePart,
FileWithBytes,
FileWithUri,
Message,
Part,
Role,
Expand Down Expand Up @@ -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.

Expand Down
140 changes: 140 additions & 0 deletions tests/utils/test_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -178,6 +183,141 @@ 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'
Comment thread
holtskinner marked this conversation as resolved.
)
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
Expand Down
Loading