From aef54f8eb7e1aaa5d3830583c876c294e0266dc3 Mon Sep 17 00:00:00 2001 From: Liang Wu Date: Thu, 10 Jul 2025 13:47:22 -0700 Subject: [PATCH] feat(config): support loading from YAML config in CLI The supported CLIs are: `adk web`, `adk run` and `adk api_server`. PiperOrigin-RevId: 781666724 --- src/google/adk/agents/config_agent_utils.py | 80 +++++++++++ src/google/adk/cli/utils/agent_loader.py | 36 ++++- .../unittests/cli/utils/test_agent_loader.py | 132 +++++++++++++++++- 3 files changed, 246 insertions(+), 2 deletions(-) create mode 100644 src/google/adk/agents/config_agent_utils.py diff --git a/src/google/adk/agents/config_agent_utils.py b/src/google/adk/agents/config_agent_utils.py new file mode 100644 index 00000000..08991146 --- /dev/null +++ b/src/google/adk/agents/config_agent_utils.py @@ -0,0 +1,80 @@ +# 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 os +from pathlib import Path + +import yaml + +from ..utils.feature_decorator import working_in_progress +from .agent_config import AgentConfig +from .base_agent import BaseAgent +from .llm_agent import LlmAgent +from .llm_agent import LlmAgentConfig +from .loop_agent import LoopAgent +from .loop_agent import LoopAgentConfig + + +@working_in_progress("from_config is not ready for use.") +def from_config(config_path: str) -> BaseAgent: + """Build agent from a configfile path. + + Args: + config: the path to a YAML config file. + + Returns: + The created agent instance. + + Raises: + FileNotFoundError: If config file doesn't exist. + ValidationError: If config file's content is invalid YAML. + ValueError: If agent type is unsupported. + """ + abs_path = os.path.abspath(config_path) + config = _load_config_from_path(abs_path) + + if isinstance(config.root, LlmAgentConfig): + return LlmAgent.from_config(config.root) + elif isinstance(config.root, LoopAgentConfig): + return LoopAgent.from_config(config.root) + else: + raise ValueError("Unsupported config type") + + +@working_in_progress("_load_config_from_path is not ready for use.") +def _load_config_from_path(config_path: str) -> AgentConfig: + """Load an agent's configuration from a YAML file. + + Args: + config_path: Path to the YAML config file. Both relative and absolute + paths are accepted. + + Returns: + The loaded and validated AgentConfig object. + + Raises: + FileNotFoundError: If config file doesn't exist. + ValidationError: If config file's content is invalid YAML. + """ + config_path = Path(config_path) + + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + + with open(config_path, "r", encoding="utf-8") as f: + config_data = yaml.safe_load(f) + + return AgentConfig.model_validate(config_data) diff --git a/src/google/adk/cli/utils/agent_loader.py b/src/google/adk/cli/utils/agent_loader.py index ca81bd23..1e206846 100644 --- a/src/google/adk/cli/utils/agent_loader.py +++ b/src/google/adk/cli/utils/agent_loader.py @@ -16,11 +16,16 @@ from __future__ import annotations import importlib import logging +import os import sys from typing import Optional +from pydantic import ValidationError + from . import envs +from ...agents import config_agent_utils from ...agents.base_agent import BaseAgent +from ...utils.feature_decorator import working_in_progress logger = logging.getLogger("google_adk." + __name__) @@ -34,6 +39,8 @@ class AgentLoader: agents_dir/{agent_name}.py (with root_agent defined in the module) c) {agent_name} as a package name agents_dir/{agent_name}/__init__.py (with root_agent in the package) + d) {agent_name} as a YAML config folder: + agents_dir/{agent_name}/root_agent.yaml defines the root agent """ @@ -128,6 +135,29 @@ class AgentLoader: return None + @working_in_progress("_load_from_yaml_config is not ready for use.") + def _load_from_yaml_config(self, agent_name: str) -> Optional[BaseAgent]: + # Load from the config file at agents_dir/{agent_name}/root_agent.yaml + config_path = os.path.join(self.agents_dir, agent_name, "root_agent.yaml") + try: + agent = config_agent_utils.from_config(config_path) + logger.info("Loaded root agent for %s from %s", agent_name, config_path) + return agent + except FileNotFoundError: + logger.debug("Config file %s not found.", config_path) + return None + except ValidationError as e: + logger.error("Config file %s is invalid YAML.", config_path) + raise e + except Exception as e: + if hasattr(e, "msg"): + e.msg = f"Fail to load '{config_path}' config. " + e.msg + raise e + e.args = ( + f"Fail to load '{config_path}' config. {e.args[0] if e.args else ''}", + ) + e.args[1:] + raise e + def _perform_load(self, agent_name: str) -> BaseAgent: """Internal logic to load an agent""" # Add self.agents_dir to sys.path @@ -145,10 +175,14 @@ class AgentLoader: if root_agent := self._load_from_submodule(agent_name): return root_agent + if root_agent := self._load_from_yaml_config(agent_name): + return root_agent + # If no root_agent was found by any pattern raise ValueError( f"No root_agent found for '{agent_name}'. Searched in" - f" '{agent_name}.agent.root_agent', '{agent_name}.root_agent'." + f" '{agent_name}.agent.root_agent', '{agent_name}.root_agent' and" + f" '{agent_name}/root_agent.yaml'." f" Ensure '{self.agents_dir}/{agent_name}' is structured correctly," " an .env file can be loaded if present, and a root_agent is" " exposed." diff --git a/tests/unittests/cli/utils/test_agent_loader.py b/tests/unittests/cli/utils/test_agent_loader.py index dafac921..2b68f3cc 100644 --- a/tests/unittests/cli/utils/test_agent_loader.py +++ b/tests/unittests/cli/utils/test_agent_loader.py @@ -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)