feat(config): add configs for ParallelAgent and SequentialAgent

PiperOrigin-RevId: 781704814
This commit is contained in:
Liang Wu
2025-07-10 15:24:22 -07:00
committed by Copybara-Service
parent d83362725d
commit a313c2c1af
4 changed files with 106 additions and 4 deletions
+4
View File
@@ -21,11 +21,15 @@ from pydantic import RootModel
from ..utils.feature_decorator import working_in_progress
from .llm_agent import LlmAgentConfig
from .loop_agent import LoopAgentConfig
from .parallel_agent import ParallelAgentConfig
from .sequential_agent import SequentialAgentConfig
# A discriminated union of all possible agent configurations.
ConfigsUnion = Union[
LlmAgentConfig,
LoopAgentConfig,
ParallelAgentConfig,
SequentialAgentConfig,
]
@@ -107,6 +107,58 @@
],
"title": "LoopAgentConfig",
"type": "object"
},
"ParallelAgentConfig": {
"additionalProperties": false,
"description": "The config for the YAML schema of a ParallelAgent.",
"properties": {
"agent_class": {
"const": "ParallelAgent",
"default": "ParallelAgent",
"title": "Agent Class",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"default": "",
"title": "Description",
"type": "string"
}
},
"required": [
"name"
],
"title": "ParallelAgentConfig",
"type": "object"
},
"SequentialAgentConfig": {
"additionalProperties": false,
"description": "The config for the YAML schema of a SequentialAgent.",
"properties": {
"agent_class": {
"const": "SequentialAgent",
"default": "SequentialAgent",
"title": "Agent Class",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"default": "",
"title": "Description",
"type": "string"
}
},
"required": [
"name"
],
"title": "SequentialAgentConfig",
"type": "object"
}
},
"anyOf": [
@@ -115,6 +167,12 @@
},
{
"$ref": "#/$defs/LoopAgentConfig"
},
{
"$ref": "#/$defs/ParallelAgentConfig"
},
{
"$ref": "#/$defs/SequentialAgentConfig"
}
],
"description": "The config for the YAML schema to create an agent.",
+23 -3
View File
@@ -18,9 +18,13 @@ from __future__ import annotations
import asyncio
from typing import AsyncGenerator
from typing import Literal
from typing import Type
from typing_extensions import override
from ..agents.base_agent import BaseAgentConfig
from ..agents.base_agent import working_in_progress
from ..agents.invocation_context import InvocationContext
from ..events.event import Event
from .base_agent import BaseAgent
@@ -33,9 +37,9 @@ def _create_branch_ctx_for_sub_agent(
) -> InvocationContext:
"""Create isolated branch for every sub-agent."""
invocation_context = invocation_context.model_copy()
branch_suffix = f"{agent.name}.{sub_agent.name}"
branch_suffix = f'{agent.name}.{sub_agent.name}'
invocation_context.branch = (
f"{invocation_context.branch}.{branch_suffix}"
f'{invocation_context.branch}.{branch_suffix}'
if invocation_context.branch
else branch_suffix
)
@@ -109,5 +113,21 @@ class ParallelAgent(BaseAgent):
async def _run_live_impl(
self, ctx: InvocationContext
) -> AsyncGenerator[Event, None]:
raise NotImplementedError("This is not supported yet for ParallelAgent.")
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,
) -> ParallelAgent:
return super().from_config(config)
@working_in_progress('ParallelAgentConfig is not ready for use.')
class ParallelAgentConfig(BaseAgentConfig):
"""The config for the YAML schema of a ParallelAgent."""
agent_class: Literal['ParallelAgent'] = 'ParallelAgent'
+21 -1
View File
@@ -17,9 +17,13 @@
from __future__ import annotations
from typing import AsyncGenerator
from typing import Literal
from typing import Type
from typing_extensions import override
from ..agents.base_agent import BaseAgentConfig
from ..agents.base_agent import working_in_progress
from ..agents.invocation_context import InvocationContext
from ..events.event import Event
from .base_agent import BaseAgent
@@ -60,7 +64,7 @@ class SequentialAgent(BaseAgent):
Signals that the model has successfully completed the user's question
or task.
"""
return "Task completion signaled."
return 'Task completion signaled.'
if isinstance(sub_agent, LlmAgent):
# Use function name to dedupe.
@@ -74,3 +78,19 @@ 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,
) -> SequentialAgent:
return super().from_config(config)
@working_in_progress('SequentialAgentConfig is not ready for use.')
class SequentialAgentConfig(BaseAgentConfig):
"""The config for the YAML schema of a SequentialAgent."""
agent_class: Literal['SequentialAgent'] = 'SequentialAgent'