Skip to content

Commit 2bb2041

Browse files
Jacksunweicopybara-github
authored andcommitted
feat(config): Adds BaseAgent.config_type field to indicate the config for the current agent and removes if-else branches against LlmAgent/LoopAgent/... in config_agent_utils::from_config
This makes the logic work with any user-defined agent with user-defined XxxAgentConfig. PiperOrigin-RevId: 789845354
1 parent 86a4487 commit 2bb2041

File tree

6 files changed

+77
-18
lines changed

6 files changed

+77
-18
lines changed

src/google/adk/agents/base_agent.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from typing import AsyncGenerator
2020
from typing import Awaitable
2121
from typing import Callable
22+
from typing import ClassVar
2223
from typing import Dict
2324
from typing import final
2425
from typing import Mapping
@@ -75,6 +76,22 @@ class BaseAgent(BaseModel):
7576
)
7677
"""The pydantic model config."""
7778

79+
config_type: ClassVar[type[BaseAgentConfig]] = BaseAgentConfig
80+
"""The config type for this agent.
81+
82+
Sub-classes should override this to specify their own config type.
83+
84+
Example:
85+
86+
```
87+
class MyAgentConfig(BaseAgentConfig):
88+
my_field: str = ''
89+
90+
class MyAgent(BaseAgent):
91+
config_type: ClassVar[type[BaseAgentConfig]] = MyAgentConfig
92+
```
93+
"""
94+
7895
name: str
7996
"""The agent's name.
8097

src/google/adk/agents/config_agent_utils.py

Lines changed: 41 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from __future__ import annotations
1616

1717
import importlib
18+
import inspect
1819
import os
1920
from typing import Any
2021
from typing import List
@@ -24,16 +25,9 @@
2425
from ..utils.feature_decorator import working_in_progress
2526
from .agent_config import AgentConfig
2627
from .base_agent import BaseAgent
28+
from .base_agent_config import BaseAgentConfig
2729
from .common_configs import AgentRefConfig
2830
from .common_configs import CodeConfig
29-
from .llm_agent import LlmAgent
30-
from .llm_agent_config import LlmAgentConfig
31-
from .loop_agent import LoopAgent
32-
from .loop_agent_config import LoopAgentConfig
33-
from .parallel_agent import ParallelAgent
34-
from .parallel_agent import ParallelAgentConfig
35-
from .sequential_agent import SequentialAgent
36-
from .sequential_agent import SequentialAgentConfig
3731

3832

3933
@working_in_progress("from_config is not ready for use.")
@@ -53,17 +47,36 @@ def from_config(config_path: str) -> BaseAgent:
5347
"""
5448
abs_path = os.path.abspath(config_path)
5549
config = _load_config_from_path(abs_path)
56-
57-
if isinstance(config.root, LlmAgentConfig):
58-
return LlmAgent.from_config(config.root, abs_path)
59-
elif isinstance(config.root, LoopAgentConfig):
60-
return LoopAgent.from_config(config.root, abs_path)
61-
elif isinstance(config.root, ParallelAgentConfig):
62-
return ParallelAgent.from_config(config.root, abs_path)
63-
elif isinstance(config.root, SequentialAgentConfig):
64-
return SequentialAgent.from_config(config.root, abs_path)
50+
agent_config = config.root
51+
52+
# pylint: disable=unidiomatic-typecheck Needs exact class matching.
53+
if type(agent_config) is BaseAgentConfig:
54+
# Resolve the concrete agent config for user-defined agent classes.
55+
agent_class = _resolve_agent_class(agent_config.agent_class)
56+
agent_config = agent_class.config_type.model_validate(
57+
agent_config.model_dump()
58+
)
59+
return agent_class.from_config(agent_config, abs_path)
6560
else:
66-
raise ValueError("Unsupported config type")
61+
# For built-in agent classes, no need to re-validate.
62+
agent_class = _resolve_agent_class(agent_config.agent_class)
63+
return agent_class.from_config(agent_config, abs_path)
64+
65+
66+
def _resolve_agent_class(agent_class: str) -> type[BaseAgent]:
67+
"""Resolve the agent class from its fully qualified name."""
68+
agent_class_name = agent_class or "LlmAgent"
69+
if "." not in agent_class_name:
70+
agent_class_name = f"google.adk.agents.{agent_class_name}"
71+
72+
agent_class = _resolve_fully_qualified_name(agent_class_name)
73+
if inspect.isclass(agent_class) and issubclass(agent_class, BaseAgent):
74+
return agent_class
75+
76+
raise ValueError(
77+
f"Invalid agent class `{agent_class_name}`. It must be a subclass of"
78+
" BaseAgent."
79+
)
6780

