mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Let part converters also return multiple parts so they can support more usecases
PiperOrigin-RevId: 830882000
This commit is contained in:
committed by
Copybara-Service
parent
fd33610e96
commit
824ab07212
@@ -301,13 +301,15 @@ def convert_a2a_message_to_event(
|
||||
)
|
||||
|
||||
try:
|
||||
parts = []
|
||||
output_parts = []
|
||||
long_running_tool_ids = set()
|
||||
|
||||
for a2a_part in a2a_message.parts:
|
||||
try:
|
||||
part = part_converter(a2a_part)
|
||||
if part is None:
|
||||
parts = part_converter(a2a_part)
|
||||
if not isinstance(parts, list):
|
||||
parts = [parts] if parts else []
|
||||
if not parts:
|
||||
logger.warning("Failed to convert A2A part, skipping: %s", a2a_part)
|
||||
continue
|
||||
|
||||
@@ -321,16 +323,18 @@ def convert_a2a_message_to_event(
|
||||
)
|
||||
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:
|
||||
logger.error("Failed to convert A2A part: %s, error: %s", a2a_part, e)
|
||||
# Continue processing other parts instead of failing completely
|
||||
continue
|
||||
|
||||
if not parts:
|
||||
if not output_parts:
|
||||
logger.warning(
|
||||
"No parts could be converted from A2A message %s", a2a_message
|
||||
)
|
||||
@@ -348,7 +352,7 @@ def convert_a2a_message_to_event(
|
||||
else None,
|
||||
content=genai_types.Content(
|
||||
role="model",
|
||||
parts=parts,
|
||||
parts=output_parts,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -387,15 +391,19 @@ def convert_event_to_a2a_message(
|
||||
return None
|
||||
|
||||
try:
|
||||
a2a_parts = []
|
||||
output_parts = []
|
||||
for part in event.content.parts:
|
||||
a2a_part = part_converter(part)
|
||||
if a2a_part:
|
||||
a2a_parts.append(a2a_part)
|
||||
a2a_parts = part_converter(part)
|
||||
if not isinstance(a2a_parts, list):
|
||||
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)
|
||||
|
||||
if a2a_parts:
|
||||
return Message(message_id=str(uuid.uuid4()), role=role, parts=a2a_parts)
|
||||
if output_parts:
|
||||
return Message(
|
||||
message_id=str(uuid.uuid4()), role=role, parts=output_parts
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to convert event to status message: %s", e)
|
||||
|
||||
@@ -22,7 +22,9 @@ import base64
|
||||
from collections.abc import Callable
|
||||
import json
|
||||
import logging
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
from .utils import _get_adk_metadata_key
|
||||
|
||||
@@ -53,10 +55,11 @@ A2A_DATA_PART_METADATA_TYPE_EXECUTABLE_CODE = 'executable_code'
|
||||
|
||||
|
||||
A2APartToGenAIPartConverter = Callable[
|
||||
[a2a_types.Part], Optional[genai_types.Part]
|
||||
[a2a_types.Part], Union[Optional[genai_types.Part], List[genai_types.Part]]
|
||||
]
|
||||
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:
|
||||
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(
|
||||
user_id=_get_user_id(request),
|
||||
session_id=request.context_id,
|
||||
new_message=genai_types.Content(
|
||||
role='user',
|
||||
parts=[part_converter(part) for part in request.message.parts],
|
||||
parts=output_parts,
|
||||
),
|
||||
run_config=RunConfig(custom_metadata=custom_metadata),
|
||||
)
|
||||
|
||||
@@ -376,9 +376,12 @@ class RemoteA2aAgent(BaseAgent):
|
||||
continue
|
||||
|
||||
for part in event.content.parts:
|
||||
converted_part = self._genai_part_converter(part)
|
||||
if converted_part:
|
||||
message_parts.append(converted_part)
|
||||
converted_parts = self._genai_part_converter(part)
|
||||
if not isinstance(converted_parts, list):
|
||||
converted_parts = [converted_parts] if converted_parts else []
|
||||
|
||||
if converted_parts:
|
||||
message_parts.extend(converted_parts)
|
||||
else:
|
||||
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.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:
|
||||
"""Test suite for A2A to Event conversion functions."""
|
||||
@@ -801,6 +832,36 @@ class TestA2AToEventConverters:
|
||||
assert result.content.parts[0].text == "test content"
|
||||
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):
|
||||
"""Test conversion with long-running tools by mocking the entire flow."""
|
||||
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_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):
|
||||
"""Test that conversion raises ValueError when message is None."""
|
||||
# Arrange
|
||||
|
||||
@@ -653,6 +653,41 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
assert result[0][0] == mock_a2a_part
|
||||
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):
|
||||
"""Test message parts construction with empty events."""
|
||||
self.mock_session.events = []
|
||||
|
||||
Reference in New Issue
Block a user