You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: New implementation of RemoteA2aAgent and A2A-ADK conversion
This change introduces compatibility of the remote_a2a_agent with the new a2a_agent_executor. New Event Converters: `to_adk_event.py`= Defines a new set of default converters for transforming A2A types (Task, Message, TaskStatusUpdateEvent, TaskArtifactUpdateEvent) into ADK Event objects. Configurable Remote Agent: The A2aRemoteAgentConfig object allows users to override the default event converters with custom ones. New AgentExecutor Compatibility: RemoteA2aAgent now checks for `agent_executor_v2` metadata in A2A responses. If detected, it delegates response handling to a new `_handle_a2a_response_impl` method, which utilizes the modular converters defined in the configuration. PiperOrigin-RevId: 878487448
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
82c2eefb27
commit
6770e419f5
@@ -25,9 +25,18 @@ from typing import Union
|
||||
from a2a.client.middleware import ClientCallContext
|
||||
from a2a.server.events import Event as A2AEvent
|
||||
from a2a.types import Message as A2AMessage
|
||||
from a2a.types import MessageSendConfiguration
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ...a2a.converters.part_converter import A2APartToGenAIPartConverter
|
||||
from ...a2a.converters.part_converter import convert_a2a_part_to_genai_part
|
||||
from ...a2a.converters.to_adk_event import A2AArtifactUpdateToEventConverter
|
||||
from ...a2a.converters.to_adk_event import A2AMessageToEventConverter
|
||||
from ...a2a.converters.to_adk_event import A2AStatusUpdateToEventConverter
|
||||
from ...a2a.converters.to_adk_event import A2ATaskToEventConverter
|
||||
from ...a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event
|
||||
from ...a2a.converters.to_adk_event import convert_a2a_message_to_event
|
||||
from ...a2a.converters.to_adk_event import convert_a2a_status_update_to_event
|
||||
from ...a2a.converters.to_adk_event import convert_a2a_task_to_event
|
||||
from ...agents.invocation_context import InvocationContext
|
||||
from ...events.event import Event
|
||||
|
||||
@@ -71,6 +80,31 @@ class RequestInterceptor(BaseModel):
|
||||
|
||||
|
||||
class A2aRemoteAgentConfig(BaseModel):
|
||||
"""Configuration for the RemoteA2aAgent."""
|
||||
"""Configuration for A2A remote agents."""
|
||||
|
||||
# Converts standard A2A Messages into ADK Event.
|
||||
a2a_message_converter: A2AMessageToEventConverter = (
|
||||
convert_a2a_message_to_event
|
||||
)
|
||||
|
||||
# Converts an A2A Task into an ADK Event.
|
||||
a2a_task_converter: A2ATaskToEventConverter = convert_a2a_task_to_event
|
||||
|
||||
# Converts A2A TaskStatusUpdateEvents into ADK Event.
|
||||
a2a_status_update_converter: A2AStatusUpdateToEventConverter = (
|
||||
convert_a2a_status_update_to_event
|
||||
)
|
||||
|
||||
# Converts A2A TaskArtifactUpdateEvents into ADK Event.
|
||||
a2a_artifact_update_converter: A2AArtifactUpdateToEventConverter = (
|
||||
convert_a2a_artifact_update_to_event
|
||||
)
|
||||
|
||||
# A low-level hook that converts individual A2A Message Parts
|
||||
# into native ADK/GenAI Part objects.
|
||||
# This is utilized internally by the other converters.
|
||||
a2a_part_converter: A2APartToGenAIPartConverter = (
|
||||
convert_a2a_part_to_genai_part
|
||||
)
|
||||
|
||||
request_interceptors: Optional[list[RequestInterceptor]] = None
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import logging
|
||||
from typing import Any
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
import uuid
|
||||
|
||||
from a2a.types import Message
|
||||
from a2a.types import Part as A2APart
|
||||
from a2a.types import Task
|
||||
from a2a.types import TaskArtifactUpdateEvent
|
||||
from a2a.types import TaskState
|
||||
from a2a.types import TaskStatusUpdateEvent
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from ...agents.invocation_context import InvocationContext
|
||||
from ...events.event import Event
|
||||
from ..experimental import a2a_experimental
|
||||
from .part_converter import A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY
|
||||
from .part_converter import A2APartToGenAIPartConverter
|
||||
from .part_converter import convert_a2a_part_to_genai_part
|
||||
from .utils import _get_adk_metadata_key
|
||||
|
||||
# Logger
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
A2AMessageToEventConverter = Callable[
|
||||
[
|
||||
Message,
|
||||
Optional[str],
|
||||
Optional[InvocationContext],
|
||||
A2APartToGenAIPartConverter,
|
||||
],
|
||||
Optional[Event],
|
||||
]
|
||||
"""A Callable that converts an A2A Message to an ADK Event.
|
||||
|
||||
Args:
|
||||
Message: The A2A message to convert.
|
||||
Optional[str]: The author of the event.
|
||||
Optional[InvocationContext]: The invocation context.
|
||||
A2APartToGenAIPartConverter: The part converter function.
|
||||
|
||||
Returns:
|
||||
Optional[Event]: The converted ADK Event.
|
||||
"""
|
||||
|
||||
A2ATaskToEventConverter = Callable[
|
||||
[
|
||||
Task,
|
||||
Optional[str],
|
||||
Optional[InvocationContext],
|
||||
A2APartToGenAIPartConverter,
|
||||
],
|
||||
Optional[Event],
|
||||
]
|
||||
"""A Callable that converts an A2A Task to an ADK Event.
|
||||
|
||||
Args:
|
||||
Task: The A2A task to convert.
|
||||
Optional[str]: The author of the event.
|
||||
Optional[InvocationContext]: The invocation context.
|
||||
A2APartToGenAIPartConverter: The part converter function.
|
||||
|
||||
Returns:
|
||||
Optional[Event]: The converted ADK Event.
|
||||
"""
|
||||
|
||||
A2AStatusUpdateToEventConverter = Callable[
|
||||
[
|
||||
TaskStatusUpdateEvent,
|
||||
Optional[str],
|
||||
Optional[InvocationContext],
|
||||
A2APartToGenAIPartConverter,
|
||||
],
|
||||
Optional[Event],
|
||||
]
|
||||
"""A Callable that converts an A2A TaskStatusUpdateEvent to an ADK Event.
|
||||
|
||||
Args:
|
||||
TaskStatusUpdateEvent: The A2A status update event to convert.
|
||||
Optional[str]: The author of the event.
|
||||
Optional[InvocationContext]: The invocation context.
|
||||
A2APartToGenAIPartConverter: The part converter function.
|
||||
|
||||
Returns:
|
||||
Optional[Event]: The converted ADK Event.
|
||||
"""
|
||||
|
||||
A2AArtifactUpdateToEventConverter = Callable[
|
||||
[
|
||||
TaskArtifactUpdateEvent,
|
||||
Optional[str],
|
||||
Optional[InvocationContext],
|
||||
A2APartToGenAIPartConverter,
|
||||
],
|
||||
Optional[Event],
|
||||
]
|
||||
"""A Callable that converts an A2A TaskArtifactUpdateEvent to an ADK Event.
|
||||
|
||||
Args:
|
||||
TaskArtifactUpdateEvent: The A2A artifact update event to convert.
|
||||
Optional[str]: The author of the event.
|
||||
Optional[InvocationContext]: The invocation context.
|
||||
A2APartToGenAIPartConverter: The part converter function.
|
||||
|
||||
Returns:
|
||||
Optional[Event]: The converted ADK Event.
|
||||
"""
|
||||
|
||||
|
||||
def _convert_a2a_parts_to_adk_parts(
|
||||
a2a_parts: List[A2APart],
|
||||
part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part,
|
||||
) -> tuple[List[genai_types.Part], set[str]]:
|
||||
"""Converts a list of A2A parts to a list of ADK parts."""
|
||||
output_parts = []
|
||||
long_running_function_ids = set()
|
||||
|
||||
for a2a_part in a2a_parts:
|
||||
try:
|
||||
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
|
||||
|
||||
# Check for long-running functions
|
||||
if (
|
||||
a2a_part.root.metadata
|
||||
and a2a_part.root.metadata.get(
|
||||
_get_adk_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY)
|
||||
)
|
||||
is True
|
||||
):
|
||||
for part in parts:
|
||||
if part.function_call:
|
||||
long_running_function_ids.add(part.function_call.id)
|
||||
|
||||
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 output_parts:
|
||||
logger.warning("No parts could be converted from A2A message")
|
||||
|
||||
return output_parts, long_running_function_ids
|
||||
|
||||
|
||||
def _create_event(
|
||||
output_parts: List[genai_types.Part],
|
||||
invocation_context: Optional[InvocationContext],
|
||||
author: Optional[str],
|
||||
long_running_function_ids: Optional[set[str]] = None,
|
||||
partial: bool = False,
|
||||
) -> Optional[Event]:
|
||||
"""Creates an ADK event from parts and metadata."""
|
||||
if not output_parts:
|
||||
return None
|
||||
|
||||
event = Event(
|
||||
invocation_id=(
|
||||
invocation_context.invocation_id
|
||||
if invocation_context
|
||||
else str(uuid.uuid4())
|
||||
),
|
||||
author=author or "a2a agent",
|
||||
branch=invocation_context.branch if invocation_context else None,
|
||||
long_running_tool_ids=(
|
||||
long_running_function_ids if long_running_function_ids else None
|
||||
),
|
||||
content=genai_types.Content(
|
||||
role="model",
|
||||
parts=output_parts,
|
||||
),
|
||||
partial=partial,
|
||||
)
|
||||
|
||||
return event
|
||||
|
||||
|
||||
@a2a_experimental
|
||||
def convert_a2a_task_to_event(
|
||||
a2a_task: Task,
|
||||
author: Optional[str] = None,
|
||||
invocation_context: Optional[InvocationContext] = None,
|
||||
part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part,
|
||||
) -> Optional[Event]:
|
||||
"""Converts an A2A task to an ADK event.
|
||||
|
||||
Args:
|
||||
a2a_task: The A2A task to convert. Must not be None.
|
||||
author: The author of the event. Defaults to "a2a agent" if not provided.
|
||||
invocation_context: The invocation context containing session information.
|
||||
If provided, the branch will be set from the context.
|
||||
part_converter: The function to convert A2A part to GenAI part.
|
||||
|
||||
Returns:
|
||||
An ADK Event object representing the converted task.
|
||||
|
||||
Raises:
|
||||
ValueError: If a2a_task is None.
|
||||
RuntimeError: If conversion of the underlying message fails.
|
||||
"""
|
||||
if a2a_task is None:
|
||||
raise ValueError("A2A task cannot be None")
|
||||
|
||||
try:
|
||||
output_parts = []
|
||||
long_running_function_ids = set()
|
||||
if a2a_task.artifacts:
|
||||
artifact_parts = [
|
||||
part for artifact in a2a_task.artifacts for part in artifact.parts
|
||||
]
|
||||
output_parts, _ = _convert_a2a_parts_to_adk_parts(
|
||||
artifact_parts, part_converter
|
||||
)
|
||||
if (
|
||||
a2a_task.status.message
|
||||
and a2a_task.status.state == TaskState.input_required
|
||||
):
|
||||
parts, ids = _convert_a2a_parts_to_adk_parts(
|
||||
a2a_task.status.message.parts, part_converter
|
||||
)
|
||||
output_parts.extend(parts)
|
||||
long_running_function_ids.update(ids)
|
||||
|
||||
return _create_event(
|
||||
output_parts,
|
||||
invocation_context,
|
||||
author,
|
||||
long_running_function_ids,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to convert A2A task to event: %s", e)
|
||||
raise
|
||||
|
||||
|
||||
@a2a_experimental
|
||||
def convert_a2a_message_to_event(
|
||||
a2a_message: Message,
|
||||
author: Optional[str] = None,
|
||||
invocation_context: Optional[InvocationContext] = None,
|
||||
part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part,
|
||||
) -> Optional[Event]:
|
||||
"""Converts an A2A message to an ADK event.
|
||||
|
||||
Args:
|
||||
a2a_message: The A2A message to convert. Must not be None.
|
||||
author: The author of the event. Defaults to "a2a agent" if not provided.
|
||||
invocation_context: The invocation context containing session information.
|
||||
If provided, the branch will be set from the context.
|
||||
part_converter: The function to convert A2A part to GenAI part.
|
||||
|
||||
Returns:
|
||||
An ADK Event object with converted content and long-running function
|
||||
metadata.
|
||||
|
||||
Raises:
|
||||
ValueError: If a2a_message is None.
|
||||
RuntimeError: If conversion of message parts fails.
|
||||
"""
|
||||
if a2a_message is None:
|
||||
raise ValueError("A2A message cannot be None")
|
||||
|
||||
try:
|
||||
output_parts, _ = _convert_a2a_parts_to_adk_parts(
|
||||
a2a_message.parts, part_converter
|
||||
)
|
||||
return _create_event(output_parts, invocation_context, author)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to convert A2A message to event: %s", e)
|
||||
raise RuntimeError(f"Failed to convert message: {e}") from e
|
||||
|
||||
|
||||
@a2a_experimental
|
||||
def convert_a2a_status_update_to_event(
|
||||
a2a_status_update: TaskStatusUpdateEvent,
|
||||
author: Optional[str] = None,
|
||||
invocation_context: Optional[InvocationContext] = None,
|
||||
part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part,
|
||||
) -> Optional[Event]:
|
||||
"""Converts an A2A task status update to an ADK event.
|
||||
|
||||
Args:
|
||||
a2a_status_update: The A2A task status update to convert.
|
||||
author: The author of the event. Defaults to "a2a agent" if not provided.
|
||||
invocation_context: The invocation context containing session information.
|
||||
part_converter: The function to convert A2A part to GenAI part.
|
||||
|
||||
Returns:
|
||||
An ADK Event object representing the converted status update.
|
||||
"""
|
||||
if a2a_status_update is None:
|
||||
raise ValueError("A2A status update cannot be None")
|
||||
|
||||
try:
|
||||
output_parts = []
|
||||
long_running_function_ids = set()
|
||||
if a2a_status_update.status.message:
|
||||
parts, ids = _convert_a2a_parts_to_adk_parts(
|
||||
a2a_status_update.status.message.parts, part_converter
|
||||
)
|
||||
output_parts.extend(parts)
|
||||
long_running_function_ids.update(ids)
|
||||
|
||||
return _create_event(
|
||||
output_parts,
|
||||
invocation_context,
|
||||
author,
|
||||
long_running_function_ids,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Failed to convert A2A status update to event: %s", e)
|
||||
raise RuntimeError(f"Failed to convert status update: {e}") from e
|
||||
|
||||
|
||||
# TODO: Add support for non-ADK Artifact Updates.
|
||||
@a2a_experimental
|
||||
def convert_a2a_artifact_update_to_event(
|
||||
a2a_artifact_update: TaskArtifactUpdateEvent,
|
||||
author: Optional[str] = None,
|
||||
invocation_context: Optional[InvocationContext] = None,
|
||||
part_converter: A2APartToGenAIPartConverter = convert_a2a_part_to_genai_part,
|
||||
) -> Optional[Event]:
|
||||
"""Converts an A2A task artifact update to an ADK event.
|
||||
|
||||
Args:
|
||||
a2a_artifact_update: The A2A task artifact update to convert.
|
||||
author: The author of the event. Defaults to "a2a agent" if not provided.
|
||||
invocation_context: The invocation context containing session information.
|
||||
part_converter: The function to convert A2A part to GenAI part.
|
||||
|
||||
Returns:
|
||||
An ADK Event object representing the converted artifact update.
|
||||
"""
|
||||
if a2a_artifact_update is None:
|
||||
raise ValueError("A2A artifact update cannot be None")
|
||||
|
||||
try:
|
||||
output_parts, _ = _convert_a2a_parts_to_adk_parts(
|
||||
a2a_artifact_update.artifact.parts, part_converter
|
||||
)
|
||||
return _create_event(
|
||||
output_parts,
|
||||
invocation_context,
|
||||
author,
|
||||
partial=not a2a_artifact_update.last_chunk,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Failed to convert A2A artifact update to event: %s", e)
|
||||
raise RuntimeError(f"Failed to convert artifact update: {e}") from e
|
||||
@@ -38,6 +38,7 @@ from a2a.types import Message as A2AMessage
|
||||
from a2a.types import MessageSendConfiguration
|
||||
from a2a.types import Part as A2APart
|
||||
from a2a.types import Role
|
||||
from a2a.types import Task as A2ATask
|
||||
from a2a.types import TaskArtifactUpdateEvent as A2ATaskArtifactUpdateEvent
|
||||
from a2a.types import TaskState
|
||||
from a2a.types import TaskStatusUpdateEvent as A2ATaskStatusUpdateEvent
|
||||
@@ -62,6 +63,7 @@ from ..a2a.converters.part_converter import A2APartToGenAIPartConverter
|
||||
from ..a2a.converters.part_converter import convert_a2a_part_to_genai_part
|
||||
from ..a2a.converters.part_converter import convert_genai_part_to_a2a_part
|
||||
from ..a2a.converters.part_converter import GenAIPartToA2APartConverter
|
||||
from ..a2a.converters.utils import _get_adk_metadata_key
|
||||
from ..a2a.experimental import a2a_experimental
|
||||
from ..a2a.logs.log_utils import build_a2a_request_log
|
||||
from ..a2a.logs.log_utils import build_a2a_response_log
|
||||
@@ -522,6 +524,76 @@ class RemoteA2aAgent(BaseAgent):
|
||||
branch=ctx.branch,
|
||||
)
|
||||
|
||||
async def _handle_a2a_response_v2(
|
||||
self, a2a_response: A2AClientEvent | A2AMessage, ctx: InvocationContext
|
||||
) -> Optional[Event]:
|
||||
"""Handle A2A response and convert to Event.
|
||||
|
||||
Args:
|
||||
a2a_response: The A2A response object
|
||||
ctx: The invocation context
|
||||
|
||||
Returns:
|
||||
Event object representing the response, or None if no event should be
|
||||
emitted.
|
||||
"""
|
||||
try:
|
||||
if isinstance(a2a_response, tuple):
|
||||
task, update = a2a_response
|
||||
event = None
|
||||
if update is None:
|
||||
# This is the initial response for a streaming task or the complete
|
||||
# response for a non-streaming task.
|
||||
event = self._config.a2a_task_converter(
|
||||
task, self.name, ctx, self._config.a2a_part_converter
|
||||
)
|
||||
elif isinstance(update, A2ATaskStatusUpdateEvent):
|
||||
# This is a streaming task status update.
|
||||
event = self._config.a2a_status_update_converter(
|
||||
update, self.name, ctx, self._config.a2a_part_converter
|
||||
)
|
||||
elif isinstance(update, A2ATaskArtifactUpdateEvent):
|
||||
# This is a streaming task artifact update.
|
||||
event = self._config.a2a_artifact_update_converter(
|
||||
update, self.name, ctx, self._config.a2a_part_converter
|
||||
)
|
||||
if not event:
|
||||
return None
|
||||
event.custom_metadata = event.custom_metadata or {}
|
||||
event.custom_metadata[A2A_METADATA_PREFIX + "task_id"] = task.id
|
||||
if task.context_id:
|
||||
event.custom_metadata[A2A_METADATA_PREFIX + "context_id"] = (
|
||||
task.context_id
|
||||
)
|
||||
|
||||
# Otherwise, it's a regular A2AMessage.
|
||||
elif isinstance(a2a_response, A2AMessage):
|
||||
event = self._config.a2a_message_converter(
|
||||
a2a_response, self.name, ctx, self._config.a2a_part_converter
|
||||
)
|
||||
event.custom_metadata = event.custom_metadata or {}
|
||||
|
||||
if a2a_response.context_id:
|
||||
event.custom_metadata[A2A_METADATA_PREFIX + "context_id"] = (
|
||||
a2a_response.context_id
|
||||
)
|
||||
else:
|
||||
event = Event(
|
||||
author=self.name,
|
||||
error_message="Unknown A2A response type",
|
||||
invocation_id=ctx.invocation_id,
|
||||
branch=ctx.branch,
|
||||
)
|
||||
return event
|
||||
except A2AClientError as e:
|
||||
logger.error("Failed to handle A2A response: %s", e)
|
||||
return Event(
|
||||
author=self.name,
|
||||
error_message=f"Failed to process A2A response: {e}",
|
||||
invocation_id=ctx.invocation_id,
|
||||
branch=ctx.branch,
|
||||
)
|
||||
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
@@ -589,7 +661,20 @@ class RemoteA2aAgent(BaseAgent):
|
||||
):
|
||||
logger.debug(build_a2a_response_log(a2a_response))
|
||||
|
||||
event = await self._handle_a2a_response(a2a_response, ctx)
|
||||
metadata = None
|
||||
if isinstance(a2a_response, tuple):
|
||||
task = a2a_response[0]
|
||||
if task:
|
||||
metadata = task.metadata
|
||||
else:
|
||||
metadata = a2a_response.metadata
|
||||
|
||||
if metadata and metadata.get(
|
||||
_get_adk_metadata_key("agent_executor_v2")
|
||||
):
|
||||
event = await self._handle_a2a_response_v2(a2a_response, ctx)
|
||||
else:
|
||||
event = await self._handle_a2a_response(a2a_response, ctx)
|
||||
if not event:
|
||||
continue
|
||||
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Round trip tests for ADK and A2A event converters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict
|
||||
from unittest.mock import Mock
|
||||
|
||||
from a2a.types import TaskArtifactUpdateEvent
|
||||
from a2a.types import TaskStatusUpdateEvent
|
||||
from google.adk.a2a.converters.from_adk_event import convert_event_to_a2a_events
|
||||
from google.adk.a2a.converters.from_adk_event import create_error_status_event
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_status_update_to_event
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.events.event import Event
|
||||
from google.genai import types as genai_types
|
||||
|
||||
|
||||
def test_round_trip_text_event():
|
||||
original_event = Event(
|
||||
invocation_id="test_invocation",
|
||||
author="test_agent",
|
||||
branch="main",
|
||||
content=genai_types.Content(
|
||||
role="model",
|
||||
parts=[genai_types.Part.from_text(text="Hello world!")],
|
||||
),
|
||||
partial=False,
|
||||
)
|
||||
agents_artifacts: Dict[str, str] = {}
|
||||
|
||||
a2a_events = convert_event_to_a2a_events(
|
||||
event=original_event,
|
||||
agents_artifacts=agents_artifacts,
|
||||
task_id="task1",
|
||||
context_id="context1",
|
||||
)
|
||||
|
||||
assert len(a2a_events) == 1
|
||||
a2a_event = a2a_events[0]
|
||||
assert isinstance(a2a_event, TaskArtifactUpdateEvent)
|
||||
|
||||
mock_context = Mock(
|
||||
spec=InvocationContext, invocation_id="test_invocation", branch="main"
|
||||
)
|
||||
|
||||
restored_event = convert_a2a_artifact_update_to_event(
|
||||
a2a_artifact_update=a2a_event,
|
||||
author="test_agent",
|
||||
invocation_context=mock_context,
|
||||
)
|
||||
|
||||
assert restored_event is not None
|
||||
assert restored_event.author == original_event.author
|
||||
assert restored_event.invocation_id == original_event.invocation_id
|
||||
assert restored_event.branch == original_event.branch
|
||||
assert restored_event.partial == original_event.partial
|
||||
assert len(restored_event.content.parts) == len(original_event.content.parts)
|
||||
assert (
|
||||
restored_event.content.parts[0].text
|
||||
== original_event.content.parts[0].text
|
||||
)
|
||||
|
||||
|
||||
def test_round_trip_error_status_event():
|
||||
original_event = Event(
|
||||
invocation_id="error_inv",
|
||||
author="error_agent",
|
||||
branch="main",
|
||||
error_message="Test Error",
|
||||
)
|
||||
|
||||
a2a_event = create_error_status_event(
|
||||
event=original_event,
|
||||
task_id="task2",
|
||||
context_id="ctx2",
|
||||
)
|
||||
|
||||
assert isinstance(a2a_event, TaskStatusUpdateEvent)
|
||||
|
||||
mock_context = Mock(
|
||||
spec=InvocationContext, invocation_id="error_inv", branch="main"
|
||||
)
|
||||
|
||||
restored_event = convert_a2a_status_update_to_event(
|
||||
a2a_status_update=a2a_event,
|
||||
author="error_agent",
|
||||
invocation_context=mock_context,
|
||||
)
|
||||
|
||||
assert restored_event is not None
|
||||
assert restored_event.author == original_event.author
|
||||
assert restored_event.invocation_id == original_event.invocation_id
|
||||
assert restored_event.branch == original_event.branch
|
||||
assert len(restored_event.content.parts) == 1
|
||||
assert restored_event.content.parts[0].text == "Test Error"
|
||||
|
||||
|
||||
def test_round_trip_function_call_event():
|
||||
original_event = Event(
|
||||
invocation_id="test_invocation",
|
||||
author="test_agent",
|
||||
branch="main",
|
||||
content=genai_types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
genai_types.Part.from_function_call(
|
||||
name="my_function",
|
||||
args={"arg1": "value1"},
|
||||
)
|
||||
],
|
||||
),
|
||||
partial=False,
|
||||
)
|
||||
agents_artifacts: Dict[str, str] = {}
|
||||
|
||||
a2a_events = convert_event_to_a2a_events(
|
||||
event=original_event,
|
||||
agents_artifacts=agents_artifacts,
|
||||
task_id="task1",
|
||||
context_id="context1",
|
||||
)
|
||||
|
||||
assert len(a2a_events) == 1
|
||||
a2a_event = a2a_events[0]
|
||||
|
||||
mock_context = Mock(
|
||||
spec=InvocationContext, invocation_id="test_invocation", branch="main"
|
||||
)
|
||||
|
||||
restored_event = convert_a2a_artifact_update_to_event(
|
||||
a2a_artifact_update=a2a_event,
|
||||
author="test_agent",
|
||||
invocation_context=mock_context,
|
||||
)
|
||||
|
||||
assert restored_event is not None
|
||||
assert restored_event.author == original_event.author
|
||||
assert restored_event.invocation_id == original_event.invocation_id
|
||||
assert restored_event.branch == original_event.branch
|
||||
assert len(restored_event.content.parts) == 1
|
||||
assert restored_event.content.parts[0].function_call.name == "my_function"
|
||||
assert restored_event.content.parts[0].function_call.args == {
|
||||
"arg1": "value1"
|
||||
}
|
||||
|
||||
|
||||
def test_round_trip_function_response_event():
|
||||
original_event = Event(
|
||||
invocation_id="test_invocation",
|
||||
author="test_agent",
|
||||
branch="main",
|
||||
content=genai_types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
genai_types.Part.from_function_response(
|
||||
name="my_function",
|
||||
response={"result": "success"},
|
||||
)
|
||||
],
|
||||
),
|
||||
partial=False,
|
||||
)
|
||||
agents_artifacts: Dict[str, str] = {}
|
||||
|
||||
a2a_events = convert_event_to_a2a_events(
|
||||
event=original_event,
|
||||
agents_artifacts=agents_artifacts,
|
||||
task_id="task1",
|
||||
context_id="context1",
|
||||
)
|
||||
|
||||
assert len(a2a_events) == 1
|
||||
a2a_event = a2a_events[0]
|
||||
|
||||
mock_context = Mock(
|
||||
spec=InvocationContext, invocation_id="test_invocation", branch="main"
|
||||
)
|
||||
|
||||
restored_event = convert_a2a_artifact_update_to_event(
|
||||
a2a_artifact_update=a2a_event,
|
||||
author="test_agent",
|
||||
invocation_context=mock_context,
|
||||
)
|
||||
|
||||
assert restored_event is not None
|
||||
assert restored_event.author == original_event.author
|
||||
assert restored_event.invocation_id == original_event.invocation_id
|
||||
assert restored_event.branch == original_event.branch
|
||||
assert len(restored_event.content.parts) == 1
|
||||
assert restored_event.content.parts[0].function_response.name == "my_function"
|
||||
assert restored_event.content.parts[0].function_response.response == {
|
||||
"result": "success"
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
from a2a.types import Artifact
|
||||
from a2a.types import Message
|
||||
from a2a.types import Part as A2APart
|
||||
from a2a.types import Task
|
||||
from a2a.types import TaskArtifactUpdateEvent
|
||||
from a2a.types import TaskState
|
||||
from a2a.types import TaskStatus
|
||||
from a2a.types import TaskStatusUpdateEvent
|
||||
from a2a.types import TextPart
|
||||
from google.adk.a2a.converters.part_converter import A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_artifact_update_to_event
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_message_to_event
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_status_update_to_event
|
||||
from google.adk.a2a.converters.to_adk_event import convert_a2a_task_to_event
|
||||
from google.adk.a2a.converters.utils import _get_adk_metadata_key
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.genai import types as genai_types
|
||||
import pytest
|
||||
|
||||
|
||||
class TestToAdk:
|
||||
"""Test suite for to_adk functions."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.mock_context = Mock(spec=InvocationContext)
|
||||
self.mock_context.invocation_id = "test-invocation"
|
||||
self.mock_context.branch = "test-branch"
|
||||
|
||||
def test_convert_a2a_message_to_event_success(self):
|
||||
"""Test successful conversion of A2A message to Event."""
|
||||
a2a_part = Mock(spec=A2APart)
|
||||
a2a_part.root = Mock()
|
||||
a2a_part.root.metadata = {}
|
||||
message = Message(message_id="msg-1", role="user", parts=[a2a_part])
|
||||
|
||||
mock_genai_part = genai_types.Part.from_text(text="hello")
|
||||
mock_part_converter = Mock(return_value=[mock_genai_part])
|
||||
|
||||
event = convert_a2a_message_to_event(
|
||||
message,
|
||||
author="test-author",
|
||||
invocation_context=self.mock_context,
|
||||
part_converter=mock_part_converter,
|
||||
)
|
||||
|
||||
assert event.author == "test-author"
|
||||
assert event.invocation_id == "test-invocation"
|
||||
assert event.branch == "test-branch"
|
||||
assert len(event.content.parts) == 1
|
||||
assert event.content.parts[0] == mock_genai_part
|
||||
|
||||
def test_convert_a2a_message_to_event_none(self):
|
||||
"""Test convert_a2a_message_to_event with None."""
|
||||
with pytest.raises(ValueError, match="A2A message cannot be None"):
|
||||
convert_a2a_message_to_event(None)
|
||||
|
||||
def test_convert_a2a_task_to_event_success(self):
|
||||
"""Test successful conversion of A2A task to Event."""
|
||||
a2a_part = Mock(spec=A2APart)
|
||||
a2a_part.root = Mock()
|
||||
a2a_part.root.metadata = {}
|
||||
task = Task(
|
||||
id="task-1",
|
||||
status=TaskStatus(
|
||||
state=TaskState.submitted, timestamp="2024-01-01T00:00:00Z"
|
||||
),
|
||||
context_id="context-1",
|
||||
history=[Message(message_id="msg-1", role="agent", parts=[a2a_part])],
|
||||
artifacts=[
|
||||
Artifact(
|
||||
artifact_id="art-1", artifact_type="message", parts=[a2a_part]
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
mock_genai_part = genai_types.Part.from_text(text="task artifact text")
|
||||
mock_part_converter = Mock(return_value=[mock_genai_part])
|
||||
|
||||
event = convert_a2a_task_to_event(
|
||||
task,
|
||||
author="test-author",
|
||||
invocation_context=self.mock_context,
|
||||
part_converter=mock_part_converter,
|
||||
)
|
||||
|
||||
assert event.author == "test-author"
|
||||
assert event.invocation_id == "test-invocation"
|
||||
assert len(event.content.parts) == 1
|
||||
assert event.content.parts[0] == mock_genai_part
|
||||
|
||||
def test_convert_a2a_task_to_event_none(self):
|
||||
"""Test convert_a2a_task_to_event with None."""
|
||||
with pytest.raises(ValueError, match="A2A task cannot be None"):
|
||||
convert_a2a_task_to_event(None)
|
||||
|
||||
def test_convert_a2a_status_update_to_event_success(self):
|
||||
"""Test successful conversion of A2A status update to Event."""
|
||||
a2a_part = Mock(spec=A2APart)
|
||||
a2a_part.root = Mock()
|
||||
a2a_part.root.metadata = {
|
||||
_get_adk_metadata_key(A2A_DATA_PART_METADATA_IS_LONG_RUNNING_KEY): True
|
||||
}
|
||||
update = TaskStatusUpdateEvent(
|
||||
task_id="task-1",
|
||||
status=TaskStatus(
|
||||
state=TaskState.input_required,
|
||||
timestamp="now",
|
||||
message=Message(
|
||||
message_id="m1",
|
||||
role="agent",
|
||||
parts=[a2a_part],
|
||||
),
|
||||
),
|
||||
context_id="context-1",
|
||||
final=False,
|
||||
)
|
||||
|
||||
mock_genai_part = genai_types.Part(
|
||||
function_call=genai_types.FunctionCall(
|
||||
name="status update text", args={"arg": "value"}, id="call-1"
|
||||
)
|
||||
)
|
||||
mock_part_converter = Mock(return_value=[mock_genai_part])
|
||||
|
||||
event = convert_a2a_status_update_to_event(
|
||||
update,
|
||||
author="test-author",
|
||||
invocation_context=self.mock_context,
|
||||
part_converter=mock_part_converter,
|
||||
)
|
||||
|
||||
assert event.author == "test-author"
|
||||
assert event.invocation_id == "test-invocation"
|
||||
assert len(event.content.parts) == 1
|
||||
assert event.content.parts[0] == mock_genai_part
|
||||
|
||||
def test_convert_a2a_status_update_to_event_none(self):
|
||||
"""Test convert_a2a_status_update_to_event with None."""
|
||||
with pytest.raises(ValueError, match="A2A status update cannot be None"):
|
||||
convert_a2a_status_update_to_event(None)
|
||||
|
||||
def test_convert_a2a_artifact_update_to_event_success(self):
|
||||
"""Test successful conversion of A2A artifact update to Event."""
|
||||
a2a_part = Mock(spec=A2APart)
|
||||
a2a_part.root = Mock()
|
||||
a2a_part.root.metadata = {}
|
||||
update = TaskArtifactUpdateEvent(
|
||||
task_id="task-1",
|
||||
artifact=Artifact(
|
||||
artifact_id="art-1", artifact_type="message", parts=[a2a_part]
|
||||
),
|
||||
append=True,
|
||||
context_id="context-1",
|
||||
last_chunk=False,
|
||||
)
|
||||
|
||||
mock_genai_part = genai_types.Part.from_text(text="artifact chunk text")
|
||||
mock_part_converter = Mock(return_value=[mock_genai_part])
|
||||
|
||||
event = convert_a2a_artifact_update_to_event(
|
||||
update,
|
||||
author="test-author",
|
||||
invocation_context=self.mock_context,
|
||||
part_converter=mock_part_converter,
|
||||
)
|
||||
|
||||
assert event.author == "test-author"
|
||||
assert event.invocation_id == "test-invocation"
|
||||
assert event.partial is True
|
||||
assert len(event.content.parts) == 1
|
||||
assert event.content.parts[0] == mock_genai_part
|
||||
|
||||
def test_convert_a2a_artifact_update_to_event_none(self):
|
||||
"""Test convert_a2a_artifact_update_to_event with None."""
|
||||
with pytest.raises(ValueError, match="A2A artifact update cannot be None"):
|
||||
convert_a2a_artifact_update_to_event(None)
|
||||
@@ -1685,6 +1685,236 @@ class TestRemoteA2aAgentMessageHandlingFromFactory:
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestRemoteA2aAgentMessageHandlingV2:
|
||||
"""Test _handle_a2a_response_impl functionality."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Setup test fixtures."""
|
||||
from google.adk.a2a.agent.config import A2aRemoteAgentConfig
|
||||
|
||||
self.agent_card = create_test_agent_card()
|
||||
self.mock_config = Mock(spec=A2aRemoteAgentConfig)
|
||||
self.mock_config.a2a_part_converter = Mock()
|
||||
self.mock_config.a2a_task_converter = Mock()
|
||||
self.mock_config.a2a_status_update_converter = Mock()
|
||||
self.mock_config.a2a_artifact_update_converter = Mock()
|
||||
self.mock_config.a2a_message_converter = Mock()
|
||||
|
||||
self.agent = RemoteA2aAgent(
|
||||
name="test_agent",
|
||||
agent_card=self.agent_card,
|
||||
config=self.mock_config,
|
||||
)
|
||||
|
||||
# Mock session and context
|
||||
self.mock_session = Mock(spec=Session)
|
||||
self.mock_session.id = "session-123"
|
||||
self.mock_session.events = []
|
||||
|
||||
self.mock_context = Mock(spec=InvocationContext)
|
||||
self.mock_context.session = self.mock_session
|
||||
self.mock_context.invocation_id = "invocation-123"
|
||||
self.mock_context.branch = "main"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_with_message(self):
|
||||
"""Test _handle_a2a_response_impl with A2AMessage."""
|
||||
mock_a2a_message = Mock(spec=A2AMessage)
|
||||
mock_a2a_message.metadata = {}
|
||||
mock_a2a_message.metadata = {}
|
||||
mock_a2a_message.context_id = "context-123"
|
||||
|
||||
mock_event = Event(
|
||||
author=self.agent.name,
|
||||
invocation_id=self.mock_context.invocation_id,
|
||||
branch=self.mock_context.branch,
|
||||
)
|
||||
self.mock_config.a2a_message_converter.return_value = mock_event
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
mock_a2a_message, self.mock_context
|
||||
)
|
||||
|
||||
assert result == mock_event
|
||||
self.mock_config.a2a_message_converter.assert_called_once_with(
|
||||
mock_a2a_message,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_config.a2a_part_converter,
|
||||
)
|
||||
assert result.custom_metadata is not None
|
||||
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
||||
assert (
|
||||
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
||||
== "context-123"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_with_task_and_no_update(self):
|
||||
"""Test _handle_a2a_response_impl with Task and no update."""
|
||||
mock_a2a_task = Mock(spec=A2ATask)
|
||||
mock_a2a_task.id = "task-123"
|
||||
mock_a2a_task.context_id = "context-123"
|
||||
|
||||
mock_event = Event(
|
||||
author=self.agent.name,
|
||||
invocation_id=self.mock_context.invocation_id,
|
||||
branch=self.mock_context.branch,
|
||||
)
|
||||
self.mock_config.a2a_task_converter.return_value = mock_event
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
(mock_a2a_task, None), self.mock_context
|
||||
)
|
||||
|
||||
assert result == mock_event
|
||||
self.mock_config.a2a_task_converter.assert_called_once_with(
|
||||
mock_a2a_task,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_config.a2a_part_converter,
|
||||
)
|
||||
assert result.custom_metadata is not None
|
||||
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
||||
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
||||
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
||||
assert (
|
||||
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
||||
== "context-123"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_with_task_status_update(self):
|
||||
"""Test _handle_a2a_response_impl with TaskStatusUpdateEvent."""
|
||||
mock_a2a_task = Mock(spec=A2ATask)
|
||||
mock_a2a_task.id = "task-123"
|
||||
mock_a2a_task.context_id = None
|
||||
|
||||
mock_update = Mock(spec=TaskStatusUpdateEvent)
|
||||
|
||||
mock_event = Event(
|
||||
author=self.agent.name,
|
||||
invocation_id=self.mock_context.invocation_id,
|
||||
branch=self.mock_context.branch,
|
||||
)
|
||||
self.mock_config.a2a_status_update_converter.return_value = mock_event
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
(mock_a2a_task, mock_update), self.mock_context
|
||||
)
|
||||
|
||||
assert result == mock_event
|
||||
self.mock_config.a2a_status_update_converter.assert_called_once_with(
|
||||
mock_update,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_config.a2a_part_converter,
|
||||
)
|
||||
assert result.custom_metadata is not None
|
||||
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
||||
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
||||
assert A2A_METADATA_PREFIX + "context_id" not in result.custom_metadata
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_with_task_artifact_update(self):
|
||||
"""Test _handle_a2a_response_impl with TaskArtifactUpdateEvent."""
|
||||
mock_a2a_task = Mock(spec=A2ATask)
|
||||
mock_a2a_task.id = "task-123"
|
||||
mock_a2a_task.context_id = "context-123"
|
||||
|
||||
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
||||
|
||||
mock_event = Event(
|
||||
author=self.agent.name,
|
||||
invocation_id=self.mock_context.invocation_id,
|
||||
branch=self.mock_context.branch,
|
||||
)
|
||||
self.mock_config.a2a_artifact_update_converter.return_value = mock_event
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
(mock_a2a_task, mock_update), self.mock_context
|
||||
)
|
||||
|
||||
assert result == mock_event
|
||||
self.mock_config.a2a_artifact_update_converter.assert_called_once_with(
|
||||
mock_update,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_config.a2a_part_converter,
|
||||
)
|
||||
assert result.custom_metadata is not None
|
||||
assert A2A_METADATA_PREFIX + "task_id" in result.custom_metadata
|
||||
assert result.custom_metadata[A2A_METADATA_PREFIX + "task_id"] == "task-123"
|
||||
assert A2A_METADATA_PREFIX + "context_id" in result.custom_metadata
|
||||
assert (
|
||||
result.custom_metadata[A2A_METADATA_PREFIX + "context_id"]
|
||||
== "context-123"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_update_converter_returns_none(self):
|
||||
"""Test _handle_a2a_response_impl when converter returns None."""
|
||||
mock_a2a_task = Mock(spec=A2ATask)
|
||||
mock_a2a_task.id = "task-123"
|
||||
|
||||
mock_update = Mock(spec=TaskArtifactUpdateEvent)
|
||||
|
||||
self.mock_config.a2a_artifact_update_converter.return_value = None
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
(mock_a2a_task, mock_update), self.mock_context
|
||||
)
|
||||
|
||||
assert result is None
|
||||
self.mock_config.a2a_artifact_update_converter.assert_called_once_with(
|
||||
mock_update,
|
||||
self.agent.name,
|
||||
self.mock_context,
|
||||
self.mock_config.a2a_part_converter,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_unknown_response_type(self):
|
||||
"""Test _handle_a2a_response_impl with unknown response type."""
|
||||
unknown_response = object()
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
unknown_response, self.mock_context
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.author == self.agent.name
|
||||
assert result.error_message == "Unknown A2A response type"
|
||||
assert result.invocation_id == self.mock_context.invocation_id
|
||||
assert result.branch == self.mock_context.branch
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_a2a_response_impl_handles_client_error(self):
|
||||
"""Test _handle_a2a_response_impl catches A2AClientError."""
|
||||
mock_a2a_message = Mock(spec=A2AMessage)
|
||||
mock_a2a_message.metadata = {}
|
||||
mock_a2a_message.metadata = {}
|
||||
|
||||
from google.adk.agents.remote_a2a_agent import A2AClientError
|
||||
|
||||
self.mock_config.a2a_message_converter.side_effect = A2AClientError(
|
||||
"Test client error"
|
||||
)
|
||||
|
||||
result = await self.agent._handle_a2a_response_v2(
|
||||
mock_a2a_message, self.mock_context
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.author == self.agent.name
|
||||
assert (
|
||||
"Failed to process A2A response: Test client error"
|
||||
in result.error_message
|
||||
)
|
||||
assert result.invocation_id == self.mock_context.invocation_id
|
||||
assert result.branch == self.mock_context.branch
|
||||
|
||||
|
||||
class TestRemoteA2aAgentExecution:
|
||||
"""Test agent execution functionality."""
|
||||
|
||||
@@ -1773,7 +2003,7 @@ class TestRemoteA2aAgentExecution:
|
||||
|
||||
# Mock A2A client
|
||||
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
||||
mock_response = Mock()
|
||||
mock_response = Mock(metadata={})
|
||||
mock_send_message = AsyncMock()
|
||||
mock_send_message.__aiter__.return_value = [mock_response]
|
||||
mock_a2a_client.send_message.return_value = mock_send_message
|
||||
@@ -1912,7 +2142,7 @@ class TestRemoteA2aAgentExecution:
|
||||
|
||||
# Mock A2A client
|
||||
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
||||
mock_response = Mock()
|
||||
mock_response = Mock(metadata={})
|
||||
mock_send_message = AsyncMock()
|
||||
mock_send_message.__aiter__.return_value = [mock_response]
|
||||
mock_a2a_client.send_message.return_value = mock_send_message
|
||||
@@ -2049,7 +2279,7 @@ class TestRemoteA2aAgentExecutionFromFactory:
|
||||
|
||||
# Mock A2A client
|
||||
mock_a2a_client = create_autospec(spec=A2AClient, instance=True)
|
||||
mock_response = Mock()
|
||||
mock_response = Mock(metadata={})
|
||||
mock_send_message = AsyncMock()
|
||||
mock_send_message.__aiter__.return_value = [mock_response]
|
||||
mock_a2a_client.send_message.return_value = mock_send_message
|
||||
@@ -2295,6 +2525,7 @@ class TestRemoteA2aAgentIntegration:
|
||||
with patch.object(agent, "_a2a_client") as mock_a2a_client:
|
||||
mock_a2a_message = create_autospec(spec=A2AMessage, instance=True)
|
||||
mock_a2a_message.context_id = "context-123"
|
||||
mock_a2a_message.metadata = {}
|
||||
mock_response = mock_a2a_message
|
||||
|
||||
mock_send_message = AsyncMock()
|
||||
@@ -2391,6 +2622,7 @@ class TestRemoteA2aAgentIntegration:
|
||||
with patch.object(agent, "_a2a_client") as mock_a2a_client:
|
||||
mock_a2a_message = create_autospec(spec=A2AMessage, instance=True)
|
||||
mock_a2a_message.context_id = "context-123"
|
||||
mock_a2a_message.metadata = {}
|
||||
mock_response = mock_a2a_message
|
||||
|
||||
mock_send_message = AsyncMock()
|
||||
|
||||
Reference in New Issue
Block a user