6881

6982
@working_in_progress("_load_config_from_path is not ready for use.")
@@ -90,6 +103,16 @@ def _load_config_from_path(config_path: str) -> AgentConfig:
90103
return AgentConfig.model_validate(config_data)
91104

92105

106+
@working_in_progress("_resolve_fully_qualified_name is not ready for use.")
107+
def _resolve_fully_qualified_name(name: str) -> Any:
108+
try:
109+
module_path, obj_name = name.rsplit(".", 1)
110+
module = importlib.import_module(module_path)
111+
return getattr(module, obj_name)
112+
except Exception as e:
113+
raise ValueError(f"Invalid fully qualified name: {name}") from e
114+
115+
93116
@working_in_progress("resolve_agent_reference is not ready for use.")
94117
def resolve_agent_reference(
95118
ref_config: AgentRefConfig, referencing_agent_config_abs_path: str

src/google/adk/agents/llm_agent.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from typing import AsyncGenerator
2323
from typing import Awaitable
2424
from typing import Callable
25+
from typing import ClassVar
2526
from typing import Literal
2627
from typing import Optional
2728
from typing import Type
@@ -55,6 +56,7 @@
5556
from ..tools.tool_context import ToolContext
5657
from ..utils.feature_decorator import working_in_progress
5758
from .base_agent import BaseAgent
59+
from .base_agent_config import BaseAgentConfig
5860
from .callback_context import CallbackContext
5961
from .common_configs import CodeConfig
6062
from .invocation_context import InvocationContext
@@ -131,6 +133,9 @@ class LlmAgent(BaseAgent):
131133
When not set, the agent will inherit the model from its ancestor.
132134
"""
133135

136+
config_type: ClassVar[type[BaseAgentConfig]] = LlmAgentConfig
137+
"""The config type for this agent."""
138+
134139
instruction: Union[str, InstructionProvider] = ''
135140
"""Instructions for the LLM model, guiding the agent's behavior."""
136141

src/google/adk/agents/loop_agent.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from __future__ import annotations
1818

1919
from typing import AsyncGenerator
20+
from typing import ClassVar
2021
from typing import Optional
2122
from typing import Type
2223

@@ -26,6 +27,7 @@
2627
from ..events.event import Event
2728
from ..utils.feature_decorator import working_in_progress
2829
from .base_agent import BaseAgent
30+
from .base_agent_config import BaseAgentConfig
2931
from .loop_agent_config import LoopAgentConfig
3032

3133

@@ -36,6 +38,9 @@ class LoopAgent(BaseAgent):
3638
reached, the loop agent will stop.
3739
"""
3840

41+
config_type: ClassVar[type[BaseAgentConfig]] = LoopAgentConfig
42+
"""The config type for this agent."""
43+
3944
max_iterations: Optional[int] = None
4045
"""The maximum number of iterations to run the loop agent.
4146

src/google/adk/agents/parallel_agent.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,15 @@
1818

1919
import asyncio
2020
from typing import AsyncGenerator
21+
from typing import ClassVar
2122
from typing import Type
2223

2324
from typing_extensions import override
2425

2526
from ..events.event import Event
2627
from ..utils.feature_decorator import working_in_progress
2728
from .base_agent import BaseAgent
29+
from .base_agent_config import BaseAgentConfig
2830
from .invocation_context import InvocationContext
2931
from .parallel_agent_config import ParallelAgentConfig
3032

@@ -95,6 +97,9 @@ class ParallelAgent(BaseAgent):
9597
- Generating multiple responses for review by a subsequent evaluation agent.
9698
"""
9799

100+
config_type: ClassVar[type[BaseAgentConfig]] = ParallelAgentConfig
101+
"""The config type for this agent."""
102+
98103
@override
99104
async def _run_async_impl(
100105
self, ctx: InvocationContext

src/google/adk/agents/sequential_agent.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from ..events.event import Event
2525
from ..utils.feature_decorator import working_in_progress
2626
from .base_agent import BaseAgent
27+
from .base_agent import BaseAgentConfig
2728
from .invocation_context import InvocationContext
2829
from .llm_agent import LlmAgent
2930
from .sequential_agent_config import SequentialAgentConfig
@@ -32,6 +33,9 @@
3233
class SequentialAgent(BaseAgent):
3334
"""A shell agent that runs its sub-agents in sequence."""
3435

36+
config_type: Type[BaseAgentConfig] = SequentialAgentConfig
37+
"""The config type for this agent."""
38+
3539
@override
3640
async def _run_async_impl(
3741
self, ctx: InvocationContext

0 commit comments

Comments
 (0)