feat(config): support loading from YAML config in CLI

The supported CLIs are: `adk web`, `adk run` and `adk api_server`.

PiperOrigin-RevId: 781666724
This commit is contained in:
Liang Wu
2025-07-10 13:47:58 -07:00
committed by Copybara-Service
parent ca396a3ab1
commit aef54f8eb7
3 changed files with 246 additions and 2 deletions
+131 -1
View File
@@ -19,6 +19,7 @@ import tempfile
from textwrap import dedent
from google.adk.cli.utils.agent_loader import AgentLoader
from pydantic import ValidationError
import pytest
@@ -30,6 +31,8 @@ class TestAgentLoader:
"""Ensure sys.path is restored after each test."""
original_path = sys.path.copy()
original_env = os.environ.copy()
# Enable WIP features for YAML agent loading tests
os.environ["ADK_ALLOW_WIP_FEATURES"] = "true"
yield
sys.path[:] = original_path
# Restore environment variables
@@ -292,7 +295,8 @@ class TestAgentLoader:
expected_msg_part_1 = "No root_agent found for 'nonexistent_agent'."
expected_msg_part_2 = (
"Searched in 'nonexistent_agent.agent.root_agent',"
" 'nonexistent_agent.root_agent'."
" 'nonexistent_agent.root_agent' and"
" 'nonexistent_agent/root_agent.yaml'."
)
expected_msg_part_3 = (
f"Ensure '{agents_dir}/nonexistent_agent' is structured correctly"
@@ -443,3 +447,129 @@ class TestAgentLoader:
# Now assert path was added
assert str(temp_path) in sys.path
assert agent.name == "path_agent"
def create_yaml_agent_structure(
self, temp_dir: Path, agent_name: str, yaml_content: str
):
"""Create an agent structure with YAML configuration.
Args:
temp_dir: The temporary directory to create the agent in
agent_name: Name of the agent
yaml_content: YAML content for the root_agent.yaml file
"""
agent_dir = temp_dir / agent_name
agent_dir.mkdir()
# Create root_agent.yaml file
yaml_file = agent_dir / "root_agent.yaml"
yaml_file.write_text(yaml_content)
def test_load_agent_from_yaml_config(self):
"""Test loading an agent from YAML configuration."""
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir)
agent_name = "yaml_agent"
# Create YAML configuration
yaml_content = dedent("""
agent_class: LlmAgent
name: yaml_test_agent
model: gemini-2.0-flash
instruction: You are a test agent loaded from YAML configuration.
description: A test agent created from YAML config
""")
self.create_yaml_agent_structure(temp_path, agent_name, yaml_content)
# Load the agent
loader = AgentLoader(str(temp_path))
agent = loader.load_agent(agent_name)
# Assert agent was loaded correctly
assert agent.name == "yaml_test_agent"
# Check if it's an LlmAgent before accessing model and instruction
from google.adk.agents.llm_agent import LlmAgent
if isinstance(agent, LlmAgent):
assert agent.model == "gemini-2.0-flash"
# Handle instruction which can be string or InstructionProvider
instruction_text = str(agent.instruction)
assert "test agent loaded from YAML" in instruction_text
def test_yaml_agent_caching_returns_same_instance(self):
"""Test that loading the same YAML agent twice returns the same instance."""
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir)
agent_name = "cached_yaml_agent"
# Create YAML configuration
yaml_content = dedent("""
agent_class: LlmAgent
name: cached_yaml_test_agent
model: gemini-2.0-flash
instruction: You are a cached test agent.
""")
self.create_yaml_agent_structure(temp_path, agent_name, yaml_content)
# Load the agent twice
loader = AgentLoader(str(temp_path))
agent1 = loader.load_agent(agent_name)
agent2 = loader.load_agent(agent_name)
# Assert same instance is returned
assert agent1 is agent2
assert agent1.name == agent2.name
def test_yaml_agent_not_found_error(self):
"""Test that appropriate error is raised when YAML agent is not found."""
with tempfile.TemporaryDirectory() as temp_dir:
loader = AgentLoader(temp_dir)
agents_dir = temp_dir # For use in the expected message string
# Try to load non-existent YAML agent
with pytest.raises(ValueError) as exc_info:
loader.load_agent("nonexistent_yaml_agent")
expected_msg_part_1 = "No root_agent found for 'nonexistent_yaml_agent'."
expected_msg_part_2 = (
"Searched in 'nonexistent_yaml_agent.agent.root_agent',"
" 'nonexistent_yaml_agent.root_agent' and"
" 'nonexistent_yaml_agent/root_agent.yaml'."
)
expected_msg_part_3 = (
f"Ensure '{agents_dir}/nonexistent_yaml_agent' is structured"
" correctly"
)
assert expected_msg_part_1 in str(exc_info.value)
assert expected_msg_part_2 in str(exc_info.value)
assert expected_msg_part_3 in str(exc_info.value)
def test_yaml_agent_invalid_yaml_error(self):
"""Test that appropriate error is raised when YAML is invalid."""
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir)
agent_name = "invalid_yaml_agent"
# Create invalid YAML content with wrong field name
invalid_yaml_content = dedent("""
agent_type: LlmAgent
name: invalid_yaml_test_agent
model: gemini-2.0-flash
instruction: You are a test agent with invalid YAML
""")
self.create_yaml_agent_structure(
temp_path, agent_name, invalid_yaml_content
)
loader = AgentLoader(str(temp_path))
# Try to load agent with invalid YAML
with pytest.raises(ValidationError) as exc_info:
loader.load_agent(agent_name)
# Should raise some form of YAML parsing error
assert "Extra inputs are not permitted" in str(exc_info.value)