You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
fix: fix bug where remote a2a agent wasn't using its a2a part converter
PiperOrigin-RevId: 836399603
This commit is contained in:
committed by
Copybara-Service
parent
a1c09b724b
commit
4eb2a11403
@@ -417,7 +417,9 @@ class RemoteA2aAgent(BaseAgent):
|
||||
# This is the initial response for a streaming task or the complete
|
||||
# response for a non-streaming task, which is the full task state.
|
||||
# We process this to get the initial message.
|
||||
event = convert_a2a_task_to_event(task, self.name, ctx)
|
||||
event = convert_a2a_task_to_event(
|
||||
task, self.name, ctx, self._a2a_part_converter
|
||||
)
|
||||
# for streaming task, we update the event with the task status.
|
||||
# We update the event as Thought updates.
|
||||
if task and task.status and task.status.state == TaskState.submitted:
|
||||
@@ -429,7 +431,7 @@ class RemoteA2aAgent(BaseAgent):
|
||||
):
|
||||
# This is a streaming task status update with a message.
|
||||
event = convert_a2a_message_to_event(
|
||||
update.status.message, self.name, ctx
|
||||
update.status.message, self.name, ctx, self._a2a_part_converter
|
||||
)
|
||||
if event.content and update.status.state in [
|
||||
TaskState.submitted,
|
||||
@@ -447,7 +449,9 @@ class RemoteA2aAgent(BaseAgent):
|
||||
# signals:
|
||||
# 1. append: True for partial updates, False for full updates.
|
||||
# 2. last_chunk: True for full updates, False for partial updates.
|
||||
event = convert_a2a_task_to_event(task, self.name, ctx)
|
||||
event = convert_a2a_task_to_event(
|
||||
task, self.name, ctx, self._a2a_part_converter
|
||||
)
|
||||
else:
|
||||
# This is a streaming update without a message (e.g. status change)
|
||||
# or a partial artifact update. We don't emit an event for these
|
||||
@@ -463,7 +467,9 @@ class RemoteA2aAgent(BaseAgent):
|
||||
|
||||
# Otherwise, it's a regular A2AMessage for non-streaming responses.
|
||||
elif isinstance(a2a_response, A2AMessage):
|
||||
event = convert_a2a_message_to_event(a2a_response, self.name, ctx)
|
||||
event = convert_a2a_message_to_event(
|
||||
a2a_response, self.name, ctx, self._a2a_part_converter
|
||||
)
|
||||
event.custom_metadata = event.custom_metadata or {}
|
||||
|
||||
if a2a_response.context_id:
|
||||
|
||||
@@ -723,6 +723,7 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -760,6 +761,7 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check the parts are not updated as Thought
|
||||
assert result.content.parts[0].thought is None
|
||||
@@ -864,6 +866,7 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check the parts are updated as Thought
|
||||
assert result.content.parts[0].thought is True
|
||||
@@ -909,6 +912,7 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -954,6 +958,7 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -1009,7 +1014,10 @@ class TestRemoteA2aAgentMessageHandling:
|
||||
|
||||
assert result == mock_event
|
||||
mock_convert.assert_called_once_with(
|
||||
mock_a2a_task, self.agent.name, self.mock_context
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.agent._a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -1039,6 +1047,8 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
|
||||
def setup_method(self):
|
||||
"""Setup test fixtures."""
|
||||
self.mock_a2a_part_converter = Mock()
|
||||
|
||||
self.agent_card = create_test_agent_card()
|
||||
self.agent = RemoteA2aAgent(
|
||||
name="test_agent",
|
||||
@@ -1046,6 +1056,7 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
a2a_client_factory=ClientFactory(
|
||||
config=ClientConfig(httpx_client=httpx.AsyncClient()),
|
||||
),
|
||||
a2a_part_converter=self.mock_a2a_part_converter,
|
||||
)
|
||||
|
||||
# Mock session and context
|
||||
@@ -1173,7 +1184,10 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
|
||||
assert result == mock_event
|
||||
mock_convert.assert_called_once_with(
|
||||
mock_a2a_message, self.agent.name, self.mock_context
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -1211,6 +1225,7 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_a2a_part_converter,
|
||||
)
|
||||
# Check the parts are not updated as Thought
|
||||
assert result.content.parts[0].thought is None
|
||||
@@ -1251,6 +1266,7 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.agent._a2a_part_converter,
|
||||
)
|
||||
# Check the parts are updated as Thought
|
||||
assert result.content.parts[0].thought is True
|
||||
@@ -1296,6 +1312,7 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.agent._a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -1341,6 +1358,7 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.agent._a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
@@ -1396,7 +1414,10 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
|
||||
assert result == mock_event
|
||||
mock_convert.assert_called_once_with(
|
||||
mock_a2a_task, self.agent.name, self.mock_context
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.agent._a2a_part_converter,
|
||||
)
|
||||
# Check that metadata was added
|
||||
assert result.custom_metadata is not None
|
||||
|
||||
Reference in New Issue
Block a user