feat(config): support sub_agents in BaseAgentConfig

Currently only support path to YAML or code reference to agent instance.

PiperOrigin-RevId: 782157110
This commit is contained in:
Liang Wu
2025-07-11 16:57:02 -07:00
committed by Copybara-Service
parent 134ec0d71e
commit b2ef9a069e
7 changed files with 236 additions and 12 deletions
+76
View File
@@ -21,6 +21,7 @@ from typing import Awaitable
from typing import Callable
from typing import Dict
from typing import final
from typing import List
from typing import Literal
from typing import Mapping
from typing import Optional
@@ -35,6 +36,7 @@ from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from pydantic import field_validator
from pydantic import model_validator
from typing_extensions import override
from typing_extensions import TypeAlias
@@ -491,6 +493,7 @@ class BaseAgent(BaseModel):
def from_config(
cls: Type[SelfAgent],
config: BaseAgentConfig,
config_abs_path: str,
) -> SelfAgent:
"""Creates an agent from a config.
@@ -506,13 +509,83 @@ class BaseAgent(BaseModel):
Returns:
The created agent.
"""
from .config_agent_utils import build_sub_agent
kwargs: Dict[str, Any] = {
'name': config.name,
'description': config.description,
}
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_agents.append(sub_agent)
kwargs['sub_agents'] = sub_agents
return cls(**kwargs)
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 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.
@@ -531,3 +604,6 @@ class BaseAgentConfig(BaseModel):
description: str = ''
"""Optional. The description of the agent."""
sub_agents: Optional[List[SubAgentConfig]] = None
"""Optional. The sub-agents of the agent."""
+60 -8
View File
@@ -14,14 +14,16 @@
from __future__ import annotations
import importlib
import os
from pathlib import Path
from typing import Any
import yaml
from ..utils.feature_decorator import working_in_progress
from .agent_config import AgentConfig
from .base_agent import BaseAgent
from .base_agent import SubAgentConfig
from .llm_agent import LlmAgent
from .llm_agent import LlmAgentConfig
from .loop_agent import LoopAgent
@@ -51,13 +53,13 @@ def from_config(config_path: str) -> BaseAgent:
config = _load_config_from_path(abs_path)
if isinstance(config.root, LlmAgentConfig):
return LlmAgent.from_config(config.root)
return LlmAgent.from_config(config.root, abs_path)
elif isinstance(config.root, LoopAgentConfig):
return LoopAgent.from_config(config.root)
return LoopAgent.from_config(config.root, abs_path)
elif isinstance(config.root, ParallelAgentConfig):
return ParallelAgent.from_config(config.root)
return ParallelAgent.from_config(config.root, abs_path)
elif isinstance(config.root, SequentialAgentConfig):
return SequentialAgent.from_config(config.root)
return SequentialAgent.from_config(config.root, abs_path)
else:
raise ValueError("Unsupported config type")
@@ -77,12 +79,62 @@ def _load_config_from_path(config_path: str) -> AgentConfig:
FileNotFoundError: If config file doesn't exist.
ValidationError: If config file's content is invalid YAML.
"""
config_path = Path(config_path)
if not config_path.exists():
if not os.path.exists(config_path):
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(config_path, "r", encoding="utf-8") as f:
config_data = yaml.safe_load(f)
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
) -> BaseAgent:
"""Build a sub-agent from configuration.
Args:
sub_config: The sub-agent configuration (SubAgentConfig).
parent_agent_folder_path: The folder path to the parent agent's YAML config.
Returns:
The created sub-agent instance.
"""
if sub_config.config:
if os.path.isabs(sub_config.config):
return from_config(sub_config.config)
else:
return from_config(
os.path.join(parent_agent_folder_path, sub_config.config)
)
elif sub_config.code:
return _resolve_sub_agent_code_reference(sub_config.code)
else:
raise ValueError("SubAgentConfig must have either 'code' or 'config'")
@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.
Args:
code: The code reference to the sub-agent.
Returns:
The resolved agent object.
Raises:
ValueError: If the code reference cannot be resolved.
"""
if "." not in code:
raise ValueError(f"Invalid code reference: {code}")
module_path, obj_name = code.rsplit(".", 1)
module = importlib.import_module(module_path)
obj = getattr(module, obj_name)
if callable(obj):
raise ValueError(f"Invalid code reference to a callable: {code}")
return obj
@@ -22,6 +22,21 @@
"title": "Description",
"type": "string"
},
"sub_agents": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Sub Agents"
},
"model": {
"anyOf": [
{
@@ -89,6 +104,21 @@
"title": "Description",
"type": "string"
},
"sub_agents": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Sub Agents"
},
"max_iterations": {
"anyOf": [
{
@@ -126,6 +156,21 @@
"default": "",
"title": "Description",
"type": "string"
},
"sub_agents": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Sub Agents"
}
},
"required": [
@@ -152,6 +197,21 @@
"default": "",
"title": "Description",
"type": "string"
},
"sub_agents": {
"anyOf": [
{
"items": {
"$ref": "#/$defs/SubAgentConfig"
},
"type": "array"
},
{
"type": "null"
}
],
"default": null,
"title": "Sub Agents"
}
},
"required": [
@@ -159,6 +219,38 @@
],
"title": "SequentialAgentConfig",
"type": "object"
},
"SubAgentConfig": {
"additionalProperties": false,
"description": "The config for a sub-agent.",
"properties": {
"config": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"default": null,
"title": "Config"
},
"code": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"default": null,
"title": "Code"
}
},
"title": "SubAgentConfig",
"type": "object"
}
},
"anyOf": [
+2 -1
View File
@@ -526,8 +526,9 @@ class LlmAgent(BaseAgent):
def from_config(
cls: Type[LlmAgent],
config: LlmAgentConfig,
config_abs_path: str,
) -> LlmAgent:
agent = super().from_config(config)
agent = super().from_config(config, config_abs_path)
if config.model:
agent.model = config.model
if config.instruction:
+2 -1
View File
@@ -73,8 +73,9 @@ class LoopAgent(BaseAgent):
def from_config(
cls: Type[LoopAgent],
config: LoopAgentConfig,
config_abs_path: str,
) -> LoopAgent:
agent = super().from_config(config)
agent = super().from_config(config, config_abs_path)
if config.max_iterations:
agent.max_iterations = config.max_iterations
return agent
+2 -1
View File
@@ -122,8 +122,9 @@ class ParallelAgent(BaseAgent):
def from_config(
cls: Type[ParallelAgent],
config: ParallelAgentConfig,
config_abs_path: str,
) -> ParallelAgent:
return super().from_config(config)
return super().from_config(config, config_abs_path)
@working_in_progress('ParallelAgentConfig is not ready for use.')
+2 -1
View File
@@ -85,8 +85,9 @@ class SequentialAgent(BaseAgent):
def from_config(
cls: Type[SequentialAgent],
config: SequentialAgentConfig,
config_abs_path: str,
) -> SequentialAgent:
return super().from_config(config)
return super().from_config(config, config_abs_path)
@working_in_progress('SequentialAgentConfig is not ready for use.')