You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
e3c2bf3062
commit
53803522b6
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user