Skip to content

Commit 6fe06e5

Browse files
Jacksunweicopybara-github
Jacksunwei
authored and
copybara-github
committed
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: 798341605
1 parent 0b89f18 commit 6fe06e5

File tree

7 files changed

+208
-80
lines changed

7 files changed

+208
-80
lines changed

‎src/google/adk/agents/agent_config.py

Lines changed: 58 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -14,53 +14,84 @@
1414

1515
from __future__ import annotations
1616

17+
from typing import Annotated
1718
from typing import Any
19+
from typing import get_args
1820
from typing import Union
1921

2022
from pydantic import Discriminator
2123
from pydantic import RootModel
24+
from pydantic import Tag
2225

2326
from ..utils.feature_decorator import experimental
24-
from .base_agent import BaseAgentConfig
27+
from .base_agent_config import BaseAgentConfig
2528
from .llm_agent_config import LlmAgentConfig
2629
from .loop_agent_config import LoopAgentConfig
27-
from .parallel_agent import ParallelAgentConfig
28-
from .sequential_agent import SequentialAgentConfig
30+
from .parallel_agent_config import ParallelAgentConfig
31+
from .sequential_agent_config import SequentialAgentConfig
2932

30-
# A discriminated union of all possible agent configurations.
31-
ConfigsUnion = Union[
32-
LlmAgentConfig,
33-
LoopAgentConfig,
34-
ParallelAgentConfig,
35-
SequentialAgentConfig,
36-
BaseAgentConfig,
37-
]
33+
34+
def _get_agent_class_values_for_config(config_class) -> set[str]:
35+
"""Get all allowed agent_class values from a config class."""
36+
if agent_class_field := config_class.model_fields.get("agent_class"):
37+
return set(filter(bool, get_args(agent_class_field.annotation)))
38+
return set()
39+
40+
41+
def _build_agent_class_mapping() -> dict[str, str]:
42+
"""Build a mapping from agent_class values to tag names."""
43+
config_mappings = [
44+
(LlmAgentConfig, "LlmAgent"),
45+
(LoopAgentConfig, "LoopAgent"),
46+
(ParallelAgentConfig, "ParallelAgent"),
47+
(SequentialAgentConfig, "SequentialAgent"),
48+
]
49+
50+
return {
51+
value: tag_name
52+
for config_class, tag_name in config_mappings
53+
for value in _get_agent_class_values_for_config(config_class)
54+
}
55+
56+
57+
# Cache the agent class mapping
58+
_AGENT_CLASS_MAPPING = _build_agent_class_mapping()
3859

3960

40-
def agent_config_discriminator(v: Any):
61+
def agent_config_discriminator(v: Any) -> str:
62+
"""Discriminator function that returns the tag name for Pydantic."""
4163
if isinstance(v, dict):
4264
agent_class = v.get("agent_class", "LlmAgent")
43-
if agent_class in [
44-
"LlmAgent",
45-
"LoopAgent",
46-
"ParallelAgent",
47-
"SequentialAgent",
48-
]:
49-
return agent_class
50-
else:
51-
return "BaseAgent"
65+
66+
# Handle empty string case - default to LlmAgent
67+
if not agent_class:
68+
return "LlmAgent"
69+
70+
# Look up the agent_class in our dynamically built mapping
71+
if agent_class in _AGENT_CLASS_MAPPING:
72+
return _AGENT_CLASS_MAPPING[agent_class]
73+
74+
# For unknown agent classes, return "BaseAgent" to use BaseAgentConfig
75+
return "BaseAgent"
5276

5377
raise ValueError(f"Invalid agent config: {v}")
5478

5579

80+
# A discriminated union of all possible agent configurations.
81+
ConfigsUnion = Annotated[
82+
Union[
83+
Annotated[LlmAgentConfig, Tag("LlmAgent")],
84+
Annotated[LoopAgentConfig, Tag("LoopAgent")],
85+
Annotated[ParallelAgentConfig, Tag("ParallelAgent")],
86+
Annotated[SequentialAgentConfig, Tag("SequentialAgent")],
87+
Annotated[BaseAgentConfig, Tag("BaseAgent")],
88+
],
89+
Discriminator(agent_config_discriminator),
90+
]
91+
92+
5693
# Use a RootModel to represent the agent directly at the top level.
5794
# The `discriminator` is applied to the union within the RootModel.
5895
@experimental
5996
class AgentConfig(RootModel[ConfigsUnion]):
6097
"""The config for the YAML schema to create an agent."""
61-
62-
class Config:
63-
# Pydantic v2 requires this for discriminated unions on RootModel
64-
# This tells the model to look at the 'agent_class' field of the input
65-
# data to decide which model from the `ConfigsUnion` to use.
66-
discriminator = Discriminator(agent_config_discriminator)

‎src/google/adk/agents/config_schemas/AgentConfig.json

