mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Allow custom part converters in A2A classes
This change introduces type descriptions for the functions which convert between A2A and GenAI `Part`s. It then allows passing instances of those functions to the various A2A-related functions/classes, effectively allowing users to inject their own logic for how part conversion should occur. The benefit of this pattern is that users can create decorators around the core `Part` conversion logic, which allows them to intercept the cases they care about while delegating the ones they do not to the core converter. This is a pattern we use a lot in the A2A Python SDK. One example where this type of logic is useful is for extensions: this allows extension logic to, for example, interpret an A2A DataPart into a FunctionResponse using extension-specific logic. PiperOrigin-RevId: 803186799
This commit is contained in:
committed by
Copybara-Service
parent
4df79dd5c9
commit
b05fef9ba7
@@ -652,6 +652,8 @@ class TestA2AToEventConverters:
|
||||
with patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_message_to_event"
|
||||
) as mock_convert_message:
|
||||
from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part
|
||||
|
||||
mock_event = Mock(spec=Event)
|
||||
mock_convert_message.return_value = mock_event
|
||||
|
||||
@@ -662,7 +664,10 @@ class TestA2AToEventConverters:
|
||||
assert result == mock_event
|
||||
# Should call convert_a2a_message_to_event with the status message
|
||||
mock_convert_message.assert_called_once_with(
|
||||
mock_status.message, "test-author", self.mock_invocation_context
|
||||
mock_status.message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
part_converter=convert_a2a_part_to_genai_part,
|
||||
)
|
||||
|
||||
def test_convert_a2a_task_to_event_with_history_message(self):
|
||||
@@ -680,6 +685,8 @@ class TestA2AToEventConverters:
|
||||
with patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_message_to_event"
|
||||
) as mock_convert_message:
|
||||
from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part
|
||||
|
||||
mock_event = Mock(spec=Event)
|
||||
mock_event.invocation_id = "test-invocation-id"
|
||||
mock_convert_message.return_value = mock_event
|
||||
@@ -688,7 +695,10 @@ class TestA2AToEventConverters:
|
||||
|
||||
# Verify the message converter was called with correct parameters
|
||||
mock_convert_message.assert_called_once_with(
|
||||
mock_message, "test-author", None
|
||||
mock_message,
|
||||
"test-author",
|
||||
None,
|
||||
part_converter=convert_a2a_part_to_genai_part,
|
||||
)
|
||||
assert result == mock_event
|
||||
|
||||
@@ -761,10 +771,7 @@ class TestA2AToEventConverters:
|
||||
with pytest.raises(RuntimeError, match="Failed to convert task message"):
|
||||
convert_a2a_task_to_event(mock_task, "test-author")
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_message_to_event_success(self, mock_convert_part):
|
||||
def test_convert_a2a_message_to_event_success(self):
|
||||
"""Test successful conversion of A2A message to event."""
|
||||
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
|
||||
from google.genai import types as genai_types
|
||||
@@ -772,13 +779,17 @@ class TestA2AToEventConverters:
|
||||
# Create mock parts and message with valid genai Part
|
||||
mock_a2a_part = Mock()
|
||||
mock_genai_part = genai_types.Part(text="test content")
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = mock_genai_part
|
||||
|
||||
mock_message = Mock(spec=Message)
|
||||
mock_message.parts = [mock_a2a_part]
|
||||
|
||||
result = convert_a2a_message_to_event(
|
||||
mock_message, "test-author", self.mock_invocation_context
|
||||
mock_message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
mock_convert_part,
|
||||
)
|
||||
|
||||
# Verify conversion was successful
|
||||
@@ -790,12 +801,7 @@ class TestA2AToEventConverters:
|
||||
assert result.content.parts[0].text == "test content"
|
||||
mock_convert_part.assert_called_once_with(mock_a2a_part)
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_message_to_event_with_long_running_tools(
|
||||
self, mock_convert_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
|
||||
|
||||
@@ -805,6 +811,7 @@ class TestA2AToEventConverters:
|
||||
mock_message.parts = [mock_a2a_part]
|
||||
|
||||
# Mock the part conversion to return None to simulate long-running tool detection logic
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = None
|
||||
|
||||
# Patch the long-running tool detection since the main logic is in the actual conversion
|
||||
@@ -812,7 +819,10 @@ class TestA2AToEventConverters:
|
||||
"google.adk.a2a.converters.event_converter.logger"
|
||||
) as mock_logger:
|
||||
result = convert_a2a_message_to_event(
|
||||
mock_message, "test-author", self.mock_invocation_context
|
||||
mock_message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
mock_convert_part,
|
||||
)
|
||||
|
||||
# Verify basic conversion worked
|
||||
@@ -845,24 +855,23 @@ class TestA2AToEventConverters:
|
||||
with pytest.raises(ValueError, match="A2A message cannot be None"):
|
||||
convert_a2a_message_to_event(None)
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_message_to_event_part_conversion_fails(
|
||||
self, mock_convert_part
|
||||
):
|
||||
def test_convert_a2a_message_to_event_part_conversion_fails(self):
|
||||
"""Test handling when part conversion returns None."""
|
||||
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
|
||||
|
||||
# Setup mock to return None (conversion failure)
|
||||
mock_a2a_part = Mock()
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = None
|
||||
|
||||
mock_message = Mock(spec=Message)
|
||||
mock_message.parts = [mock_a2a_part]
|
||||
|
||||
result = convert_a2a_message_to_event(
|
||||
mock_message, "test-author", self.mock_invocation_context
|
||||
mock_message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
mock_convert_part,
|
||||
)
|
||||
|
||||
# Verify event was created but with no parts
|
||||
@@ -871,12 +880,7 @@ class TestA2AToEventConverters:
|
||||
assert result.content.role == "model"
|
||||
assert len(result.content.parts) == 0
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_message_to_event_part_conversion_exception(
|
||||
self, mock_convert_part
|
||||
):
|
||||
def test_convert_a2a_message_to_event_part_conversion_exception(self):
|
||||
"""Test handling when part conversion raises exception."""
|
||||
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
|
||||
from google.genai import types as genai_types
|
||||
@@ -886,6 +890,7 @@ class TestA2AToEventConverters:
|
||||
mock_a2a_part2 = Mock()
|
||||
mock_genai_part = genai_types.Part(text="successful conversion")
|
||||
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.side_effect = [
|
||||
Exception("Conversion failed"), # First part fails
|
||||
mock_genai_part, # Second part succeeds
|
||||
@@ -895,7 +900,10 @@ class TestA2AToEventConverters:
|
||||
mock_message.parts = [mock_a2a_part1, mock_a2a_part2]
|
||||
|
||||
result = convert_a2a_message_to_event(
|
||||
mock_message, "test-author", self.mock_invocation_context
|
||||
mock_message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
mock_convert_part,
|
||||
)
|
||||
|
||||
# Verify event was created with only the successfully converted part
|
||||
@@ -905,12 +913,7 @@ class TestA2AToEventConverters:
|
||||
assert len(result.content.parts) == 1
|
||||
assert result.content.parts[0].text == "successful conversion"
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.event_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_message_to_event_missing_tool_id(
|
||||
self, mock_convert_part
|
||||
):
|
||||
def test_convert_a2a_message_to_event_missing_tool_id(self):
|
||||
"""Test handling of message conversion when part conversion fails."""
|
||||
from google.adk.a2a.converters.event_converter import convert_a2a_message_to_event
|
||||
|
||||
@@ -920,10 +923,14 @@ class TestA2AToEventConverters:
|
||||
mock_message.parts = [mock_a2a_part]
|
||||
|
||||
# Mock the part conversion to return None
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = None
|
||||
|
||||
result = convert_a2a_message_to_event(
|
||||
mock_message, "test-author", self.mock_invocation_context
|
||||
mock_message,
|
||||
"test-author",
|
||||
self.mock_invocation_context,
|
||||
mock_convert_part,
|
||||
)
|
||||
|
||||
# Verify basic conversion worked
|
||||
|
||||
@@ -146,10 +146,7 @@ class TestGetUserId:
|
||||
class TestConvertA2aRequestToAdkRunArgs:
|
||||
"""Test cases for convert_a2a_request_to_adk_run_args function."""
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_request_basic(self, mock_convert_part):
|
||||
def test_convert_a2a_request_basic(self):
|
||||
"""Test basic conversion of A2A request to ADK run args."""
|
||||
# Arrange
|
||||
mock_part1 = Mock()
|
||||
@@ -172,10 +169,11 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
# 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_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
@@ -201,14 +199,12 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
with pytest.raises(ValueError, match="Request message cannot be None"):
|
||||
convert_a2a_request_to_adk_run_args(request)
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_request_empty_parts(self, mock_convert_part):
|
||||
def test_convert_a2a_request_empty_parts(self):
|
||||
"""Test conversion with empty parts list."""
|
||||
# Arrange
|
||||
mock_message = Mock()
|
||||
mock_message.parts = []
|
||||
mock_convert_part = Mock()
|
||||
|
||||
request = Mock(spec=RequestContext)
|
||||
request.message = mock_message
|
||||
@@ -216,7 +212,7 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
request.call_context = None
|
||||
|
||||
# Act
|
||||
result = convert_a2a_request_to_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
@@ -230,10 +226,7 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
# Verify convert_part wasn't called
|
||||
mock_convert_part.assert_not_called()
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_request_none_context_id(self, mock_convert_part):
|
||||
def test_convert_a2a_request_none_context_id(self):
|
||||
"""Test conversion when context_id is None."""
|
||||
# Arrange
|
||||
mock_part = Mock()
|
||||
@@ -247,10 +240,11 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
|
||||
# Create proper genai_types.Part object instead of mock
|
||||
mock_genai_part = genai_types.Part(text="test part")
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = mock_genai_part
|
||||
|
||||
# Act
|
||||
result = convert_a2a_request_to_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
@@ -261,10 +255,7 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
assert result["new_message"].parts == [mock_genai_part]
|
||||
assert isinstance(result["run_config"], RunConfig)
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_convert_a2a_request_no_auth(self, mock_convert_part):
|
||||
def test_convert_a2a_request_no_auth(self):
|
||||
"""Test conversion when no authentication is available."""
|
||||
# Arrange
|
||||
mock_part = Mock()
|
||||
@@ -278,10 +269,11 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
|
||||
# Create proper genai_types.Part object instead of mock
|
||||
mock_genai_part = genai_types.Part(text="test part")
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = mock_genai_part
|
||||
|
||||
# Act
|
||||
result = convert_a2a_request_to_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
@@ -296,10 +288,7 @@ class TestConvertA2aRequestToAdkRunArgs:
|
||||
class TestIntegration:
|
||||
"""Integration test cases combining both functions."""
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_end_to_end_conversion_with_auth_user(self, mock_convert_part):
|
||||
def test_end_to_end_conversion_with_auth_user(self):
|
||||
"""Test end-to-end conversion with authenticated user."""
|
||||
# Arrange
|
||||
mock_user = Mock()
|
||||
@@ -319,10 +308,11 @@ class TestIntegration:
|
||||
|
||||
# Create proper genai_types.Part object instead of mock
|
||||
mock_genai_part = genai_types.Part(text="test part")
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = mock_genai_part
|
||||
|
||||
# Act
|
||||
result = convert_a2a_request_to_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
@@ -333,10 +323,7 @@ class TestIntegration:
|
||||
assert result["new_message"].parts == [mock_genai_part]
|
||||
assert isinstance(result["run_config"], RunConfig)
|
||||
|
||||
@patch(
|
||||
"google.adk.a2a.converters.request_converter.convert_a2a_part_to_genai_part"
|
||||
)
|
||||
def test_end_to_end_conversion_with_fallback_user(self, mock_convert_part):
|
||||
def test_end_to_end_conversion_with_fallback_user(self):
|
||||
"""Test end-to-end conversion with fallback user ID."""
|
||||
# Arrange
|
||||
mock_part = Mock()
|
||||
@@ -350,10 +337,11 @@ class TestIntegration:
|
||||
|
||||
# Create proper genai_types.Part object instead of mock
|
||||
mock_genai_part = genai_types.Part(text="test part")
|
||||
mock_convert_part = Mock()
|
||||
mock_convert_part.return_value = mock_genai_part
|
||||
|
||||
# Act
|
||||
result = convert_a2a_request_to_adk_run_args(request)
|
||||
result = convert_a2a_request_to_adk_run_args(request, mock_convert_part)
|
||||
|
||||
# Assert
|
||||
assert result is not None
|
||||
|
||||
Reference in New Issue
Block a user