chore(config): Reimplements AgentConfig with pydantic v2 convention and allow all possible values for agent_class field in all Agent Configs

All below are valid values now.

```
agent_class: LlmAgent
agent_class: google.adk.agents.LlmAgent
agent_class: google.adk.agents.llm_agent.LlmAgent
```
PiperOrigin-RevId: 800228114
This commit is contained in:
Wei Sun (Jack)
2025-08-27 17:07:53 -07:00
committed by Copybara-Service
parent f743c29d00
commit 3bc2d77b4d
7 changed files with 300 additions and 130 deletions
+34 -27
View File
@@ -14,53 +14,60 @@
from __future__ import annotations
from typing import Annotated
from typing import Any
from typing import get_args
from typing import Union
from pydantic import Discriminator
from pydantic import RootModel
from pydantic import Tag
from ..utils.feature_decorator import experimental
from .base_agent import BaseAgentConfig
from .base_agent_config import BaseAgentConfig
from .llm_agent_config import LlmAgentConfig
from .loop_agent_config import LoopAgentConfig
from .parallel_agent import ParallelAgentConfig
from .sequential_agent import SequentialAgentConfig
from .parallel_agent_config import ParallelAgentConfig
from .sequential_agent_config import SequentialAgentConfig
# A discriminated union of all possible agent configurations.
ConfigsUnion = Union[
LlmAgentConfig,
LoopAgentConfig,
ParallelAgentConfig,
SequentialAgentConfig,
BaseAgentConfig,
]
_ADK_AGENT_CLASSES: set[str] = {
"LlmAgent",
"LoopAgent",
"ParallelAgent",
"SequentialAgent",
}
def agent_config_discriminator(v: Any):
def agent_config_discriminator(v: Any) -> str:
"""Discriminator function that returns the tag name for Pydantic."""
if isinstance(v, dict):
agent_class = v.get("agent_class", "LlmAgent")
if agent_class in [
"LlmAgent",
"LoopAgent",
"ParallelAgent",
"SequentialAgent",
]:
agent_class: str = v.get("agent_class", "LlmAgent")
# Look up the agent_class in our dynamically built mapping
if agent_class in _ADK_AGENT_CLASSES:
return agent_class
else:
return "BaseAgent"
# For non ADK agent classes, use BaseAgent to handle it.
return "BaseAgent"
raise ValueError(f"Invalid agent config: {v}")
# A discriminated union of all possible agent configurations.
ConfigsUnion = Annotated[
Union[
Annotated[LlmAgentConfig, Tag("LlmAgent")],
Annotated[LoopAgentConfig, Tag("LoopAgent")],
Annotated[ParallelAgentConfig, Tag("ParallelAgent")],
Annotated[SequentialAgentConfig, Tag("SequentialAgent")],
Annotated[BaseAgentConfig, Tag("BaseAgent")],
],
Discriminator(agent_config_discriminator),
]
# Use a RootModel to represent the agent directly at the top level.
# The `discriminator` is applied to the union within the RootModel.
@experimental
class AgentConfig(RootModel[ConfigsUnion]):
"""The config for the YAML schema to create an agent."""
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 = Discriminator(agent_config_discriminator)
@@ -681,12 +681,29 @@
"EnterpriseWebSearch": {
"additionalProperties": false,
"description": "Tool to search public web data, powered by Vertex AI Search and Sec4 compliance.",
"properties": {},
"properties": {
"excludeDomains": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional. List of domains to be excluded from the search results. The default limit is 2000 domains.",
"title": "Excludedomains"
}
},
"title": "EnterpriseWebSearch",
"type": "object"
},
"Environment": {
"description": "Required. The environment being operated.",
"description": "The environment being operated.",
"enum": [
"ENVIRONMENT_UNSPECIFIED",
"ENVIRONMENT_BROWSER"
@@ -1476,20 +1493,7 @@
"$ref": "#/$defs/Content"
},
{
"items": {
"anyOf": [
{
"$ref": "#/$defs/File"
},
{
"$ref": "#/$defs/Part"
},
{
"type": "string"
}
]
},
"type": "array"
"type": "string"
},
{
"$ref": "#/$defs/File"
@@ -1498,7 +1502,20 @@
"$ref": "#/$defs/Part"
},
{
"type": "string"
"items": {
"anyOf": [
{
"type": "string"
},
{
"$ref": "#/$defs/File"
},
{
"$ref": "#/$defs/Part"
}
]
},
"type": "array"
},
{
"type": "null"
@@ -1830,10 +1847,10 @@
"speechConfig": {
"anyOf": [
{
"$ref": "#/$defs/SpeechConfig"
"type": "string"
},
{
"type": "string"
"$ref": "#/$defs/SpeechConfig"
},
{
"type": "null"
@@ -1999,6 +2016,22 @@
],
"default": null,
"description": "Optional. Filter search results to a specific time range.\n If customers set a start time, they must set an end time (and vice versa).\n "
},
"excludeDomains": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional. List of domains to be excluded from the search results.\n The default limit is 2000 domains.",
"title": "Excludedomains"
}
},
"title": "GoogleSearch",
@@ -2356,10 +2389,6 @@
"agent_class": {
"default": "LlmAgent",
"description": "The value is used to uniquely identify the LlmAgent class. If it is empty, it is by default an LlmAgent.",
"enum": [
"LlmAgent",
""
],
"title": "Agent Class",
"type": "string"
},
@@ -2618,7 +2647,6 @@
"description": "The config for the YAML schema of a LoopAgent.",
"properties": {
"agent_class": {
"const": "LoopAgent",
"default": "LoopAgent",
"description": "The value is used to uniquely identify the LoopAgent class.",
"title": "Agent Class",
@@ -2774,7 +2802,6 @@
"description": "The config for the YAML schema of a ParallelAgent.",
"properties": {
"agent_class": {
"const": "ParallelAgent",
"default": "ParallelAgent",
"description": "The value is used to uniquely identify the ParallelAgent class.",
"title": "Agent Class",
@@ -3680,7 +3707,6 @@
"description": "The config for the YAML schema of a SequentialAgent.",
"properties": {
"agent_class": {
"const": "SequentialAgent",
"default": "SequentialAgent",
"description": "The value is used to uniquely identify the SequentialAgent class.",
"title": "Agent Class",
@@ -4413,6 +4439,18 @@
"default": null,
"description": "Optional. Tool to support URL context retrieval."
},
"computerUse": {
"anyOf": [
{
"$ref": "#/$defs/ToolComputerUse"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional. Tool to support the model interacting directly with the\n computer. If enabled, it automatically populates computer-use specific\n Function Declarations."
},
"codeExecution": {
"anyOf": [
{
@@ -4424,18 +4462,6 @@
],
"default": null,
"description": "Optional. CodeExecution tool type. Enables the model to execute code as part of generation."
},
"computerUse": {
"anyOf": [
{
"$ref": "#/$defs/ToolComputerUse"
},
{
"type": "null"
}
],
"default": null,
"description": "Optional. Tool to support the model interacting directly with the computer. If enabled, it automatically populates computer-use specific Function Declarations."
}
},
"title": "Tool",
@@ -4556,7 +4582,8 @@
"type": "object"
}
},
"anyOf": [
"description": "The config for the YAML schema to create an agent.",
"oneOf": [
{
"$ref": "#/$defs/LlmAgentConfig"
},
@@ -4573,6 +4600,5 @@
"$ref": "#/$defs/BaseAgentConfig"
}
],
"description": "The config for the YAML schema to create an agent.",
"title": "AgentConfig"
}
+1 -1
View File
@@ -37,7 +37,7 @@ class LlmAgentConfig(BaseAgentConfig):
extra='forbid',
)
agent_class: Literal['LlmAgent', ''] = Field(
agent_class: str = Field(
default='LlmAgent',
description=(
'The value is used to uniquely identify the LlmAgent class. If it is'
+1 -2
View File
@@ -16,7 +16,6 @@
from __future__ import annotations
from typing import Literal
from typing import Optional
from pydantic import ConfigDict
@@ -34,7 +33,7 @@ class LoopAgentConfig(BaseAgentConfig):
extra='forbid',
)
agent_class: Literal['LoopAgent'] = Field(
agent_class: str = Field(
default='LoopAgent',
description='The value is used to uniquely identify the LoopAgent class.',
)
@@ -16,8 +16,6 @@
from __future__ import annotations
from typing import Literal
from pydantic import ConfigDict
from pydantic import Field
@@ -30,12 +28,12 @@ class ParallelAgentConfig(BaseAgentConfig):
"""The config for the YAML schema of a ParallelAgent."""
model_config = ConfigDict(
extra='forbid',
extra="forbid",
)
agent_class: Literal['ParallelAgent'] = Field(
default='ParallelAgent',
agent_class: str = Field(
default="ParallelAgent",
description=(
'The value is used to uniquely identify the ParallelAgent class.'
"The value is used to uniquely identify the ParallelAgent class."
),
)
@@ -16,8 +16,6 @@
from __future__ import annotations
from typing import Literal
from pydantic import ConfigDict
from pydantic import Field
@@ -30,12 +28,12 @@ class SequentialAgentConfig(BaseAgentConfig):
"""The config for the YAML schema of a SequentialAgent."""
model_config = ConfigDict(
extra='forbid',
extra="forbid",
)
agent_class: Literal['SequentialAgent'] = Field(
default='SequentialAgent',
agent_class: str = Field(
default="SequentialAgent",
description=(
'The value is used to uniquely identify the SequentialAgent class.'
"The value is used to uniquely identify the SequentialAgent class."
),
)
+190 -48
View File
@@ -12,18 +12,23 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path
from typing import Literal
from typing import Type
from google.adk.agents import config_agent_utils
from google.adk.agents.agent_config import AgentConfig
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.base_agent_config import BaseAgentConfig
from google.adk.agents.llm_agent_config import LlmAgentConfig
from google.adk.agents.loop_agent_config import LoopAgentConfig
from google.adk.agents.parallel_agent_config import ParallelAgentConfig
from google.adk.agents.sequential_agent_config import SequentialAgentConfig
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent
import pytest
import yaml
def test_agent_config_discriminator_default_is_llm_agent():
def test_agent_config_discriminator_default_is_llm_agent(tmp_path: Path):
yaml_content = """\
name: search_agent
model: gemini-2.0-flash
@@ -32,17 +37,29 @@ instruction: a fake instruction
tools:
- name: google_search
"""
config_data = yaml.safe_load(yaml_content)
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(config_data)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(config.root, LlmAgentConfig)
assert isinstance(agent, LlmAgent)
assert config.root.agent_class == "LlmAgent"
def test_agent_config_discriminator_llm_agent():
yaml_content = """\
agent_class: LlmAgent
@pytest.mark.parametrize(
"agent_class_value",
[
"LlmAgent",
"google.adk.agents.LlmAgent",
"google.adk.agents.llm_agent.LlmAgent",
],
)
def test_agent_config_discriminator_llm_agent(
agent_class_value: str, tmp_path: Path
):
yaml_content = f"""\
agent_class: {agent_class_value}
name: search_agent
model: gemini-2.0-flash
description: a sample description
@@ -50,66 +67,191 @@ instruction: a fake instruction
tools:
- name: google_search
"""
config_data = yaml.safe_load(yaml_content)
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(config_data)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(config.root, LlmAgentConfig)
assert config.root.agent_class == "LlmAgent"
assert isinstance(agent, LlmAgent)
assert config.root.agent_class == agent_class_value
def test_agent_config_discriminator_loop_agent():
yaml_content = """\
agent_class: LoopAgent
@pytest.mark.parametrize(
"agent_class_value",
[
"LoopAgent",
"google.adk.agents.LoopAgent",
"google.adk.agents.loop_agent.LoopAgent",
],
)
def test_agent_config_discriminator_loop_agent(
agent_class_value: str, tmp_path: Path
):
yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
sub_agents: []
"""
config_data = yaml.safe_load(yaml_content)
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(config_data)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(config.root, LoopAgentConfig)
assert config.root.agent_class == "LoopAgent"
assert isinstance(agent, LoopAgent)
assert config.root.agent_class == agent_class_value
def test_agent_config_discriminator_parallel_agent():
yaml_content = """\
agent_class: ParallelAgent
@pytest.mark.parametrize(
"agent_class_value",
[
"ParallelAgent",
"google.adk.agents.ParallelAgent",
"google.adk.agents.parallel_agent.ParallelAgent",
],
)
def test_agent_config_discriminator_parallel_agent(
agent_class_value: str, tmp_path: Path
):
yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
sub_agents: []
"""
config_data = yaml.safe_load(yaml_content)
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(config_data)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(config.root, ParallelAgentConfig)
assert config.root.agent_class == "ParallelAgent"
assert isinstance(agent, ParallelAgent)
assert config.root.agent_class == agent_class_value
def test_agent_config_discriminator_sequential_agent():
yaml_content = """\
agent_class: SequentialAgent
@pytest.mark.parametrize(
"agent_class_value",
[
"SequentialAgent",
"google.adk.agents.SequentialAgent",
"google.adk.agents.sequential_agent.SequentialAgent",
],
)
def test_agent_config_discriminator_sequential_agent(
agent_class_value: str, tmp_path: Path
):
yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
sub_agents: []
"""
config_data = yaml.safe_load(yaml_content)
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(config_data)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(config.root, SequentialAgentConfig)
assert config.root.agent_class == "SequentialAgent"
assert isinstance(agent, SequentialAgent)
assert config.root.agent_class == agent_class_value
@pytest.mark.parametrize(
("agent_class_value", "expected_agent_type"),
[
("LoopAgent", LoopAgent),
("google.adk.agents.LoopAgent", LoopAgent),
("google.adk.agents.loop_agent.LoopAgent", LoopAgent),
("ParallelAgent", ParallelAgent),
("google.adk.agents.ParallelAgent", ParallelAgent),
("google.adk.agents.parallel_agent.ParallelAgent", ParallelAgent),
("SequentialAgent", SequentialAgent),
("google.adk.agents.SequentialAgent", SequentialAgent),
("google.adk.agents.sequential_agent.SequentialAgent", SequentialAgent),
],
)
def test_agent_config_discriminator_with_sub_agents(
agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
):
# Create sub-agent config files
sub_agent_dir = tmp_path / "sub_agents"
sub_agent_dir.mkdir()
sub_agent_config = """\
name: sub_agent_{index}
model: gemini-2.0-flash
description: a sub agent
instruction: sub agent instruction
"""
(sub_agent_dir / "sub_agent1.yaml").write_text(
sub_agent_config.format(index=1)
)
(sub_agent_dir / "sub_agent2.yaml").write_text(
sub_agent_config.format(index=2)
)
yaml_content = f"""\
agent_class: {agent_class_value}
name: main_agent
description: main agent with sub agents
sub_agents:
- config_path: sub_agents/sub_agent1.yaml
- config_path: sub_agents/sub_agent2.yaml
"""
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(agent, expected_agent_type)
assert config.root.agent_class == agent_class_value
@pytest.mark.parametrize(
("agent_class_value", "expected_agent_type"),
[
("LlmAgent", LlmAgent),
("google.adk.agents.LlmAgent", LlmAgent),
("google.adk.agents.llm_agent.LlmAgent", LlmAgent),
],
)
def test_agent_config_discriminator_llm_agent_with_sub_agents(
agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
):
# Create sub-agent config files
sub_agent_dir = tmp_path / "sub_agents"
sub_agent_dir.mkdir()
sub_agent_config = """\
name: sub_agent_{index}
model: gemini-2.0-flash
description: a sub agent
instruction: sub agent instruction
"""
(sub_agent_dir / "sub_agent1.yaml").write_text(
sub_agent_config.format(index=1)
)
(sub_agent_dir / "sub_agent2.yaml").write_text(
sub_agent_config.format(index=2)
)
yaml_content = f"""\
agent_class: {agent_class_value}
name: main_agent
model: gemini-2.0-flash
description: main agent with sub agents
instruction: main agent instruction
sub_agents:
- config_path: sub_agents/sub_agent1.yaml
- config_path: sub_agents/sub_agent2.yaml
"""
config_file = tmp_path / "test_config.yaml"
config_file.write_text(yaml_content)
config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
agent = config_agent_utils.from_config(str(config_file))
assert isinstance(agent, expected_agent_type)
assert config.root.agent_class == agent_class_value
def test_agent_config_discriminator_custom_agent():