feat: Add an to_a2a util to convert adk agent to A2A ASGI application

Users can do :
```
a2a_app = to_a2a(root_agent)
```
then use below command to start a2a server:
```
uvicorn user_module:a2a_app --host localhost --port 8000
```

PiperOrigin-RevId: 785625048
This commit is contained in:
Xiang (Sean) Zhou
2025-07-21 16:54:53 -07:00
committed by Copybara-Service
parent d1f182e8e6
commit a77d68964a
3 changed files with 899 additions and 37 deletions
+118
View File
@@ -0,0 +1,118 @@
# Copyright 2025 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
import logging
import sys
try:
from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
except ImportError as e:
if sys.version_info < (3, 10):
raise ImportError(
"A2A requires Python 3.10 or above. Please upgrade your Python version."
) from e
else:
raise e
from starlette.applications import Starlette
from ...agents.base_agent import BaseAgent
from ...artifacts.in_memory_artifact_service import InMemoryArtifactService
from ...auth.credential_service.in_memory_credential_service import InMemoryCredentialService
from ...cli.utils.logs import setup_adk_logger
from ...memory.in_memory_memory_service import InMemoryMemoryService
from ...runners import Runner
from ...sessions.in_memory_session_service import InMemorySessionService
from ..executor.a2a_agent_executor import A2aAgentExecutor
from .agent_card_builder import AgentCardBuilder
def to_a2a(
agent: BaseAgent, *, host: str = "localhost", port: int = 8000
) -> Starlette:
"""Convert an ADK agent to a A2A Starlette application.
Args:
agent: The ADK agent to convert
host: The host for the A2A RPC URL (default: "localhost")
port: The port for the A2A RPC URL (default: 8000)
Returns:
A Starlette application that can be run with uvicorn
Example:
agent = MyAgent()
app = to_a2a(agent, host="localhost", port=8000)
# Then run with: uvicorn module:app --host localhost --port 8000
"""
# Set up ADK logging to ensure logs are visible when using uvicorn directly
setup_adk_logger(logging.INFO)
async def create_runner() -> Runner:
"""Create a runner for the agent."""
return Runner(
app_name=agent.name or "adk_agent",
agent=agent,
# Use minimal services - in a real implementation these could be configured
artifact_service=InMemoryArtifactService(),
session_service=InMemorySessionService(),
memory_service=InMemoryMemoryService(),
credential_service=InMemoryCredentialService(),
)
# Create A2A components
task_store = InMemoryTaskStore()
agent_executor = A2aAgentExecutor(
runner=create_runner,
)
request_handler = DefaultRequestHandler(
agent_executor=agent_executor, task_store=task_store
)
# Build agent card
rpc_url = f"http://{host}:{port}/"
card_builder = AgentCardBuilder(
agent=agent,
rpc_url=rpc_url,
)
# Create a Starlette app that will be configured during startup
app = Starlette()
# Add startup handler to build the agent card and configure A2A routes
async def setup_a2a():
# Build the agent card asynchronously
agent_card = await card_builder.build()
# Create the A2A Starlette application
a2a_app = A2AStarletteApplication(
agent_card=agent_card,
http_handler=request_handler,
)
# Add A2A routes to the main app
a2a_app.add_routes_to_app(
app,
)
# Store the setup function to be called during startup
app.add_event_handler("startup", setup_a2a)
return app
File diff suppressed because it is too large Load Diff
@@ -14,7 +14,7 @@
"""Unit tests for LlmAgent output saving functionality."""
from unittest.mock import Mock
import logging
from unittest.mock import patch
from google.adk.agents.llm_agent import LlmAgent
@@ -61,20 +61,33 @@ class TestLlmAgentOutputSave:
def test_maybe_save_output_to_state_skips_different_author(self, caplog):
"""Test that output is not saved when event author differs from agent name."""
agent = LlmAgent(name="agent_a", output_key="result")
event = create_test_event(author="agent_b", content_text="Response from B")
with caplog.at_level("DEBUG"):
agent._LlmAgent__maybe_save_output_to_state(event)
# Should not add anything to state_delta
assert len(event.actions.state_delta) == 0
# Should log the skip
assert (
"Skipping output save for agent agent_a: event authored by agent_b"
in caplog.text
# Set the LlmAgent logger to DEBUG level
llm_agent_logger = logging.getLogger(
"google_adk.google.adk.agents.llm_agent"
)
original_level = llm_agent_logger.level
llm_agent_logger.setLevel(logging.DEBUG)
try:
agent = LlmAgent(name="agent_a", output_key="result")
event = create_test_event(
author="agent_b", content_text="Response from B"
)
with caplog.at_level("DEBUG"):
agent._LlmAgent__maybe_save_output_to_state(event)
# Should not add anything to state_delta
assert len(event.actions.state_delta) == 0
# Should log the skip
assert (
"Skipping output save for agent agent_a: event authored by agent_b"
in caplog.text
)
finally:
# Restore original logger level
llm_agent_logger.setLevel(original_level)
def test_maybe_save_output_to_state_saves_same_author(self):
"""Test that output is saved when event author matches agent name."""
@@ -163,36 +176,61 @@ class TestLlmAgentOutputSave:
# Scenario: Agent A transfers to Agent B, Agent B produces output
# Agent A should not save Agent B's output
agent_a = LlmAgent(name="support_agent", output_key="support_result")
agent_b_event = create_test_event(
author="billing_agent", content_text="Your bill is $100"
# Set the LlmAgent logger to DEBUG level
llm_agent_logger = logging.getLogger(
"google_adk.google.adk.agents.llm_agent"
)
original_level = llm_agent_logger.level
llm_agent_logger.setLevel(logging.DEBUG)
with caplog.at_level("DEBUG"):
agent_a._LlmAgent__maybe_save_output_to_state(agent_b_event)
try:
agent_a = LlmAgent(name="support_agent", output_key="support_result")
agent_b_event = create_test_event(
author="billing_agent", content_text="Your bill is $100"
)
# Agent A should not save Agent B's output
assert len(agent_b_event.actions.state_delta) == 0
assert (
"Skipping output save for agent support_agent: event authored by"
" billing_agent"
in caplog.text
)
with caplog.at_level("DEBUG"):
agent_a._LlmAgent__maybe_save_output_to_state(agent_b_event)
# Agent A should not save Agent B's output
assert len(agent_b_event.actions.state_delta) == 0
assert (
"Skipping output save for agent support_agent: event authored by"
" billing_agent"
in caplog.text
)
finally:
# Restore original logger level
llm_agent_logger.setLevel(original_level)
def test_maybe_save_output_to_state_case_sensitive_names(self, caplog):
"""Test that agent name comparison is case-sensitive."""
agent = LlmAgent(name="TestAgent", output_key="result")
event = create_test_event(author="testagent", content_text="Test response")
with caplog.at_level("DEBUG"):
agent._LlmAgent__maybe_save_output_to_state(event)
# Should not save due to case mismatch
assert len(event.actions.state_delta) == 0
assert (
"Skipping output save for agent TestAgent: event authored by testagent"
in caplog.text
# Set the LlmAgent logger to DEBUG level
llm_agent_logger = logging.getLogger(
"google_adk.google.adk.agents.llm_agent"
)
original_level = llm_agent_logger.level
llm_agent_logger.setLevel(logging.DEBUG)
try:
agent = LlmAgent(name="TestAgent", output_key="result")
event = create_test_event(
author="testagent", content_text="Test response"
)
with caplog.at_level("DEBUG"):
agent._LlmAgent__maybe_save_output_to_state(event)
# Should not save due to case mismatch
assert len(event.actions.state_delta) == 0
assert (
"Skipping output save for agent TestAgent: event authored by"
" testagent"
in caplog.text
)
finally:
# Restore original logger level
llm_agent_logger.setLevel(original_level)
@patch("google.adk.agents.llm_agent.logger")
def test_maybe_save_output_to_state_logging(self, mock_logger):