feat(config): support ADK built-in and custom tools in config

PiperOrigin-RevId: 787735915
This commit is contained in:
Liang Wu
2025-07-27 11:14:04 -07:00
committed by Copybara-Service
parent c69dcf8779
commit f1889ae440
10 changed files with 441 additions and 133 deletions
+5 -4
View File
@@ -41,6 +41,7 @@ from ..events.event import Event
from ..utils.feature_decorator import working_in_progress
from .base_agent_config import BaseAgentConfig
from .callback_context import CallbackContext
from .common_configs import AgentRefConfig
if TYPE_CHECKING:
from .invocation_context import InvocationContext
@@ -503,11 +504,13 @@ class BaseAgent(BaseModel):
Args:
config: The config to create the agent from.
config_abs_path: The absolute path to the config file that contains the
agent config.
Returns:
The created agent.
"""
from .config_agent_utils import build_sub_agent
from .config_agent_utils import resolve_agent_reference
from .config_agent_utils import resolve_callbacks
kwargs: Dict[str, Any] = {
@@ -517,9 +520,7 @@ class BaseAgent(BaseModel):
if config.sub_agents:
sub_agents = []
for sub_agent_config in config.sub_agents:
sub_agent = build_sub_agent(
sub_agent_config, config_abs_path.rsplit('/', 1)[0]
)
sub_agent = resolve_agent_reference(sub_agent_config, config_abs_path)
sub_agents.append(sub_agent)
kwargs['sub_agents'] = sub_agents
+2 -61
View File
@@ -43,6 +43,7 @@ from typing_extensions import TypeAlias
from ..events.event import Event
from ..utils.feature_decorator import working_in_progress
from .callback_context import CallbackContext
from .common_configs import AgentRefConfig
from .common_configs import CodeConfig
if TYPE_CHECKING:
@@ -52,66 +53,6 @@ if TYPE_CHECKING:
TBaseAgentConfig = TypeVar('TBaseAgentConfig', bound='BaseAgentConfig')
class SubAgentConfig(BaseModel):
"""The config for a sub-agent."""
model_config = ConfigDict(extra='forbid')
config: Optional[str] = None
"""The YAML config file path of the sub-agent.
Only one of `config` or `code` can be set.
Example:
```
sub_agents:
- config: search_agent.yaml
- config: my_library/my_custom_agent.yaml
```
"""
code: Optional[str] = None
"""The agent instance defined in the code.
Only one of `config` or `code` can be set.
Example:
For the following agent defined in Python code:
```
# my_library/custom_agents.py
from google.adk.agents.llm_agent import LlmAgent
my_custom_agent = LlmAgent(
name="my_custom_agent",
instruction="You are a helpful custom agent.",
model="gemini-2.0-flash",
)
```
The yaml config should be:
```
sub_agents:
- code: my_library.custom_agents.my_custom_agent
```
"""
@model_validator(mode='after')
def validate_exactly_one_field(self):
code_provided = self.code is not None
config_provided = self.config is not None
if code_provided and config_provided:
raise ValueError('Only one of code or config should be provided')
if not code_provided and not config_provided:
raise ValueError('Exactly one of code or config must be provided')
return self
@working_in_progress('BaseAgentConfig is not ready for use.')
class BaseAgentConfig(BaseModel):
"""The config for the YAML schema of a BaseAgent.
@@ -133,7 +74,7 @@ class BaseAgentConfig(BaseModel):
description: str = ''
"""Optional. The description of the agent."""
sub_agents: Optional[List[SubAgentConfig]] = None
sub_agents: Optional[List[AgentRefConfig]] = None
"""Optional. The sub-agents of the agent."""
before_agent_callbacks: Optional[List[CodeConfig]] = None
+63
View File
@@ -21,6 +21,7 @@ from typing import Optional
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import model_validator
from ..utils.feature_decorator import working_in_progress
@@ -77,3 +78,65 @@ class CodeConfig(BaseModel):
value: True
```
"""
class AgentRefConfig(BaseModel):
"""The config for the reference to another agent."""
model_config = ConfigDict(extra="forbid")
config_path: Optional[str] = None
"""The YAML config file path of the sub-agent.
Only one of `config_path` or `code` can be set.
Example:
```
sub_agents:
- config_path: search_agent.yaml
- config_path: my_library/my_custom_agent.yaml
```
"""
code: Optional[str] = None
"""The agent instance defined in the code.
Only one of `config` or `code` can be set.
Example:
For the following agent defined in Python code:
```
# my_library/custom_agents.py
from google.adk.agents.llm_agent import LlmAgent
my_custom_agent = LlmAgent(
name="my_custom_agent",
instruction="You are a helpful custom agent.",
model="gemini-2.0-flash",
)
```
The yaml config should be:
```
sub_agents:
- code: my_library.custom_agents.my_custom_agent
```
"""
@model_validator(mode="after")
def validate_exactly_one_field(self) -> AgentRefConfig:
code_provided = self.code is not None
config_path_provided = self.config_path is not None
if code_provided and config_path_provided:
raise ValueError("Only one of `code` or `config_path` should be provided")
if not code_provided and not config_path_provided:
raise ValueError(
"Exactly one of `code` or `config_path` must be provided"
)
return self
+29 -22
View File
@@ -24,7 +24,7 @@ import yaml
from ..utils.feature_decorator import working_in_progress
from .agent_config import AgentConfig
from .base_agent import BaseAgent
from .base_agent_config import SubAgentConfig
from .common_configs import AgentRefConfig
from .common_configs import CodeConfig
from .llm_agent import LlmAgent
from .llm_agent_config import LlmAgentConfig
@@ -90,44 +90,48 @@ def _load_config_from_path(config_path: str) -> AgentConfig:
return AgentConfig.model_validate(config_data)
@working_in_progress("build_sub_agent is not ready for use.")
def build_sub_agent(
sub_config: SubAgentConfig, parent_agent_folder_path: str
@working_in_progress("resolve_agent_reference is not ready for use.")
def resolve_agent_reference(
ref_config: AgentRefConfig, referencing_agent_config_abs_path: str
) -> BaseAgent:
"""Build a sub-agent from configuration.
"""Build an agent from a reference.
Args:
sub_config: The sub-agent configuration (SubAgentConfig).
parent_agent_folder_path: The folder path to the parent agent's YAML config.
ref_config: The agent reference configuration (AgentRefConfig).
referencing_agent_config_abs_path: The absolute path to the agent config
that contains the reference.
Returns:
The created sub-agent instance.
The created agent instance.
"""
if sub_config.config:
if os.path.isabs(sub_config.config):
return from_config(sub_config.config)
if ref_config.config_path:
if os.path.isabs(ref_config.config_path):
return from_config(ref_config.config_path)
else:
return from_config(
os.path.join(parent_agent_folder_path, sub_config.config)
os.path.join(
referencing_agent_config_abs_path.rsplit("/", 1)[0],
ref_config.config_path,
)
)
elif sub_config.code:
return _resolve_sub_agent_code_reference(sub_config.code)
elif ref_config.code:
return _resolve_agent_code_reference(ref_config.code)
else:
raise ValueError("SubAgentConfig must have either 'code' or 'config'")
raise ValueError("AgentRefConfig must have either 'code' or 'config_path'")
@working_in_progress("_resolve_sub_agent_code_reference is not ready for use.")
def _resolve_sub_agent_code_reference(code: str) -> Any:
"""Resolve a code reference to an actual agent object.
@working_in_progress("_resolve_agent_code_reference is not ready for use.")
def _resolve_agent_code_reference(code: str) -> Any:
"""Resolve a code reference to an actual agent instance.
Args:
code: The code reference to the sub-agent.
code: The fully-qualified path to an agent instance.
Returns:
The resolved agent object.
The resolved agent instance.
Raises:
ValueError: If the code reference cannot be resolved.
ValueError: If the agent reference cannot be resolved.
"""
if "." not in code:
raise ValueError(f"Invalid code reference: {code}")
@@ -137,7 +141,10 @@ def _resolve_sub_agent_code_reference(code: str) -> Any:
obj = getattr(module, obj_name)
if callable(obj):
raise ValueError(f"Invalid code reference to a callable: {code}")
raise ValueError(f"Invalid agent reference to a callable: {code}")
if not isinstance(obj, BaseAgent):
raise ValueError(f"Invalid agent reference to a non-agent instance: {code}")
return obj
@@ -1,5 +1,37 @@
{
"$defs": {
"AgentRefConfig": {
"additionalProperties": false,
"description": "The config for the reference to another agent.",
"properties": {
"config_path": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"default": null,
"title": "Config Path"
},
"code": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"default": null,
"title": "Code"
}
},
"title": "AgentRefConfig",
"type": "object"
},
"ArgumentConfig": {
"additionalProperties": false,
"description": "An argument passed to a function or a class's constructor.",
@@ -26,6 +58,84 @@
"title": "ArgumentConfig",
"type": "object"
},
"BaseAgentConfig": {
"additionalProperties": true,
"description": "The config for the YAML schema of a BaseAgent.\n\nDo not use this class directly. It's the base class for all agent configs.",
"properties": {
"agent_class": {
"anyOf": [
{
"const": "BaseAgent",
"type": "string"
},
{
"type": "string"
}
],
"default": "BaseAgent",
"title": "Agent Class"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"default": "",
"title": "Description",
"type": "string"
},
"sub_agents": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/AgentRefConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Sub Agents"
},
"before_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Before Agent Callbacks"
},
"after_agent_callbacks": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "After Agent Callbacks"
}
},
"required": [
"name"
],
"title": "BaseAgentConfig",
"type": "object"
},
"CodeConfig": {
"additionalProperties": false,
"description": "Code reference config for a variable, a function, or a class.\n\nThis config is used for configuring callbacks and tools.",
@@ -82,7 +192,7 @@
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
"$ref": "#/$defs/AgentRefConfig"
},
"type": "array"
},
@@ -210,7 +320,7 @@
"anyOf": [
{
"items": {
"$ref": "#/$defs/CodeConfig"
"$ref": "#/$defs/ToolConfig"
},
"type": "array"
},
@@ -312,7 +422,7 @@
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
"$ref": "#/$defs/AgentRefConfig"
},
"type": "array"
},
@@ -395,7 +505,7 @@
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
"$ref": "#/$defs/AgentRefConfig"
},
"type": "array"
},
@@ -466,7 +576,7 @@
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
"$ref": "#/$defs/AgentRefConfig"
},
"type": "array"
},
@@ -514,36 +624,37 @@
"title": "SequentialAgentConfig",
"type": "object"
},
"SubAgentConfig": {
"ToolArgsConfig": {
"additionalProperties": true,
"description": "The configuration for tool arguments.\n\nThis config allows arbitrary key-value pairs as tool arguments.",
"properties": {},
"title": "ToolArgsConfig",
"type": "object"
},
"ToolConfig": {
"additionalProperties": false,
"description": "The config for a sub-agent.",
"description": "The configuration for a tool.\n\nThe config supports these types of tools:\n1. ADK built-in tools\n2. User-defined tool instances\n3. User-defined tool classes\n4. User-defined functions that generate tool instances\n5. User-defined function tools\n\nFor examples:\n\n 1. For ADK built-in tool instances or classes in `google.adk.tools` package,\n they can be referenced directly with the `name` and optionally with\n `config`.\n\n ```\n tools:\n - name: google_search\n - name: AgentTool\n config:\n agent: ./another_agent.yaml\n skip_summarization: true\n ```\n\n 2. For user-defined tool instances, the `name` is the fully qualified path\n to the tool instance.\n\n ```\n tools:\n - name: my_package.my_module.my_tool\n ```\n\n 3. For user-defined tool classes (custom tools), the `name` is the fully\n qualified path to the tool class and `config` is the arguments for the tool.\n\n ```\n tools:\n - name: my_package.my_module.my_tool_class\n config:\n my_tool_arg1: value1\n my_tool_arg2: value2\n ```\n\n 4. For user-defined functions that generate tool instances, the `name` is the\n fully qualified path to the function and `config` is passed to the function\n as arguments.\n\n ```\n tools:\n - name: my_package.my_module.my_tool_function\n config:\n my_function_arg1: value1\n my_function_arg2: value2\n ```\n\n The function must have the following signature:\n ```\n def my_function(config: ToolArgsConfig) -> BaseTool:\n ...\n ```\n\n 5. For user-defined function tools, the `name` is the fully qualified path\n to the function.\n\n ```\n tools:\n - name: my_package.my_module.my_function_tool\n ```",
"properties": {
"config": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"default": null,
"title": "Config"
"name": {
"title": "Name",
"type": "string"
},
"code": {
"args": {
"anyOf": [
{
"type": "string"
"$ref": "#/$defs/ToolArgsConfig"
},
{
"type": "null"
}
],
"default": null,
"title": "Code"
"default": null
}
},
"title": "SubAgentConfig",
"required": [
"name"
],
"title": "ToolConfig",
"type": "object"
}
},
@@ -559,6 +670,9 @@
},
{
"$ref": "#/$defs/SequentialAgentConfig"
},
{
"$ref": "#/$defs/BaseAgentConfig"
}
],
"description": "The config for the YAML schema to create an agent.",
+43 -12
View File
@@ -17,6 +17,7 @@ from __future__ import annotations
import importlib
import inspect
import logging
import os
from typing import Any
from typing import AsyncGenerator
from typing import Awaitable
@@ -46,7 +47,9 @@ 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
from ..tools.function_tool import FunctionTool
from ..tools.tool_context import ToolContext
@@ -525,31 +528,59 @@ class LlmAgent(BaseAgent):
@classmethod
@working_in_progress('LlmAgent._resolve_tools is not ready for use.')
def _resolve_tools(cls, tools_config: list[CodeConfig]) -> list[Any]:
def _resolve_tools(
cls, tool_configs: list[ToolConfig], config_abs_path: str
) -> list[Any]:
"""Resolve tools from configuration.
Args:
tools_config: List of tool configurations (CodeConfig objects).
tool_configs: List of tool configurations (ToolConfig objects).
config_abs_path: The absolute path to the agent config file.
Returns:
List of resolved tool objects.
"""
resolved_tools = []
for tool_config in tools_config:
for tool_config in tool_configs:
if '.' not in tool_config.name:
# ADK built-in tools
module = importlib.import_module('google.adk.tools')
obj = getattr(module, tool_config.name)
if isinstance(obj, ToolUnion):
resolved_tools.append(obj)
else:
raise ValueError(
f'Invalid tool name: {tool_config.name} is not a built-in tool.'
)
else:
from .config_agent_utils import resolve_code_reference
# User-defined tools
module_path, obj_name = tool_config.name.rsplit('.', 1)
module = importlib.import_module(module_path)
obj = getattr(module, obj_name)
resolved_tools.append(resolve_code_reference(tool_config))
if isinstance(obj, BaseTool) or isinstance(obj, BaseToolset):
logger.debug(
'Tool %s is an instance of BaseTool/BaseToolset.', tool_config.name
)
resolved_tools.append(obj)
elif inspect.isclass(obj) and (
issubclass(obj, BaseTool) or issubclass(obj, BaseToolset)
):
logger.debug(
'Tool %s is a sub-class of BaseTool/BaseToolset.', tool_config.name
)
resolved_tools.append(
obj.from_config(tool_config.args, config_abs_path)
)
elif callable(obj):
if tool_config.args:
logger.debug(
'Tool %s is a user-defined tool-generating function.',
tool_config.name,
)
resolved_tools.append(obj(tool_config.args))
else:
logger.debug(
'Tool %s is a user-defined function tool.', tool_config.name
)
resolved_tools.append(obj)
else:
raise ValueError(f'Invalid tool YAML config: {tool_config}.')
return resolved_tools
@@ -582,7 +613,7 @@ class LlmAgent(BaseAgent):
if config.output_key:
agent.output_key = config.output_key
if config.tools:
agent.tools = cls._resolve_tools(config.tools)
agent.tools = cls._resolve_tools(config.tools, config_abs_path)
if config.before_model_callbacks:
agent.before_model_callback = resolve_callbacks(
config.before_model_callbacks
+2 -1
View File
@@ -21,6 +21,7 @@ from typing import Optional
from pydantic import ConfigDict
from ..tools.base_tool import ToolConfig
from .base_agent_config import BaseAgentConfig
from .common_configs import CodeConfig
@@ -63,7 +64,7 @@ class LlmAgentConfig(BaseAgentConfig):
include_contents: Literal['default', 'none'] = 'default'
"""Optional. LlmAgent.include_contents."""
tools: Optional[list[CodeConfig]] = None
tools: Optional[list[ToolConfig]] = None
"""Optional. LlmAgent.tools.
Examples:
+29
View File
@@ -18,13 +18,16 @@ from typing import Any
from typing import TYPE_CHECKING
from google.genai import types
from pydantic import BaseModel
from pydantic import model_validator
from typing_extensions import override
from . import _automatic_function_calling_util
from ..agents.common_configs import AgentRefConfig
from ..memory.in_memory_memory_service import InMemoryMemoryService
from ._forwarding_artifact_service import ForwardingArtifactService
from .base_tool import BaseTool
from .base_tool import ToolArgsConfig
from .tool_context import ToolContext
if TYPE_CHECKING:
@@ -154,3 +157,29 @@ class AgentTool(BaseTool):
else:
tool_result = merged_text
return tool_result
@classmethod
@override
def from_config(
cls, config: ToolArgsConfig, config_abs_path: str
) -> AgentTool:
from ..agents import config_agent_utils
agent_tool_config = AgentToolConfig.model_validate(config.model_dump())
agent = config_agent_utils.resolve_agent_reference(
agent_tool_config.agent, config_abs_path
)
return cls(
agent=agent, skip_summarization=agent_tool_config.skip_summarization
)
class AgentToolConfig(BaseModel):
"""The config for the AgentTool."""
agent: AgentRefConfig
"""The reference to the agent instance."""
skip_summarization: bool = False
"""Whether to skip summarization of the agent output."""
+122 -1
View File
@@ -17,9 +17,13 @@ from __future__ import annotations
from abc import ABC
from typing import Any
from typing import Optional
from typing import Type
from typing import TYPE_CHECKING
from typing import TypeVar
from google.genai import types
from pydantic import BaseModel
from pydantic import ConfigDict
from ..utils.variant_utils import get_google_llm_variant
from ..utils.variant_utils import GoogleLLMVariant
@@ -28,6 +32,8 @@ from .tool_context import ToolContext
if TYPE_CHECKING:
from ..models.llm_request import LlmRequest
SelfTool = TypeVar("SelfTool", bound="BaseTool")
class BaseTool(ABC):
"""The base class for all tools."""
@@ -78,7 +84,7 @@ class BaseTool(ABC):
Returns:
The result of running the tool.
"""
raise NotImplementedError(f'{type(self)} is not implemented')
raise NotImplementedError(f"{type(self)} is not implemented")
async def process_llm_request(
self, *, tool_context: ToolContext, llm_request: LlmRequest
@@ -122,6 +128,25 @@ class BaseTool(ABC):
def _api_variant(self) -> GoogleLLMVariant:
return get_google_llm_variant()
@classmethod
def from_config(
cls: Type[SelfTool], config: ToolArgsConfig, config_abs_path: str
) -> SelfTool:
"""Creates a tool instance from a config.
Subclasses should override and implement this method to do custom
initialization from a config.
Args:
config: The config for the tool.
config_abs_path: The absolute path to the config file that contains the
tool config.
Returns:
The tool instance.
"""
raise NotImplementedError(f"from_config for {cls} not implemented.")
def _find_tool_with_function_declarations(
llm_request: LlmRequest,
@@ -138,3 +163,99 @@ def _find_tool_with_function_declarations(
),
None,
)
class ToolArgsConfig(BaseModel):
"""The configuration for tool arguments.
This config allows arbitrary key-value pairs as tool arguments.
"""
model_config = ConfigDict(extra="allow")
class ToolConfig(BaseModel):
"""The configuration for a tool.
The config supports these types of tools:
1. ADK built-in tools
2. User-defined tool instances
3. User-defined tool classes
4. User-defined functions that generate tool instances
5. User-defined function tools
For examples:
1. For ADK built-in tool instances or classes in `google.adk.tools` package,
they can be referenced directly with the `name` and optionally with
`config`.
```
tools:
- name: google_search
- name: AgentTool
config:
agent: ./another_agent.yaml
skip_summarization: true
```
2. For user-defined tool instances, the `name` is the fully qualified path
to the tool instance.
```
tools:
- name: my_package.my_module.my_tool
```
3. For user-defined tool classes (custom tools), the `name` is the fully
qualified path to the tool class and `config` is the arguments for the tool.
```
tools:
- name: my_package.my_module.my_tool_class
config:
my_tool_arg1: value1
my_tool_arg2: value2
```
4. For user-defined functions that generate tool instances, the `name` is the
fully qualified path to the function and `config` is passed to the function
as arguments.
```
tools:
- name: my_package.my_module.my_tool_function
config:
my_function_arg1: value1
my_function_arg2: value2
```
The function must have the following signature:
```
def my_function(config: ToolArgsConfig) -> BaseTool:
...
```
5. For user-defined function tools, the `name` is the fully qualified path
to the function.
```
tools:
- name: my_package.my_module.my_function_tool
```
"""
model_config = ConfigDict(extra="forbid")
name: str
"""The name of the tool.
For ADK built-in tools, the name is the name of the tool, e.g. `google_search`
or `AgentTool`.
For user-defined tools, the name is the fully qualified path to the tool, e.g.
`my_package.my_module.my_tool`.
"""
args: Optional[ToolArgsConfig] = None
"""The args for the tool."""
+9 -9
View File
@@ -50,9 +50,9 @@ agent_class: LoopAgent
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config: sub_agents/code_writer_agent.yaml
- config: sub_agents/code_reviewer_agent.yaml
- config: sub_agents/code_refactorer_agent.yaml
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
"""
config_data = yaml.safe_load(yaml_content)
@@ -68,9 +68,9 @@ agent_class: ParallelAgent
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config: sub_agents/code_writer_agent.yaml
- config: sub_agents/code_reviewer_agent.yaml
- config: sub_agents/code_refactorer_agent.yaml
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
"""
config_data = yaml.safe_load(yaml_content)
@@ -86,9 +86,9 @@ agent_class: SequentialAgent
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents:
- config: sub_agents/code_writer_agent.yaml
- config: sub_agents/code_reviewer_agent.yaml
- config: sub_agents/code_refactorer_agent.yaml
- config_path: sub_agents/code_writer_agent.yaml
- config_path: sub_agents/code_reviewer_agent.yaml
- config_path: sub_agents/code_refactorer_agent.yaml
"""
config_data = yaml.safe_load(yaml_content)