From 8222fd8a14d3623e548d491ce5f64a198642d109 Mon Sep 17 00:00:00 2001 From: jayy-77 <1427jay@gmail.com> Date: Fri, 6 Feb 2026 20:09:29 +0530 Subject: [PATCH] added conditionalcgent for runtime workflow branching. --- src/google/adk/agents/__init__.py | 2 + src/google/adk/agents/agent_config.py | 3 + src/google/adk/agents/if_agent.py | 262 +++++++++++++++++++++++ src/google/adk/agents/if_agent_config.py | 80 +++++++ 4 files changed, 347 insertions(+) create mode 100644 src/google/adk/agents/if_agent.py create mode 100644 src/google/adk/agents/if_agent_config.py diff --git a/src/google/adk/agents/__init__.py b/src/google/adk/agents/__init__.py index 35198179a5..8d37ab9777 100644 --- a/src/google/adk/agents/__init__.py +++ b/src/google/adk/agents/__init__.py @@ -13,6 +13,7 @@ # limitations under the License. from .base_agent import BaseAgent +from .if_agent import IfAgent from .invocation_context import InvocationContext from .live_request_queue import LiveRequest from .live_request_queue import LiveRequestQueue @@ -27,6 +28,7 @@ __all__ = [ 'Agent', 'BaseAgent', + 'IfAgent', 'LlmAgent', 'LoopAgent', 'McpInstructionProvider', diff --git a/src/google/adk/agents/agent_config.py b/src/google/adk/agents/agent_config.py index add31f4bb3..b69b865430 100644 --- a/src/google/adk/agents/agent_config.py +++ b/src/google/adk/agents/agent_config.py @@ -25,12 +25,14 @@ from ..utils.feature_decorator import experimental from .base_agent_config import BaseAgentConfig +from .if_agent_config import IfAgentConfig from .llm_agent_config import LlmAgentConfig from .loop_agent_config import LoopAgentConfig from .parallel_agent_config import ParallelAgentConfig from .sequential_agent_config import SequentialAgentConfig _ADK_AGENT_CLASSES: set[str] = { + "IfAgent", "LlmAgent", "LoopAgent", "ParallelAgent", @@ -56,6 +58,7 @@ def agent_config_discriminator(v: Any) -> str: # A discriminated union of all possible agent configurations. ConfigsUnion = Annotated[ Union[ + Annotated[IfAgentConfig, Tag("IfAgent")], Annotated[LlmAgentConfig, Tag("LlmAgent")], Annotated[LoopAgentConfig, Tag("LoopAgent")], Annotated[ParallelAgentConfig, Tag("ParallelAgent")], diff --git a/src/google/adk/agents/if_agent.py b/src/google/adk/agents/if_agent.py new file mode 100644 index 0000000000..84899fe2fe --- /dev/null +++ b/src/google/adk/agents/if_agent.py @@ -0,0 +1,262 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Conditional branching agent implementation (IfAgent/ConditionalAgent).""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +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 + +from ..agents.base_agent import BaseAgent +from ..agents.base_agent import BaseAgentState +from ..agents.base_agent_config import BaseAgentConfig +from ..agents.config_agent_utils import resolve_code_reference +from ..agents.if_agent_config import IfAgentConfig +from ..agents.invocation_context import InvocationContext +from ..events.event import Event +from ..utils.context_utils import Aclosing +from ..utils.feature_decorator import experimental + +logger = logging.getLogger('google_adk.' + __name__) + + +@experimental +class IfAgentState(BaseAgentState): + """State for IfAgent.""" + + condition_result: bool = False + """The result of the condition evaluation.""" + + chosen_agent: str = '' + """The name of the chosen sub-agent.""" + + +@experimental +class IfAgent(BaseAgent): + """Conditional branching agent for runtime workflow decisions. + + IfAgent evaluates a condition function at runtime and delegates execution + to one of two sub-agents (if_true or if_false) based on the result. + + The condition function receives the InvocationContext and must return a + boolean. This enables dynamic routing based on: + - Session state and history + - User message content + - Custom metadata + - External state or API calls + + Example Python usage: + def is_premium_user(ctx: InvocationContext) -> bool: + return ctx.session.custom_metadata.get('tier') == 'premium' + + if_agent = IfAgent( + name='user_router', + description='Routes to specialized agents based on user tier', + condition=is_premium_user, + sub_agents=[premium_agent, standard_agent] + ) + + Example YAML usage: + name: user_router + agent_class: IfAgent + description: Routes to specialized agents based on user tier + condition: + module: my_module.conditions + function: is_premium_user + sub_agents: + - agent_ref: premium_agent + - agent_ref: standard_agent + """ + + config_type: ClassVar[Type[BaseAgentConfig]] = IfAgentConfig + """The config type for this agent.""" + + def __init__( + self, + *, + condition: Callable[[InvocationContext], bool] | None = None, + **kwargs: Any, + ): + """Initializes an IfAgent. + + Args: + condition: A callable that takes InvocationContext and returns bool. + If True, executes the first sub-agent (if_true). + If False, executes the second sub-agent (if_false). + If None, defaults to always True. + **kwargs: Additional arguments passed to BaseAgent. + + Raises: + ValueError: If sub_agents is not provided or doesn't have exactly 2 + agents. + """ + super().__init__(**kwargs) + self._condition = condition or (lambda ctx: True) + + if not self.sub_agents or len(self.sub_agents) != 2: + raise ValueError( + f'IfAgent {self.name} requires exactly 2 sub-agents ' + f'(if_true and if_false), but got {len(self.sub_agents) if self.sub_agents else 0}.' + ) + + @override + @classmethod + @experimental + def _parse_config( + cls: Type[IfAgent], + config: IfAgentConfig, + config_abs_path: str, + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: + """Parses IfAgent-specific configuration from YAML. + + Args: + config: The IfAgentConfig parsed from YAML. + config_abs_path: Absolute path to the config file. + kwargs: Keyword arguments being built for agent instantiation. + + Returns: + Updated kwargs dictionary with resolved condition. + """ + if config.condition: + kwargs['condition'] = resolve_code_reference(config.condition) + return kwargs + + @property + def if_true_agent(self) -> BaseAgent: + """Returns the agent to execute when condition is True.""" + return self.sub_agents[0] + + @property + def if_false_agent(self) -> BaseAgent: + """Returns the agent to execute when condition is False.""" + return self.sub_agents[1] + + def _evaluate_condition(self, ctx: InvocationContext) -> bool: + """Evaluates the condition function safely. + + Args: + ctx: The invocation context. + + Returns: + The result of the condition evaluation, or False if an error occurs. + """ + try: + result = self._condition(ctx) + logger.debug( + 'IfAgent %s: condition evaluated to %s', self.name, result + ) + return bool(result) + except Exception as e: # pylint: disable=broad-except + logger.warning( + 'IfAgent %s: condition evaluation failed with error: %s. ' + 'Defaulting to False.', + self.name, + e, + ) + return False + + @override + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + """Runs the IfAgent by evaluating condition and delegating to chosen agent. + + Args: + ctx: The invocation context. + + Yields: + Events from the chosen sub-agent. + """ + # Load or initialize state + agent_state = self._load_agent_state(ctx, IfAgentState) + + if agent_state: + # Resuming from a saved state + condition_result = agent_state.condition_result + chosen_agent_name = agent_state.chosen_agent + logger.debug( + 'IfAgent %s: resuming with condition_result=%s, chosen_agent=%s', + self.name, + condition_result, + chosen_agent_name, + ) + else: + # First invocation: evaluate condition + condition_result = self._evaluate_condition(ctx) + chosen_agent = ( + self.if_true_agent if condition_result else self.if_false_agent + ) + chosen_agent_name = chosen_agent.name + + # Save state before executing sub-agent + if ctx.is_resumable: + agent_state = IfAgentState( + condition_result=condition_result, + chosen_agent=chosen_agent_name, + ) + ctx.set_agent_state(self.name, agent_state=agent_state) + yield self._create_agent_state_event(ctx) + + # Execute the chosen sub-agent + chosen_agent = ( + self.if_true_agent if condition_result else self.if_false_agent + ) + + async with Aclosing(chosen_agent.run_async(ctx)) as agen: + async for event in agen: + yield event + + # Mark completion + if ctx.is_resumable: + ctx.set_agent_state(self.name, end_of_agent=True) + yield self._create_agent_state_event(ctx) + + @override + async def _run_live_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + """Runs the IfAgent in live mode with bidirectional streaming. + + Args: + ctx: The invocation context. + + Yields: + Events from the chosen sub-agent. + """ + # Evaluate condition + condition_result = self._evaluate_condition(ctx) + chosen_agent = ( + self.if_true_agent if condition_result else self.if_false_agent + ) + + logger.debug( + 'IfAgent %s (live): condition=%s, executing agent=%s', + self.name, + condition_result, + chosen_agent.name, + ) + + # Execute the chosen sub-agent in live mode + async with Aclosing(chosen_agent.run_live(ctx)) as agen: + async for event in agen: + yield event diff --git a/src/google/adk/agents/if_agent_config.py b/src/google/adk/agents/if_agent_config.py new file mode 100644 index 0000000000..0a503ca350 --- /dev/null +++ b/src/google/adk/agents/if_agent_config.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Config definition for IfAgent.""" + +from __future__ import annotations + +from typing import Optional + +from pydantic import ConfigDict +from pydantic import Field + +from ..agents.base_agent_config import BaseAgentConfig +from ..agents.common_configs import CodeConfig +from ..utils.feature_decorator import experimental + + +@experimental +class IfAgentConfig(BaseAgentConfig): + """The config for the YAML schema of an IfAgent. + + IfAgent enables conditional branching in agent workflows based on runtime + conditions. It evaluates a condition function against the invocation context + and delegates to one of two sub-agents based on the result. + + Example YAML: + name: conditional_router + agent_class: IfAgent + description: Routes to specialized agents based on user intent + condition: + module: my_module.conditions + function: is_technical_query + sub_agents: + - agent_ref: technical_agent + - agent_ref: general_agent + + Example with inline condition arguments: + name: priority_router + agent_class: IfAgent + description: Routes based on priority threshold + condition: + module: my_module.conditions + function: check_priority_threshold + args: + threshold: 5 + sub_agents: + - agent_ref: high_priority_agent + - agent_ref: normal_priority_agent + """ + + model_config = ConfigDict( + extra='forbid', + ) + + agent_class: str = Field( + default='IfAgent', + description='The value is used to uniquely identify the IfAgent class.', + ) + + condition: Optional[CodeConfig] = Field( + default=None, + description=( + 'Optional. IfAgent.condition. A CodeConfig that references a callable' + ' (function, lambda, or class with __call__) that takes an' + ' InvocationContext and returns a boolean. If True, executes' + ' if_true agent; if False, executes if_false agent. If not provided,' + ' defaults to always True.' + ), + )