fix: Let part converters also return multiple parts so they can support more usecases

PiperOrigin-RevId: 830882000
This commit is contained in:
Google Team Member
2025-11-11 06:15:00 -08:00
committed by Copybara-Service
parent fd33610e96
commit 824ab07212
7 changed files with 188 additions and 19 deletions
@@ -301,13 +301,15 @@ def convert_a2a_message_to_event(
) )
try: try:
parts = [] output_parts = []
long_running_tool_ids = set() long_running_tool_ids = set()
for a2a_part in a2a_message.parts: for a2a_part in a2a_message.parts:
try: try:
part = part_converter(a2a_part) parts = part_converter(a2a_part)
if part is None: if not isinstance(parts, list):
parts = [parts] if parts else []
if not parts:
logger.warning("Failed to convert A2A part, skipping: %s", a2a_part) logger.warning("Failed to convert A2A part, skipping: %s", a2a_part)
continue continue
@@ -321,16 +323,18 @@ def convert_a2a_message_to_event(
) )
is True is True
): ):
long_running_tool_ids.add(part.function_call.id) for part in parts:
if part.function_call:
long_running_tool_ids.add(part.function_call.id)
parts.append(part) output_parts.extend(parts)
except Exception as e: except Exception as e:
logger.error("Failed to convert A2A part: %s, error: %s", a2a_part, e) logger.error("Failed to convert A2A part: %s, error: %s", a2a_part, e)
# Continue processing other parts instead of failing completely # Continue processing other parts instead of failing completely
continue continue
if not parts: if not output_parts:
logger.warning( logger.warning(
"No parts could be converted from A2A message %s", a2a_message "No parts could be converted from A2A message %s", a2a_message
) )
@@ -348,7 +352,7 @@ def convert_a2a_message_to_event(
else None, else None,
content=genai_types.Content( content=genai_types.Content(
role="model", role="model",
parts=parts, parts=output_parts,
), ),
) )
@@ -387,15 +391,19 @@ def convert_event_to_a2a_message(
return None return None
try: try:
a2a_parts = [] output_parts = []
for part in event.content.parts: for part in event.content.parts:
a2a_part = part_converter(part) a2a_parts = part_converter(part)
if a2a_part: if not isinstance(a2a_parts, list):
a2a_parts.append(a2a_part) a2a_parts = [a2a_parts] if a2a_parts else []
for a2a_part in a2a_parts:
output_parts.append(a2a_part)
_process_long_running_tool(a2a_part, event) _process_long_running_tool(a2a_part, event)
if a2a_parts: if output_parts:
return Message(message_id=str(uuid.uuid4()), role=role, parts=a2a_parts) return Message(
message_id=str(uuid.uuid4()), role=role, parts=output_parts
)
except Exception as e: except Exception as e:
logger.error("Failed to convert event to status message: %s", e) logger.error("Failed to convert event to status message: %s", e)
@@ -22,7 +22,9 @@ import base64
from collections.abc import Callable from collections.abc import Callable
import json import json
import logging import logging
from typing import List
from typing import Optional from typing import Optional
from typing import Union
from .utils import _get_adk_metadata_key from .utils import _get_adk_metadata_key
@@ -53,10 +55,11 @@ A2A_DATA_PART_METADATA_TYPE_EXECUTABLE_CODE = 'executable_code'
A2APartToGenAIPartConverter = Callable[ A2APartToGenAIPartConverter = Callable[
[a2a_types.Part], Optional[genai_types.Part] [a2a_types.Part], Union[Optional[genai_types.Part], List[genai_types.Part]]
] ]
GenAIPartToA2APartConverter = Callable[ GenAIPartToA2APartConverter = Callable[
[genai_types.Part], Optional[a2a_types.Part] [genai_types.Part],
Union[Optional[a2a_types.Part], List[a2a_types.Part]],
] ]
@@ -110,12 +110,19 @@ def convert_a2a_request_to_agent_run_request(
if request.metadata: if request.metadata:
custom_metadata['a2a_metadata'] = request.metadata custom_metadata['a2a_metadata'] = request.metadata
output_parts = []
for a2a_part in request.message.parts:
genai_parts = part_converter(a2a_part)
if not isinstance(genai_parts, list):
genai_parts = [genai_parts] if genai_parts else []
output_parts.extend(genai_parts)
return AgentRunRequest( return AgentRunRequest(
user_id=_get_user_id(request), user_id=_get_user_id(request),
session_id=request.context_id, session_id=request.context_id,
new_message=genai_types.Content( new_message=genai_types.Content(
role='user', role='user',
parts=[part_converter(part) for part in request.message.parts], parts=output_parts,
), ),
run_config=RunConfig(custom_metadata=custom_metadata), run_config=RunConfig(custom_metadata=custom_metadata),
) )
+6 -3
View File
@@ -376,9 +376,12 @@ class RemoteA2aAgent(BaseAgent):
continue continue
for part in event.content.parts: for part in event.content.parts:
converted_part = self._genai_part_converter(part) converted_parts = self._genai_part_converter(part)
if converted_part: if not isinstance(converted_parts, list):
message_parts.append(converted_part) converted_parts = [converted_parts] if converted_parts else []
if converted_parts:
message_parts.extend(converted_parts)
else: else:
logger.warning("Failed to convert part to A2A format: %s", part) logger.warning("Failed to convert part to A2A format: %s", part)
@@ -576,6 +576,37 @@ class TestEventConverter:
assert result.context_id == context_id assert result.context_id == context_id
assert result.status.state == TaskState.input_required assert result.status.state == TaskState.input_required
def test_convert_event_to_a2a_message_with_multiple_parts_returned(self):
"""Test event to message conversion when part_converter returns multiple parts."""
from a2a import types as a2a_types
from google.adk.a2a.converters.event_converter import convert_event_to_a2a_message
from google.genai import types as genai_types
# Arrange
mock_genai_part = genai_types.Part(text="source part")
mock_a2a_part1 = a2a_types.Part(root=a2a_types.TextPart(text="part 1"))
mock_a2a_part2 = a2a_types.Part(root=a2a_types.TextPart(text="part 2"))
mock_convert_part = Mock()
mock_convert_part.return_value = [mock_a2a_part1, mock_a2a_part2]
self.mock_event.content = genai_types.Content(
parts=[mock_genai_part], role="model"
)
# Act
result = convert_event_to_a2a_message(
self.mock_event,
self.mock_invocation_context,
part_converter=mock_convert_part,
)
# Assert
assert result is not None
assert len(result.parts) == 2
assert result.parts[0].root.text == "part 1"
assert result.parts[1].root.text == "part 2"
mock_convert_part.assert_called_once_with(mock_genai_part)
class TestA2AToEventConverters: class TestA2AToEventConverters:
"""Test suite for A2A to Event conversion functions.""" """Test suite for A2A to Event conversion functions."""
@@ -801,6 +832,36 @@ class TestA2AToEventConverters:
assert result.content.parts[0].text == "test content" assert result.content.parts[0].text == "test content"
mock_convert_part.assert_called_once_with(mock_a2a_part) mock_convert_part.assert_called_once_with(mock_a2a_part)
def test_convert_a2a_message_to_event_with_multiple_parts_returned(self):
"""Test message to event conversion when part_converter returns multiple parts."""
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
from google.genai import types as genai_types
# Arrange
mock_a2a_part = Mock()
mock_genai_part1 = genai_types.Part(text="part 1")
mock_genai_part2 = genai_types.Part(text="part 2")
mock_convert_part = Mock()
mock_convert_part.return_value = [mock_genai_part1, mock_genai_part2]
mock_message = Mock(spec=Message)
mock_message.parts = [mock_a2a_part]
# Act
result = convert_a2a_message_to_event(
mock_message,
"test-author",
self.mock_invocation_context,
mock_convert_part,
)
# Assert
assert result.content.role == "model"
assert len(result.content.parts) == 2
assert result.content.parts[0].text == "part 1"
assert result.content.parts[1].text == "part 2"
mock_convert_part.assert_called_once_with(mock_a2a_part)
def test_convert_a2a_message_to_event_with_long_running_tools(self): def test_convert_a2a_message_to_event_with_long_running_tools(self):
"""Test conversion with long-running tools by mocking the entire flow.""" """Test conversion with long-running tools by mocking the entire flow."""
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
@@ -195,6 +195,58 @@ class TestConvertA2aRequestToAgentRunRequest:
mock_convert_part.assert_any_call(mock_part1) mock_convert_part.assert_any_call(mock_part1)
mock_convert_part.assert_any_call(mock_part2) mock_convert_part.assert_any_call(mock_part2)
def test_convert_a2a_request_multiple_parts(self):
"""Test basic conversion of A2A request to ADK AgentRunRequest."""
# Arrange
mock_part1 = Mock()
mock_part2 = Mock()
mock_message = Mock()
mock_message.parts = [mock_part1, mock_part2]
mock_user = Mock()
mock_user.user_name = "test_user"
mock_call_context = Mock()
mock_call_context.user = mock_user
request = Mock(spec=RequestContext)
request.message = mock_message
request.context_id = "test_context_123"
request.call_context = mock_call_context
request.metadata = {"test_key": "test_value"}
# Create proper genai_types.Part objects instead of mocks
mock_genai_part1 = genai_types.Part(text="test part 1")
mock_genai_part2 = genai_types.Part(text="test part 2")
mock_convert_part = Mock()
mock_convert_part.side_effect = [mock_genai_part1, mock_genai_part2]
# Act
result = convert_a2a_request_to_agent_run_request(
request, mock_convert_part
)
# Assert
assert result is not None
assert result.user_id == "test_user"
assert result.session_id == "test_context_123"
assert isinstance(result.new_message, genai_types.Content)
assert result.new_message.role == "user"
assert result.new_message.parts == [
mock_genai_part1,
mock_genai_part2,
]
assert isinstance(result.run_config, RunConfig)
assert result.run_config.custom_metadata == {
"a2a_metadata": {"test_key": "test_value"}
}
# Verify calls
assert mock_convert_part.call_count == 2
mock_convert_part.assert_any_call(mock_part1)
mock_convert_part.assert_any_call(mock_part2)
def test_convert_a2a_request_no_message_raises_error(self): def test_convert_a2a_request_no_message_raises_error(self):
"""Test that conversion raises ValueError when message is None.""" """Test that conversion raises ValueError when message is None."""
# Arrange # Arrange
@@ -653,6 +653,41 @@ class TestRemoteA2aAgentMessageHandling:
assert result[0][0] == mock_a2a_part assert result[0][0] == mock_a2a_part
assert result[1] is None # context_id assert result[1] is None # context_id
def test_construct_message_parts_from_session_success_multiple_parts(self):
"""Test successful message parts construction from session."""
# Mock event with text content
mock_part = Mock()
mock_part.text = "Hello world"
mock_content = Mock()
mock_content.parts = [mock_part]
mock_event = Mock()
mock_event.content = mock_content
self.mock_session.events = [mock_event]
with patch(
"google.adk.agents.remote_a2a_agent._present_other_agent_message"
) as mock_convert:
mock_convert.return_value = mock_event
mock_a2a_part1 = Mock()
mock_a2a_part2 = Mock()
self.mock_genai_part_converter.return_value = [
mock_a2a_part1,
mock_a2a_part2,
]
result = self.agent._construct_message_parts_from_session(
self.mock_context
)
assert len(result) == 2 # Returns tuple of (parts, context_id)
assert len(result[0]) == 2 # parts list
assert result[0] == [mock_a2a_part1, mock_a2a_part2]
assert result[1] is None # context_id
def test_construct_message_parts_from_session_empty_events(self): def test_construct_message_parts_from_session_empty_events(self):
"""Test message parts construction with empty events.""" """Test message parts construction with empty events."""
self.mock_session.events = [] self.mock_session.events = []