Lines changed: 74 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -681,12 +681,29 @@
681681
"EnterpriseWebSearch": {
682682
"additionalProperties": false,
683683
"description": "Tool to search public web data, powered by Vertex AI Search and Sec4 compliance.",
684-
"properties": {},
684+
"properties": {
685+
"excludeDomains": {
686+
"anyOf": [
687+
{
688+
"items": {
689+
"type": "string"
690+
},
691+
"type": "array"
692+
},
693+
{
694+
"type": "null"
695+
}
696+
],
697+
"default": null,
698+
"description": "Optional. List of domains to be excluded from the search results. The default limit is 2000 domains.",
699+
"title": "Excludedomains"
700+
}
701+
},
685702
"title": "EnterpriseWebSearch",
686703
"type": "object"
687704
},
688705
"Environment": {
689-
"description": "Required. The environment being operated.",
706+
"description": "The environment being operated.",
690707
"enum": [
691708
"ENVIRONMENT_UNSPECIFIED",
692709
"ENVIRONMENT_BROWSER"
@@ -1475,31 +1492,31 @@
14751492
{
14761493
"$ref": "#/$defs/Content"
14771494
},
1495+
{
1496+
"type": "string"
1497+
},
1498+
{
1499+
"$ref": "#/$defs/File"
1500+
},
1501+
{
1502+
"$ref": "#/$defs/Part"
1503+
},
14781504
{
14791505
"items": {
14801506
"anyOf": [
14811507
{
1482-
"$ref": "#/$defs/File"
1508+
"type": "string"
14831509
},
14841510
{
1485-
"$ref": "#/$defs/Part"
1511+
"$ref": "#/$defs/File"
14861512
},
14871513
{
1488-
"type": "string"
1514+
"$ref": "#/$defs/Part"
14891515
}
14901516
]
14911517
},
14921518
"type": "array"
14931519
},
1494-
{
1495-
"$ref": "#/$defs/File"
1496-
},
1497-
{
1498-
"$ref": "#/$defs/Part"
1499-
},
1500-
{
1501-
"type": "string"
1502-
},
15031520
{
15041521
"type": "null"
15051522
}
@@ -1830,10 +1847,10 @@
18301847
"speechConfig": {
18311848
"anyOf": [
18321849
{
1833-
"$ref": "#/$defs/SpeechConfig"
1850+
"type": "string"
18341851
},
18351852
{
1836-
"type": "string"
1853+
"$ref": "#/$defs/SpeechConfig"
18371854
},
18381855
{
18391856
"type": "null"
@@ -1999,6 +2016,22 @@
19992016
],
20002017
"default": null,
20012018
"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 "
2019+
},
2020+
"excludeDomains": {
2021+
"anyOf": [
2022+
{
2023+
"items": {
2024+
"type": "string"
2025+
},
2026+
"type": "array"
2027+
},
2028+
{
2029+
"type": "null"
2030+
}
2031+
],
2032+
"default": null,
2033+
"description": "Optional. List of domains to be excluded from the search results.\n The default limit is 2000 domains.",
2034+
"title": "Excludedomains"
20022035
}
20032036
},
20042037
"title": "GoogleSearch",
@@ -2358,6 +2391,8 @@
23582391
"description": "The value is used to uniquely identify the LlmAgent class. If it is empty, it is by default an LlmAgent.",
23592392
"enum": [
23602393
"LlmAgent",
2394+
"google.adk.agents.LlmAgent",
2395+
"google.adk.agents.llm_agent.LlmAgent",
23612396
""
23622397
],
23632398
"title": "Agent Class",
@@ -2618,9 +2653,13 @@
26182653
"description": "The config for the YAML schema of a LoopAgent.",
26192654
"properties": {
26202655
"agent_class": {
2621-
"const": "LoopAgent",
26222656
"default": "LoopAgent",
26232657
"description": "The value is used to uniquely identify the LoopAgent class.",
2658+
"enum": [
2659+
"LoopAgent",
2660+
"google.adk.agents.LoopAgent",
2661+
"google.adk.agents.loop_agent.LoopAgent"
2662+
],
26242663
"title": "Agent Class",
26252664
"type": "string"
26262665
},
@@ -2774,9 +2813,13 @@
27742813
"description": "The config for the YAML schema of a ParallelAgent.",
27752814
"properties": {
27762815
"agent_class": {
2777-
"const": "ParallelAgent",
27782816
"default": "ParallelAgent",
27792817
"description": "The value is used to uniquely identify the ParallelAgent class.",
2818+
"enum": [
2819+
"ParallelAgent",
2820+
"google.adk.agents.ParallelAgent",
2821+
"google.adk.agents.parallel_agent.ParallelAgent"
2822+
],
27802823
"title": "Agent Class",
27812824
"type": "string"
27822825
},
@@ -3680,9 +3723,13 @@
36803723
"description": "The config for the YAML schema of a SequentialAgent.",
36813724
"properties": {
36823725
"agent_class": {
3683-
"const": "SequentialAgent",
36843726
"default": "SequentialAgent",
36853727
"description": "The value is used to uniquely identify the SequentialAgent class.",
3728+
"enum": [
3729+
"SequentialAgent",
3730+
"google.adk.agents.SequentialAgent",
3731+
"google.adk.agents.sequential_agent.SequentialAgent"
3732+
],
36863733
"title": "Agent Class",
36873734
"type": "string"
36883735
},
@@ -4413,29 +4460,29 @@
44134460
"default": null,
44144461
"description": "Optional. Tool to support URL context retrieval."
44154462
},
4416-
"codeExecution": {
4463+
"computerUse": {
44174464
"anyOf": [
44184465
{
4419-
"$ref": "#/$defs/ToolCodeExecution"
4466+
"$ref": "#/$defs/ToolComputerUse"
44204467
},
44214468
{
44224469
"type": "null"
44234470
}
44244471
],
44254472
"default": null,
4426-
"description": "Optional. CodeExecution tool type. Enables the model to execute code as part of generation."
4473+
"description": "Optional. Tool to support the model interacting directly with the\n computer. If enabled, it automatically populates computer-use specific\n Function Declarations."
44274474
},
4428-
"computerUse": {
4475+
"codeExecution": {
44294476
"anyOf": [
44304477
{
4431-
"$ref": "#/$defs/ToolComputerUse"
4478+
"$ref": "#/$defs/ToolCodeExecution"
44324479
},
44334480
{
44344481
"type": "null"
44354482
}
44364483
],
44374484
"default": null,
4438-
"description": "Optional. Tool to support the model interacting directly with the computer. If enabled, it automatically populates computer-use specific Function Declarations."
4485+
"description": "Optional. CodeExecution tool type. Enables the model to execute code as part of generation."
44394486
}
44404487
},
44414488
"title": "Tool",
@@ -4556,7 +4603,8 @@
45564603
"type": "object"
45574604
}
45584605
},
4559-
"anyOf": [
4606+
"description": "The config for the YAML schema to create an agent.",
4607+
"oneOf": [
45604608
{
45614609
"$ref": "#/$defs/LlmAgentConfig"
45624610
},
@@ -4573,6 +4621,5 @@
45734621
"$ref": "#/$defs/BaseAgentConfig"
45744622
}
45754623
],
4576-
"description": "The config for the YAML schema to create an agent.",
45774624
"title": "AgentConfig"
45784625
}

‎src/google/adk/agents/llm_agent_config.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,12 @@ class LlmAgentConfig(BaseAgentConfig):
3737
extra='forbid',
3838
)
3939

40-
agent_class: Literal['LlmAgent', ''] = Field(
40+
agent_class: Literal[
41+
'LlmAgent',
42+
'google.adk.agents.LlmAgent',
43+
'google.adk.agents.llm_agent.LlmAgent',
44+
'',
45+
] = Field(
4146
default='LlmAgent',
4247
description=(
4348
'The value is used to uniquely identify the LlmAgent class. If it is'

‎src/google/adk/agents/loop_agent_config.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@ class LoopAgentConfig(BaseAgentConfig):
3434
extra='forbid',
3535
)
3636

37-
agent_class: Literal['LoopAgent'] = Field(
37+
agent_class: Literal[
38+
'LoopAgent',
39+
'google.adk.agents.LoopAgent',
40+
'google.adk.agents.loop_agent.LoopAgent',
41+
] = Field(
3842
default='LoopAgent',
3943
description='The value is used to uniquely identify the LoopAgent class.',
4044
)

‎src/google/adk/agents/parallel_agent_config.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,11 @@ class ParallelAgentConfig(BaseAgentConfig):
3333
extra='forbid',
3434
)
3535

36-
agent_class: Literal['ParallelAgent'] = Field(
36+
agent_class: Literal[
37+
'ParallelAgent',
38+
'google.adk.agents.ParallelAgent',
39+
'google.adk.agents.parallel_agent.ParallelAgent',
40+
] = Field(
3741
default='ParallelAgent',
3842
description=(
3943
'The value is used to uniquely identify the ParallelAgent class.'

‎src/google/adk/agents/sequential_agent_config.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,11 @@ class SequentialAgentConfig(BaseAgentConfig):
3333
extra='forbid',
3434
)
3535

36-
agent_class: Literal['SequentialAgent'] = Field(
36+
agent_class: Literal[
37+
'SequentialAgent',
38+
'google.adk.agents.SequentialAgent',
39+
'google.adk.agents.sequential_agent.SequentialAgent',
40+
] = Field(
3741
default='SequentialAgent',
3842
description=(
3943
'The value is used to uniquely identify the SequentialAgent class.'

0 commit comments

Comments
 (0)