mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
d1f182e8e6
commit
a77d68964a
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user