From a3ff21eb0bf67ddd3e1cfcd83ec31d692d15271a Mon Sep 17 00:00:00 2001 From: "Wei Sun (Jack)" Date: Fri, 25 Jul 2025 09:55:17 -0700 Subject: [PATCH] feat(config): Adds CustomAgentConfig to support user-defined agents in config PiperOrigin-RevId: 787148485 --- src/google/adk/agents/agent_config.py | 22 ++- src/google/adk/agents/base_agent.py | 102 +---------- src/google/adk/agents/base_agent_config.py | 160 ++++++++++++++++++ src/google/adk/agents/config_agent_utils.py | 2 +- src/google/adk/agents/llm_agent.py | 7 +- src/google/adk/agents/loop_agent.py | 9 +- src/google/adk/agents/parallel_agent.py | 7 +- src/google/adk/agents/sequential_agent.py | 7 +- tests/unittests/agents/test_agent_config.py | 123 ++++++++++++++ .../unittests/cli/utils/test_agent_loader.py | 3 +- 10 files changed, 331 insertions(+), 111 deletions(-) create mode 100644 src/google/adk/agents/base_agent_config.py create mode 100644 tests/unittests/agents/test_agent_config.py diff --git a/src/google/adk/agents/agent_config.py b/src/google/adk/agents/agent_config.py index f32f0f96..1d9f36fb 100644 --- a/src/google/adk/agents/agent_config.py +++ b/src/google/adk/agents/agent_config.py @@ -14,11 +14,14 @@ from __future__ import annotations +from typing import Any from typing import Union +from pydantic import Discriminator from pydantic import RootModel from ..utils.feature_decorator import working_in_progress +from .base_agent import BaseAgentConfig from .llm_agent import LlmAgentConfig from .loop_agent import LoopAgentConfig from .parallel_agent import ParallelAgentConfig @@ -30,9 +33,26 @@ ConfigsUnion = Union[ LoopAgentConfig, ParallelAgentConfig, SequentialAgentConfig, + BaseAgentConfig, ] +def agent_config_discriminator(v: Any): + if isinstance(v, dict): + agent_class = v.get("agent_class", "LlmAgent") + if agent_class in [ + "LlmAgent", + "LoopAgent", + "ParallelAgent", + "SequentialAgent", + ]: + return agent_class + else: + return "BaseAgent" + + raise ValueError(f"Invalid agent config: {v}") + + # Use a RootModel to represent the agent directly at the top level. # The `discriminator` is applied to the union within the RootModel. @working_in_progress("AgentConfig is not ready for use.") @@ -43,4 +63,4 @@ class AgentConfig(RootModel[ConfigsUnion]): # Pydantic v2 requires this for discriminated unions on RootModel # This tells the model to look at the 'agent_class' field of the input # data to decide which model from the `ConfigsUnion` to use. - discriminator = "agent_class" + discriminator = Discriminator(agent_config_discriminator) diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index 80b58ff1..1ea63f28 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -21,8 +21,6 @@ from typing import Awaitable from typing import Callable from typing import Dict from typing import final -from typing import List -from typing import Literal from typing import Mapping from typing import Optional from typing import Type @@ -36,14 +34,13 @@ from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field from pydantic import field_validator -from pydantic import model_validator from typing_extensions import override from typing_extensions import TypeAlias from ..events.event import Event from ..utils.feature_decorator import working_in_progress +from .base_agent_config import BaseAgentConfig from .callback_context import CallbackContext -from .common_configs import CodeConfig if TYPE_CHECKING: from .invocation_context import InvocationContext @@ -535,100 +532,3 @@ class BaseAgent(BaseModel): config.after_agent_callbacks ) return cls(**kwargs) - - -class SubAgentConfig(BaseModel): - """The config for a sub-agent.""" - - model_config = ConfigDict(extra='forbid') - - config: Optional[str] = None - """The YAML config file path of the sub-agent. - - Only one of `config` or `code` can be set. - - Example: - - ``` - sub_agents: - - config: search_agent.yaml - - config: my_library/my_custom_agent.yaml - ``` - """ - - code: Optional[str] = None - """The agent instance defined in the code. - - Only one of `config` or `code` can be set. - - Example: - - For the following agent defined in Python code: - - ``` - # my_library/custom_agents.py - from google.adk.agents.llm_agent import LlmAgent - - my_custom_agent = LlmAgent( - name="my_custom_agent", - instruction="You are a helpful custom agent.", - model="gemini-2.0-flash", - ) - ``` - - The yaml config should be: - - ``` - sub_agents: - - code: my_library.custom_agents.my_custom_agent - ``` - """ - - @model_validator(mode='after') - def validate_exactly_one_field(self): - code_provided = self.code is not None - config_provided = self.config is not None - - if code_provided and config_provided: - raise ValueError('Only one of code or config should be provided') - if not code_provided and not config_provided: - raise ValueError('Exactly one of code or config must be provided') - - return self - - -@working_in_progress('BaseAgentConfig is not ready for use.') -class BaseAgentConfig(BaseModel): - """The config for the YAML schema of a BaseAgent. - - Do not use this class directly. It's the base class for all agent configs. - """ - - model_config = ConfigDict(extra='forbid') - - agent_class: Literal['BaseAgent'] = 'BaseAgent' - """Required. The class of the agent. The value is used to differentiate - among different agent classes.""" - - name: str - """Required. The name of the agent.""" - - description: str = '' - """Optional. The description of the agent.""" - - sub_agents: Optional[List[SubAgentConfig]] = None - """Optional. The sub-agents of the agent.""" - - before_agent_callbacks: Optional[List[CodeConfig]] = None - """Optional. The before_agent_callbacks of the agent. - - Example: - - ``` - before_agent_callbacks: - - name: my_library.security_callbacks.before_agent_callback - ``` - """ - - after_agent_callbacks: Optional[List[CodeConfig]] = None - """Optional. The after_agent_callbacks of the agent.""" diff --git a/src/google/adk/agents/base_agent_config.py b/src/google/adk/agents/base_agent_config.py new file mode 100644 index 00000000..04ef0e7d --- /dev/null +++ b/src/google/adk/agents/base_agent_config.py @@ -0,0 +1,160 @@ +# 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 inspect +from typing import Any +from typing import AsyncGenerator +from typing import Awaitable +from typing import Callable +from typing import Dict +from typing import final +from typing import List +from typing import Literal +from typing import Mapping +from typing import Optional +from typing import Type +from typing import TYPE_CHECKING +from typing import TypeVar +from typing import Union + +from google.genai import types +from opentelemetry import trace +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field +from pydantic import field_validator +from pydantic import model_validator +from typing_extensions import override +from typing_extensions import TypeAlias + +from ..events.event import Event +from ..utils.feature_decorator import working_in_progress +from .callback_context import CallbackContext +from .common_configs import CodeConfig + +if TYPE_CHECKING: + from .invocation_context import InvocationContext + + +TBaseAgentConfig = TypeVar('TBaseAgentConfig', bound='BaseAgentConfig') + + +class SubAgentConfig(BaseModel): + """The config for a sub-agent.""" + + model_config = ConfigDict(extra='forbid') + + config: Optional[str] = None + """The YAML config file path of the sub-agent. + + Only one of `config` or `code` can be set. + + Example: + + ``` + sub_agents: + - config: search_agent.yaml + - config: my_library/my_custom_agent.yaml + ``` + """ + + code: Optional[str] = None + """The agent instance defined in the code. + + Only one of `config` or `code` can be set. + + Example: + + For the following agent defined in Python code: + + ``` + # my_library/custom_agents.py + from google.adk.agents.llm_agent import LlmAgent + + my_custom_agent = LlmAgent( + name="my_custom_agent", + instruction="You are a helpful custom agent.", + model="gemini-2.0-flash", + ) + ``` + + The yaml config should be: + + ``` + sub_agents: + - code: my_library.custom_agents.my_custom_agent + ``` + """ + + @model_validator(mode='after') + def validate_exactly_one_field(self): + code_provided = self.code is not None + config_provided = self.config is not None + + if code_provided and config_provided: + raise ValueError('Only one of code or config should be provided') + if not code_provided and not config_provided: + raise ValueError('Exactly one of code or config must be provided') + + return self + + +@working_in_progress('BaseAgentConfig is not ready for use.') +class BaseAgentConfig(BaseModel): + """The config for the YAML schema of a BaseAgent. + + Do not use this class directly. It's the base class for all agent configs. + """ + + model_config = ConfigDict( + extra='allow', + ) + + agent_class: Union[Literal['BaseAgent'], str] = 'BaseAgent' + """Required. The class of the agent. The value is used to differentiate + among different agent classes.""" + + name: str + """Required. The name of the agent.""" + + description: str = '' + """Optional. The description of the agent.""" + + sub_agents: Optional[List[SubAgentConfig]] = None + """Optional. The sub-agents of the agent.""" + + before_agent_callbacks: Optional[List[CodeConfig]] = None + """Optional. The before_agent_callbacks of the agent. + + Example: + + ``` + before_agent_callbacks: + - name: my_library.security_callbacks.before_agent_callback + ``` + """ + + after_agent_callbacks: Optional[List[CodeConfig]] = None + """Optional. The after_agent_callbacks of the agent.""" + + def to_agent_config( + self, custom_agent_config_cls: Type[TBaseAgentConfig] + ) -> TBaseAgentConfig: + """Converts this config to the concrete agent config type. + + NOTE: this is for ADK framework use only. + """ + return custom_agent_config_cls.model_validate(self.model_dump()) diff --git a/src/google/adk/agents/config_agent_utils.py b/src/google/adk/agents/config_agent_utils.py index 00b12ff6..5da6a211 100644 --- a/src/google/adk/agents/config_agent_utils.py +++ b/src/google/adk/agents/config_agent_utils.py @@ -24,7 +24,7 @@ import yaml from ..utils.feature_decorator import working_in_progress from .agent_config import AgentConfig from .base_agent import BaseAgent -from .base_agent import SubAgentConfig +from .base_agent_config import SubAgentConfig from .common_configs import CodeConfig from .llm_agent import LlmAgent from .llm_agent import LlmAgentConfig diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index c20d2696..b193f577 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -29,6 +29,7 @@ from typing import Union from google.genai import types from pydantic import BaseModel +from pydantic import ConfigDict from pydantic import Field from pydantic import field_validator from pydantic import model_validator @@ -53,7 +54,7 @@ from ..tools.function_tool import FunctionTool from ..tools.tool_context import ToolContext from ..utils.feature_decorator import working_in_progress from .base_agent import BaseAgent -from .base_agent import BaseAgentConfig +from .base_agent_config import BaseAgentConfig from .callback_context import CallbackContext from .common_configs import CodeConfig from .invocation_context import InvocationContext @@ -607,6 +608,10 @@ Agent: TypeAlias = LlmAgent class LlmAgentConfig(BaseAgentConfig): """The config for the YAML schema of a LlmAgent.""" + model_config = ConfigDict( + extra='forbid', + ) + agent_class: Literal['LlmAgent', ''] = 'LlmAgent' """The value is used to uniquely identify the LlmAgent class. If it is empty, it is by default an LlmAgent.""" diff --git a/src/google/adk/agents/loop_agent.py b/src/google/adk/agents/loop_agent.py index e5822786..4fc0f166 100644 --- a/src/google/adk/agents/loop_agent.py +++ b/src/google/adk/agents/loop_agent.py @@ -16,20 +16,19 @@ from __future__ import annotations -from typing import Any from typing import AsyncGenerator -from typing import Dict from typing import Literal from typing import Optional from typing import Type +from pydantic import ConfigDict from typing_extensions import override from ..agents.invocation_context import InvocationContext from ..events.event import Event from ..utils.feature_decorator import working_in_progress from .base_agent import BaseAgent -from .base_agent import BaseAgentConfig +from .base_agent_config import BaseAgentConfig class LoopAgent(BaseAgent): @@ -90,6 +89,10 @@ class LoopAgent(BaseAgent): class LoopAgentConfig(BaseAgentConfig): """The config for the YAML schema of a LoopAgent.""" + model_config = ConfigDict( + extra='forbid', + ) + agent_class: Literal['LoopAgent'] = 'LoopAgent' max_iterations: Optional[int] = None diff --git a/src/google/adk/agents/parallel_agent.py b/src/google/adk/agents/parallel_agent.py index 36034056..99f4f486 100644 --- a/src/google/adk/agents/parallel_agent.py +++ b/src/google/adk/agents/parallel_agent.py @@ -21,10 +21,11 @@ from typing import AsyncGenerator from typing import Literal from typing import Type +from pydantic import ConfigDict from typing_extensions import override -from ..agents.base_agent import BaseAgentConfig from ..agents.base_agent import working_in_progress +from ..agents.base_agent_config import BaseAgentConfig from ..agents.invocation_context import InvocationContext from ..events.event import Event from .base_agent import BaseAgent @@ -131,4 +132,8 @@ class ParallelAgent(BaseAgent): class ParallelAgentConfig(BaseAgentConfig): """The config for the YAML schema of a ParallelAgent.""" + model_config = ConfigDict( + extra='forbid', + ) + agent_class: Literal['ParallelAgent'] = 'ParallelAgent' diff --git a/src/google/adk/agents/sequential_agent.py b/src/google/adk/agents/sequential_agent.py index 51dff22c..4ca282de 100644 --- a/src/google/adk/agents/sequential_agent.py +++ b/src/google/adk/agents/sequential_agent.py @@ -20,10 +20,11 @@ from typing import AsyncGenerator from typing import Literal from typing import Type +from pydantic import ConfigDict from typing_extensions import override -from ..agents.base_agent import BaseAgentConfig from ..agents.base_agent import working_in_progress +from ..agents.base_agent_config import BaseAgentConfig from ..agents.invocation_context import InvocationContext from ..events.event import Event from .base_agent import BaseAgent @@ -94,4 +95,8 @@ class SequentialAgent(BaseAgent): class SequentialAgentConfig(BaseAgentConfig): """The config for the YAML schema of a SequentialAgent.""" + model_config = ConfigDict( + extra='forbid', + ) + agent_class: Literal['SequentialAgent'] = 'SequentialAgent' diff --git a/tests/unittests/agents/test_agent_config.py b/tests/unittests/agents/test_agent_config.py new file mode 100644 index 00000000..38954eab --- /dev/null +++ b/tests/unittests/agents/test_agent_config.py @@ -0,0 +1,123 @@ +from typing import Literal + +from google.adk.agents.agent_config import AgentConfig +from google.adk.agents.agent_config import LlmAgentConfig +from google.adk.agents.agent_config import LoopAgentConfig +from google.adk.agents.agent_config import ParallelAgentConfig +from google.adk.agents.agent_config import SequentialAgentConfig +from google.adk.agents.base_agent_config import BaseAgentConfig +import yaml + + +def test_agent_config_discriminator_default_is_llm_agent(): + yaml_content = """\ +name: search_agent +model: gemini-2.0-flash +description: a sample description +instruction: a fake instruction +tools: + - name: google_search +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, LlmAgentConfig) + assert config.root.agent_class == "LlmAgent" + + +def test_agent_config_discriminator_llm_agent(): + yaml_content = """\ +agent_class: LlmAgent +name: search_agent +model: gemini-2.0-flash +description: a sample description +instruction: a fake instruction +tools: + - name: google_search +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, LlmAgentConfig) + assert config.root.agent_class == "LlmAgent" + + +def test_agent_config_discriminator_loop_agent(): + yaml_content = """\ +agent_class: LoopAgent +name: CodePipelineAgent +description: Executes a sequence of code writing, reviewing, and refactoring. +sub_agents: + - config: sub_agents/code_writer_agent.yaml + - config: sub_agents/code_reviewer_agent.yaml + - config: sub_agents/code_refactorer_agent.yaml +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, LoopAgentConfig) + assert config.root.agent_class == "LoopAgent" + + +def test_agent_config_discriminator_parallel_agent(): + yaml_content = """\ +agent_class: ParallelAgent +name: CodePipelineAgent +description: Executes a sequence of code writing, reviewing, and refactoring. +sub_agents: + - config: sub_agents/code_writer_agent.yaml + - config: sub_agents/code_reviewer_agent.yaml + - config: sub_agents/code_refactorer_agent.yaml +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, ParallelAgentConfig) + assert config.root.agent_class == "ParallelAgent" + + +def test_agent_config_discriminator_sequential_agent(): + yaml_content = """\ +agent_class: SequentialAgent +name: CodePipelineAgent +description: Executes a sequence of code writing, reviewing, and refactoring. +sub_agents: + - config: sub_agents/code_writer_agent.yaml + - config: sub_agents/code_reviewer_agent.yaml + - config: sub_agents/code_refactorer_agent.yaml +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, SequentialAgentConfig) + assert config.root.agent_class == "SequentialAgent" + + +def test_agent_config_discriminator_custom_agent(): + class MyCustomAgentConfig(BaseAgentConfig): + agent_class: Literal["mylib.agents.MyCustomAgent"] = ( + "mylib.agents.MyCustomAgent" + ) + other_field: str + + yaml_content = """\ +agent_class: mylib.agents.MyCustomAgent +name: CodePipelineAgent +description: Executes a sequence of code writing, reviewing, and refactoring. +other_field: other value +""" + config_data = yaml.safe_load(yaml_content) + + config = AgentConfig.model_validate(config_data) + + assert isinstance(config.root, BaseAgentConfig) + assert config.root.agent_class == "mylib.agents.MyCustomAgent" + assert config.root.model_extra == {"other_field": "other value"} + + my_custom_config = config.root.to_agent_config(MyCustomAgentConfig) + assert my_custom_config.other_field == "other value" diff --git a/tests/unittests/cli/utils/test_agent_loader.py b/tests/unittests/cli/utils/test_agent_loader.py index 2b68f3cc..81d6baae 100644 --- a/tests/unittests/cli/utils/test_agent_loader.py +++ b/tests/unittests/cli/utils/test_agent_loader.py @@ -555,8 +555,7 @@ class TestAgentLoader: # Create invalid YAML content with wrong field name invalid_yaml_content = dedent(""" - agent_type: LlmAgent - name: invalid_yaml_test_agent + not_exist_field: invalid_yaml_test_agent model: gemini-2.0-flash instruction: You are a test agent with invalid YAML """)