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
This commit is contained in:
Wei Sun (Jack)
2025-08-05 23:52:50 -07:00
committed by Copybara-Service
parent e3c2bf3062
commit 53803522b6
5 changed files with 72 additions and 58 deletions
+38 -7
View File
@@ -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
+23 -23
View File
@@ -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
+9 -8
View File
@@ -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
+2 -10
View File
@@ -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)
-10
View File
@@ -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)