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

PiperOrigin-RevId: 786456046

Author: Jacksunwei

src/google/adk/agents/agent_config.py

@@ -14,25 +14,88 @@
 
 from __future__ import annotations
 
+from typing import Any
+from typing import Literal
+from typing import Type
+from typing import TypeVar
 from typing import Union
 
+from pydantic import ConfigDict
+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
 from .sequential_agent import SequentialAgentConfig
 
+TBaseAgentConfig = TypeVar("TBaseAgentConfig", bound=BaseAgentConfig)
+
+
+@working_in_progress("AgentConfig is not ready for use.")
+class CustomAgentConfig(BaseAgentConfig):
+  """Used for configs for user-defined custom agents."""
+
+  model_config = ConfigDict(
+      extra="allow",
+  )
+  agent_class: Union[Literal["CustomAgent"], str] = "CustomAgent"
+
+  def to_agent_config(
+      self, custom_agent_config_cls: Type[TBaseAgentConfig]
+  ) -> TBaseAgentConfig:
+    """Converts the this config to the concrete agent config type.
+
+    ```
+    # In my_custom_agent.py
+    class MyCustomAgentConfig(BaseAgentConfig):
+      agent_class: Literal["mylib.agents.MyCustomAgent"] = "mylib.agents.MyCustomAgent"
+      other_field: str
+
+    class MyCustomAgent(BaseAgent):
+      ...
+
+      def from_config(
+          cls: Type[MyCustomAgent],
+          config: CustomAgentConfig,
+          config_abs_path: str,
+      ) -> MyCustomAgent:
+        my_custom_agent_config = config.to_agent_config(MyCustomAgentConfig)
+        # use my_custom_agent_config for the remaining logic...
+
+    ```
+    """
+    return custom_agent_config_cls.model_validate(self.model_dump())
+
+
 # A discriminated union of all possible agent configurations.
 ConfigsUnion = Union[
     LlmAgentConfig,
     LoopAgentConfig,
     ParallelAgentConfig,
     SequentialAgentConfig,
+    CustomAgentConfig,
 ]
 
 
+def agent_config_discriminator(v: Any):
+  if isinstance(v, dict):
+    agent_class = v.get("agent_class", "LlmAgentConfig")
+    if agent_class in [
+        "LlmAgent",
+        "LoopAgent",
+        "ParallelAgent",
+        "SequentialAgent",
+    ]:
+      return agent_class
+    else:
+      return "CustomAgent"
+
+  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 +106,4 @@ class Config:
     # 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)

tests/unittests/agents/test_agent_config.py

@@ -0,0 +1,124 @@
+from typing import Literal
+
+from google.adk.agents.agent_config import AgentConfig
+from google.adk.agents.agent_config import CustomAgentConfig
+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 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 valud
+"""
+  config_data = yaml.safe_load(yaml_content)
+
+  config = AgentConfig.model_validate(config_data)
+
+  assert isinstance(config.root, CustomAgentConfig)
+  assert config.root.agent_class == "mylib.agents.MyCustomAgent"
+  assert config.root.model_extra == {"other_field": "other valud"}
+
+  my_custom_config = config.root.to_agent_config(MyCustomAgentConfig)
+  assert my_custom_config.other_field == "other valud"