feat(config): Adds CustomAgentConfig to support user-defined agents in config

PiperOrigin-RevId: 787148485

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)

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."""

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())

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

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."""

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

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'

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'

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"

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
       """)