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:
Wei Sun (Jack)
committed by
Copybara-Service
Copybara-Service
parent
f743c29d00
commit
3bc2d77b4d
@@ -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