mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
f743c29d00
commit
3bc2d77b4d
@@ -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"
|
||||
}
|
||||
@@ -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'
|
||||
|
||||
@@ -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."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user