From 53803522b6acbdfd929ca93565a27477ddeae855 Mon Sep 17 00:00:00 2001 From: "Wei Sun (Jack)" Date: Tue, 5 Aug 2025 23:52:50 -0700 Subject: [PATCH] refactor(config): Makes `BaseAgent.from_config` a final method and let sub-class to optionally override `_parse_config` to update kwargs if needed This ensures that the pydantic hooks (e.g. model_validators) are triggered correctly. PiperOrigin-RevId: 791545704 --- src/google/adk/agents/base_agent.py | 45 ++++++++++++++++++---- src/google/adk/agents/llm_agent.py | 46 +++++++++++------------ src/google/adk/agents/loop_agent.py | 17 +++++---- src/google/adk/agents/parallel_agent.py | 12 +----- src/google/adk/agents/sequential_agent.py | 10 ----- 5 files changed, 72 insertions(+), 58 deletions(-) diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py index 1d2d8c02..98f7b125 100644 --- a/src/google/adk/agents/base_agent.py +++ b/src/google/adk/agents/base_agent.py @@ -504,8 +504,8 @@ class BaseAgent(BaseModel): sub_agent.parent_agent = self return self + @final @classmethod - @working_in_progress('BaseAgent.from_config is not ready for use.') def from_config( cls: Type[SelfAgent], config: BaseAgentConfig, @@ -513,11 +513,8 @@ class BaseAgent(BaseModel): ) -> SelfAgent: """Creates an agent from a config. - This method converts fields in a config to the corresponding - fields in an agent. - - Child classes should re-implement this method to support loading from their - custom config types. + If sub-classes uses a custom agent config, override `_from_config_kwargs` + method to return an updated kwargs for agent construstor. Args: config: The config to create the agent from. @@ -527,6 +524,40 @@ class BaseAgent(BaseModel): Returns: The created agent. """ + kwargs = cls.__create_kwargs(config, config_abs_path) + kwargs = cls._parse_config(config, config_abs_path, kwargs) + return cls(**kwargs) + + @classmethod + def _parse_config( + cls: Type[SelfAgent], + config: BaseAgentConfig, + config_abs_path: str, + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: + """Parses the config and returns updated kwargs to construct the agent. + + Sub-classes should override this method to use a custome agent config class. + + Args: + config: The config to parse. + config_abs_path: The absolute path to the config file that contains the + agent config. + kwargs: The keyword arguments used for agent constructor. + + Returns: + The updated keyword arguments used for agent constructor. + """ + return kwargs + + @classmethod + def __create_kwargs( + cls, + config: BaseAgentConfig, + config_abs_path: str, + ) -> Dict[str, Any]: + """Creates kwargs for the fields of BaseAgent.""" + from .config_agent_utils import resolve_agent_reference from .config_agent_utils import resolve_callbacks @@ -549,4 +580,4 @@ class BaseAgent(BaseModel): kwargs['after_agent_callback'] = resolve_callbacks( config.after_agent_callbacks ) - return cls(**kwargs) + return kwargs diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 86ace586..db957f93 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -17,12 +17,12 @@ from __future__ import annotations import importlib import inspect import logging -import os from typing import Any from typing import AsyncGenerator from typing import Awaitable from typing import Callable from typing import ClassVar +from typing import Dict from typing import Literal from typing import Optional from typing import Type @@ -46,7 +46,6 @@ from ..models.llm_request import LlmRequest from ..models.llm_response import LlmResponse from ..models.registry import LLMRegistry from ..planners.base_planner import BasePlanner -from ..tools.agent_tool import AgentTool from ..tools.base_tool import BaseTool from ..tools.base_tool import ToolConfig from ..tools.base_toolset import BaseToolset @@ -56,7 +55,6 @@ from ..utils.feature_decorator import working_in_progress from .base_agent import BaseAgent from .base_agent_config import BaseAgentConfig from .callback_context import CallbackContext -from .common_configs import CodeConfig from .invocation_context import InvocationContext from .llm_agent_config import LlmAgentConfig from .readonly_context import ReadonlyContext @@ -586,53 +584,55 @@ class LlmAgent(BaseAgent): return resolved_tools - @classmethod @override - @working_in_progress('LlmAgent.from_config is not ready for use.') - def from_config( + @classmethod + def _parse_config( cls: Type[LlmAgent], config: LlmAgentConfig, config_abs_path: str, - ) -> LlmAgent: + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: from .config_agent_utils import resolve_callbacks from .config_agent_utils import resolve_code_reference - agent = super().from_config(config, config_abs_path) if config.model: - agent.model = config.model + kwargs['model'] = config.model if config.instruction: - agent.instruction = config.instruction + kwargs['instruction'] = config.instruction if config.disallow_transfer_to_parent: - agent.disallow_transfer_to_parent = config.disallow_transfer_to_parent + kwargs['disallow_transfer_to_parent'] = config.disallow_transfer_to_parent if config.disallow_transfer_to_peers: - agent.disallow_transfer_to_peers = config.disallow_transfer_to_peers + kwargs['disallow_transfer_to_peers'] = config.disallow_transfer_to_peers if config.include_contents != 'default': - agent.include_contents = config.include_contents + kwargs['include_contents'] = config.include_contents if config.input_schema: - agent.input_schema = resolve_code_reference(config.input_schema) + kwargs['input_schema'] = resolve_code_reference(config.input_schema) if config.output_schema: - agent.output_schema = resolve_code_reference(config.output_schema) + kwargs['output_schema'] = resolve_code_reference(config.output_schema) if config.output_key: - agent.output_key = config.output_key + kwargs['output_key'] = config.output_key if config.tools: - agent.tools = cls._resolve_tools(config.tools, config_abs_path) + kwargs['tools'] = cls._resolve_tools(config.tools, config_abs_path) if config.before_model_callbacks: - agent.before_model_callback = resolve_callbacks( + kwargs['before_model_callback'] = resolve_callbacks( config.before_model_callbacks ) if config.after_model_callbacks: - agent.after_model_callback = resolve_callbacks( + kwargs['after_model_callback'] = resolve_callbacks( config.after_model_callbacks ) if config.before_tool_callbacks: - agent.before_tool_callback = resolve_callbacks( + kwargs['before_tool_callback'] = resolve_callbacks( config.before_tool_callbacks ) if config.after_tool_callbacks: - agent.after_tool_callback = resolve_callbacks(config.after_tool_callbacks) + kwargs['after_tool_callback'] = resolve_callbacks( + config.after_tool_callbacks + ) if config.generate_content_config: - agent.generate_content_config = config.generate_content_config - return agent + kwargs['generate_content_config'] = config.generate_content_config + + return kwargs Agent: TypeAlias = LlmAgent diff --git a/src/google/adk/agents/loop_agent.py b/src/google/adk/agents/loop_agent.py index 812361a3..de4f3438 100644 --- a/src/google/adk/agents/loop_agent.py +++ b/src/google/adk/agents/loop_agent.py @@ -16,8 +16,10 @@ from __future__ import annotations +from typing import Any from typing import AsyncGenerator from typing import ClassVar +from typing import Dict from typing import Optional from typing import Type @@ -74,15 +76,14 @@ class LoopAgent(BaseAgent): raise NotImplementedError('This is not supported yet for LoopAgent.') yield # AsyncGenerator requires having at least one yield statement - @classmethod @override - @working_in_progress('LoopAgent.from_config is not ready for use.') - def from_config( - cls: Type[LoopAgent], + @classmethod + def _parse_config( + cls: type[LoopAgent], config: LoopAgentConfig, config_abs_path: str, - ) -> LoopAgent: - agent = super().from_config(config, config_abs_path) + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: if config.max_iterations: - agent.max_iterations = config.max_iterations - return agent + kwargs['max_iterations'] = config.max_iterations + return kwargs diff --git a/src/google/adk/agents/parallel_agent.py b/src/google/adk/agents/parallel_agent.py index f8c4c28e..bb8e0a46 100644 --- a/src/google/adk/agents/parallel_agent.py +++ b/src/google/adk/agents/parallel_agent.py @@ -17,8 +17,10 @@ from __future__ import annotations import asyncio +from typing import Any from typing import AsyncGenerator from typing import ClassVar +from typing import Dict from typing import Type from typing_extensions import override @@ -119,13 +121,3 @@ class ParallelAgent(BaseAgent): ) -> AsyncGenerator[Event, None]: raise NotImplementedError('This is not supported yet for ParallelAgent.') yield # AsyncGenerator requires having at least one yield statement - - @classmethod - @override - @working_in_progress('ParallelAgent.from_config is not ready for use.') - def from_config( - cls: Type[ParallelAgent], - config: ParallelAgentConfig, - config_abs_path: str, - ) -> ParallelAgent: - return super().from_config(config, config_abs_path) diff --git a/src/google/adk/agents/sequential_agent.py b/src/google/adk/agents/sequential_agent.py index c0c832ff..10d1e7c2 100644 --- a/src/google/adk/agents/sequential_agent.py +++ b/src/google/adk/agents/sequential_agent.py @@ -81,13 +81,3 @@ class SequentialAgent(BaseAgent): for sub_agent in self.sub_agents: async for event in sub_agent.run_live(ctx): yield event - - @classmethod - @override - @working_in_progress('SequentialAgent.from_config is not ready for use.') - def from_config( - cls: Type[SequentialAgent], - config: SequentialAgentConfig, - config_abs_path: str, - ) -> SequentialAgent: - return super().from_config(config, config_abs_path)