Skip to content

Commit fa21fd0

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 2dd432c commit fa21fd0

File tree

7 files changed

+180
-95
lines changed

7 files changed

+180
-95
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: 57 additions & 31 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",
@@ -2356,10 +2389,6 @@
23562389
"agent_class": {
23572390
"default": "LlmAgent",
23582391
"description": "The value is used to uniquely identify the LlmAgent class. If it is empty, it is by default an LlmAgent.",
2359-
"enum": [
2360-
"LlmAgent",
2361-
""
2362-
],
23632392
"title": "Agent Class",
23642393
"type": "string"
23652394
},
@@ -2618,7 +2647,6 @@
26182647
"description": "The config for the YAML schema of a LoopAgent.",
26192648
"properties": {
26202649
"agent_class": {
2621-
"const": "LoopAgent",
26222650
"default": "LoopAgent",
26232651
"description": "The value is used to uniquely identify the LoopAgent class.",
26242652
"title": "Agent Class",
@@ -2774,7 +2802,6 @@
27742802
"description": "The config for the YAML schema of a ParallelAgent.",
27752803
"properties": {
27762804
"agent_class": {
2777-
"const": "ParallelAgent",
27782805
"default": "ParallelAgent",
27792806
"description": "The value is used to uniquely identify the ParallelAgent class.",
27802807
"title": "Agent Class",
@@ -3680,7 +3707,6 @@
36803707
"description": "The config for the YAML schema of a SequentialAgent.",
36813708
"properties": {
36823709
"agent_class": {
3683-
"const": "SequentialAgent",
36843710
"default": "SequentialAgent",
36853711
"description": "The value is used to uniquely identify the SequentialAgent class.",
36863712
"title": "Agent Class",
@@ -4413,29 +4439,29 @@
44134439
"default": null,
44144440
"description": "Optional. Tool to support URL context retrieval."
44154441
},
4416-
"codeExecution": {
4442+
"computerUse": {
44174443
"anyOf": [
44184444
{
4419-
"$ref": "#/$defs/ToolCodeExecution"
4445+
"$ref": "#/$defs/ToolComputerUse"
44204446
},
44214447
{
44224448
"type": "null"
44234449
}
44244450
],
44254451
"default": null,
4426-
"description": "Optional. CodeExecution tool type. Enables the model to execute code as part of generation."
4452+
"description": "Optional. Tool to support the model interacting directly with the\n computer. If enabled, it automatically populates computer-use specific\n Function Declarations."
44274453
},
4428-
"computerUse": {
4454+
"codeExecution": {
44294455
"anyOf": [
44304456
{
4431-
"$ref": "#/$defs/ToolComputerUse"
4457+
"$ref": "#/$defs/ToolCodeExecution"
44324458
},
44334459
{
44344460
"type": "null"
44354461
}
44364462
],
44374463
"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."
4464+
"description": "Optional. CodeExecution tool type. Enables the model to execute code as part of generation."
44394465
}
44404466
},
44414467
"title": "Tool",
@@ -4556,7 +4582,8 @@
45564582
"type": "object"
45574583
}
45584584
},
4559-
"anyOf": [
4585+
"description": "The config for the YAML schema to create an agent.",
4586+
"oneOf": [
45604587
{
45614588
"$ref": "#/$defs/LlmAgentConfig"
45624589
},
@@ -4573,6 +4600,5 @@
45734600
"$ref": "#/$defs/BaseAgentConfig"
45744601
}
45754602
],
4576-
"description": "The config for the YAML schema to create an agent.",
45774603
"title": "AgentConfig"
45784604
}

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

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

40-
agent_class: Literal['LlmAgent', ''] = Field(
40+
agent_class: str = Field(
4141
default='LlmAgent',
4242
description=(
4343
'The value is used to uniquely identify the LlmAgent class. If it is'

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
from __future__ import annotations
1818

19-
from typing import Literal
2019
from typing import Optional
2120

2221
from pydantic import ConfigDict
@@ -34,7 +33,7 @@ class LoopAgentConfig(BaseAgentConfig):
3433
extra='forbid',
3534
)
3635

37-
agent_class: Literal['LoopAgent'] = Field(
36+
agent_class: str = Field(
3837
default='LoopAgent',
3938
description='The value is used to uniquely identify the LoopAgent class.',
4039
)

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616

1717
from __future__ import annotations
1818

19-
from typing import Literal
20-
2119
from pydantic import ConfigDict
2220
from pydantic import Field
2321

@@ -30,12 +28,12 @@ class ParallelAgentConfig(BaseAgentConfig):
3028
"""The config for the YAML schema of a ParallelAgent."""
3129

3230
model_config = ConfigDict(
33-
extra='forbid',
31+
extra="forbid",
3432
)
3533

36-
agent_class: Literal['ParallelAgent'] = Field(
37-
default='ParallelAgent',
34+
agent_class: str = Field(
35+
default="ParallelAgent",
3836
description=(
39-
'The value is used to uniquely identify the ParallelAgent class.'
37+
"The value is used to uniquely identify the ParallelAgent class."
4038
),
4139
)

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616

1717
from __future__ import annotations
1818

19-
from typing import Literal
20-
2119
from pydantic import ConfigDict
2220
from pydantic import Field
2321

@@ -30,12 +28,12 @@ class SequentialAgentConfig(BaseAgentConfig):
3028
"""The config for the YAML schema of a SequentialAgent."""
3129

3230
model_config = ConfigDict(
33-
extra='forbid',
31+
extra="forbid",
3432
)
3533

36-
agent_class: Literal['SequentialAgent'] = Field(
37-
default='SequentialAgent',
34+
agent_class: str = Field(
35+
default="SequentialAgent",
3836
description=(
39-
'The value is used to uniquely identify the SequentialAgent class.'
37+
"The value is used to uniquely identify the SequentialAgent class."
4038
),
4139
)

0 commit comments

Comments
 (0)