nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/cli/register_workflow.py

@@ -27,6 +27,8 @@ from nat.cli.type_registry import EvaluatorRegisteredCallableT
 from nat.cli.type_registry import FrontEndBuildCallableT
 from nat.cli.type_registry import FrontEndRegisteredCallableT
 from nat.cli.type_registry import FunctionBuildCallableT
+from nat.cli.type_registry import FunctionGroupBuildCallableT
+from nat.cli.type_registry import FunctionGroupRegisteredCallableT
 from nat.cli.type_registry import FunctionRegisteredCallableT
 from nat.cli.type_registry import LLMClientBuildCallableT
 from nat.cli.type_registry import LLMClientRegisteredCallableT
@@ -60,6 +62,7 @@ from nat.data_models.embedder import EmbedderBaseConfigT
 from nat.data_models.evaluator import EvaluatorBaseConfigT
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionConfigT
+from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.memory import MemoryBaseConfigT
 from nat.data_models.object_store import ObjectStoreBaseConfigT
@@ -155,10 +158,7 @@ def register_function(config_type: type[FunctionConfigT],
 
         context_manager_fn = asynccontextmanager(fn)
 
-        if framework_wrappers is None:
-            framework_wrappers_list: list[str] = []
-        else:
-            framework_wrappers_list = list(framework_wrappers)
+        framework_wrappers_list = list(framework_wrappers or [])
 
         discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
                                                                 component_type=ComponentEnum.FUNCTION)
@@ -177,6 +177,40 @@ def register_function(config_type: type[FunctionConfigT],
     return register_function_inner
 
 
+def register_function_group(config_type: type[FunctionGroupConfigT],
+                            framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
+    """
+    Register a function group with optional framework_wrappers for automatic profiler hooking.
+    Function groups share configuration/resources across multiple functions.
+    """
+
+    def register_function_group_inner(
+        fn: FunctionGroupBuildCallableT[FunctionGroupConfigT]
+    ) -> FunctionGroupRegisteredCallableT[FunctionGroupConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredFunctionGroupInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        framework_wrappers_list = list(framework_wrappers or [])
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.FUNCTION_GROUP)
+
+        GlobalTypeRegistry.get().register_function_group(
+            RegisteredFunctionGroupInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                framework_wrappers=framework_wrappers_list,
+                discovery_metadata=discovery_metadata,
+            ))
+
+        return context_manager_fn
+
+    return register_function_group_inner
+
+
 def register_llm_provider(config_type: type[LLMBaseConfigT]):
 
     def register_llm_provider_inner(

nat/cli/type_registry.py

@@ -37,6 +37,7 @@ from nat.builder.embedder import EmbedderProviderInfo
 from nat.builder.evaluator import EvaluatorInfo
 from nat.builder.front_end import FrontEndBase
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function_base import FunctionBase
 from nat.builder.function_info import FunctionInfo
 from nat.builder.llm import LLMProviderInfo
@@ -55,6 +56,8 @@ from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionBaseConfig
 from nat.data_models.function import FunctionConfigT
+from nat.data_models.function import FunctionGroupBaseConfig
+from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.logging import LoggingBaseConfig
@@ -85,6 +88,7 @@ EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncI
 EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
 FrontEndBuildCallableT = Callable[[FrontEndConfigT, Config], AsyncIterator[FrontEndBase]]
 FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
+FunctionGroupBuildCallableT = Callable[[FunctionGroupConfigT, Builder], AsyncIterator[FunctionGroup]]
 TTCStrategyBuildCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AsyncIterator[StrategyBase]]
 LLMClientBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[typing.Any]]
 LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
@@ -106,6 +110,7 @@ EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], Abs
 FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, Config], AbstractAsyncContextManager[FrontEndBase]]
 FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
                                        AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
+FunctionGroupRegisteredCallableT = Callable[[FunctionGroupConfigT, Builder], AbstractAsyncContextManager[FunctionGroup]]
 TTCStrategyRegisterCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
 LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
 LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
@@ -178,6 +183,16 @@ class RegisteredFunctionInfo(RegisteredInfo[FunctionBaseConfig]):
     framework_wrappers: list[str] = Field(default_factory=list)
 
 
+class RegisteredFunctionGroupInfo(RegisteredInfo[FunctionGroupBaseConfig]):
+    """
+    Represents a registered function group. Function groups are collections of functions that share configuration
+    and resources.
+    """
+
+    build_fn: FunctionGroupRegisteredCallableT = Field(repr=False)
+    framework_wrappers: list[str] = Field(default_factory=list)
+
+
 class RegisteredLLMProviderInfo(RegisteredInfo[LLMBaseConfig]):
     """
     Represents a registered LLM provider. LLM Providers are the operators of the LLMs. i.e. NIMs, OpenAI, Anthropic,
@@ -298,7 +313,7 @@ class RegisteredPackage(BaseModel):
     discovery_metadata: DiscoveryMetadata
 
 
-class TypeRegistry:  # pylint: disable=too-many-public-methods
+class TypeRegistry:
 
     def __init__(self) -> None:
         # Telemetry Exporters
@@ -313,6 +328,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         # Functions
         self._registered_functions: dict[type[FunctionBaseConfig], RegisteredFunctionInfo] = {}
 
+        # Function Groups
+        self._registered_function_groups: dict[type[FunctionGroupBaseConfig], RegisteredFunctionGroupInfo] = {}
+
         # LLMs
         self._registered_llm_provider_infos: dict[type[LLMBaseConfig], RegisteredLLMProviderInfo] = {}
         self._llm_client_provider_to_framework: dict[type[LLMBaseConfig], dict[str, RegisteredLLMClientInfo]] = {}
@@ -478,6 +496,50 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         return list(self._registered_functions.values())
 
+    def register_function_group(self, registration: RegisteredFunctionGroupInfo):
+        """Register a function group with the type registry.
+
+        Args:
+            registration: The function group registration information
+
+        Raises:
+            ValueError: If a function group with the same config type is already registered
+        """
+        if (registration.config_type in self._registered_function_groups):
+            raise ValueError(
+                f"A function group with the same config type `{registration.config_type}` has already been "
+                "registered.")
+
+        self._registered_function_groups[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def get_function_group(self, config_type: type[FunctionGroupBaseConfig]) -> RegisteredFunctionGroupInfo:
+        """Get a registered function group by its config type.
+
+        Args:
+            config_type: The function group configuration type
+
+        Returns:
+            RegisteredFunctionGroupInfo: The registered function group information
+
+        Raises:
+            KeyError: If no function group is registered for the given config type
+        """
+        try:
+            return self._registered_function_groups[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find a registered function group for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_function_groups.keys())}") from err
+
+    def get_registered_function_groups(self) -> list[RegisteredInfo[FunctionGroupBaseConfig]]:
+        """Get all registered function groups.
+
+        Returns:
+            list[RegisteredInfo[FunctionGroupBaseConfig]]: List of all registered function groups
+        """
+        return list(self._registered_function_groups.values())
+
     def register_llm_provider(self, info: RegisteredLLMProviderInfo):
 
         if (info.config_type in self._registered_llm_provider_infos):
@@ -779,7 +841,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         self._registration_changed()
 
-    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:  # pylint: disable=R0911
+    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:
 
         if component_type == ComponentEnum.FRONT_END:
             return self._registered_front_end_infos
@@ -790,6 +852,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         if component_type == ComponentEnum.FUNCTION:
             return self._registered_functions
 
+        if component_type == ComponentEnum.FUNCTION_GROUP:
+            return self._registered_function_groups
+
         if component_type == ComponentEnum.TOOL_WRAPPER:
             return self._registered_tool_wrappers
 
@@ -849,12 +914,14 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         raise ValueError(f"Supplied an unsupported component type {component_type}")
 
-    def get_registered_types_by_component_type(  # pylint: disable=R0911
-            self, component_type: ComponentEnum) -> list[str]:
+    def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]:
 
         if component_type == ComponentEnum.FUNCTION:
             return [i.static_type() for i in self._registered_functions]
 
+        if component_type == ComponentEnum.FUNCTION_GROUP:
+            return [i.static_type() for i in self._registered_function_groups]
+
         if component_type == ComponentEnum.TOOL_WRAPPER:
             return list(self._registered_tool_wrappers)
 
@@ -925,8 +992,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
             if (short_names[key.local_name] == 1):
                 type_list.append((key.local_name, key.config_type))
 
-        # pylint: disable=consider-alternative-union-syntax
-        return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+        return typing.Union[*tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
     def compute_annotation(self, cls: type[TypedBaseModelT]):
 
@@ -945,6 +1011,9 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         if issubclass(cls, FunctionBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_functions())
 
+        if issubclass(cls, FunctionGroupBaseConfig):
+            return self._do_compute_annotation(cls, self.get_registered_function_groups())
+
         if issubclass(cls, LLMBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_llm_providers())
 

nat/control_flow/register.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# flake8: noqa
+
+# Import any control flows which need to be automatically registered here
+from . import sequential_executor
+from .router_agent import register

nat/control_flow/router_agent/agent.py

@@ -0,0 +1,329 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+import typing
+
+from langchain_core.callbacks.base import AsyncCallbackHandler
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages.base import BaseMessage
+from langchain_core.messages.human import HumanMessage
+from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.agent.base import AGENT_CALL_LOG_MESSAGE
+from nat.agent.base import AGENT_LOG_PREFIX
+from nat.agent.base import BaseAgent
+
+if typing.TYPE_CHECKING:
+    from nat.control_flow.router_agent.register import RouterAgentWorkflowConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RouterAgentGraphState(BaseModel):
+    """State schema for the Router Agent Graph.
+
+    This class defines the state structure used throughout the Router Agent's
+    execution graph, containing messages, routing information, and branch selection.
+
+    Attributes:
+        messages: A list of messages representing the conversation history.
+        forward_message: The message to be forwarded to the chosen branch.
+        chosen_branch: The name of the branch selected by the router agent.
+    """
+    messages: list[BaseMessage] = Field(default_factory=list)
+    forward_message: BaseMessage = Field(default_factory=lambda: HumanMessage(content=""))
+    chosen_branch: str = Field(default="")
+
+
+class RouterAgentGraph(BaseAgent):
+    """Configurable Router Agent for routing requests to different branches.
+
+    A Router Agent analyzes incoming requests and routes them to one of the
+    configured branches based on the conte nt and context. It makes a single
+    routing decision and executes only the selected branch before returning.
+
+    This agent is useful for creating multi-path workflows where different
+    types of requests need to be handled by specialized sub-agents or tools.
+    """
+
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        branches: list[BaseTool],
+        prompt: ChatPromptTemplate,
+        max_router_retries: int = 3,
+        callbacks: list[AsyncCallbackHandler] | None = None,
+        detailed_logs: bool = False,
+        log_response_max_chars: int = 1000,
+    ):
+        """Initialize the Router Agent.
+
+        Args:
+            llm: The language model to use for routing decisions.
+            branches: List of tools/branches that the agent can route to.
+            prompt: The chat prompt template for the routing agent.
+            max_router_retries: Maximum number of retries if branch selection fails.
+            callbacks: Optional list of async callback handlers.
+            detailed_logs: Whether to enable detailed logging.
+            log_response_max_chars: Maximum characters to log in responses.
+        """
+        super().__init__(llm=llm,
+                         tools=branches,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
+
+        self._branches = branches
+        self._branches_dict = {branch.name: branch for branch in branches}
+        branch_names = ",".join([branch.name for branch in branches])
+        branch_names_and_descriptions = "\n".join([f"{branch.name}: {branch.description}" for branch in branches])
+
+        prompt = prompt.partial(branches=branch_names_and_descriptions, branch_names=branch_names)
+        self.agent = prompt | self.llm
+
+        self.max_router_retries = max_router_retries
+
+    def _get_branch(self, branch_name: str) -> BaseTool | None:
+        return self._branches_dict.get(branch_name, None)
+
+    async def agent_node(self, state: RouterAgentGraphState):
+        """Execute the agent node to select a branch for routing.
+
+        This method processes the incoming request and determines which branch
+        should handle it. It uses the configured LLM to analyze the request
+        and select the most appropriate branch.
+
+        Args:
+            state: The current state of the router agent graph.
+
+        Returns:
+            RouterAgentGraphState: Updated state with the chosen branch.
+
+        Raises:
+            RuntimeError: If the agent fails to choose a branch after max retries.
+        """
+        logger.debug("%s Starting the Router Agent Node", AGENT_LOG_PREFIX)
+        chat_history = self._get_chat_history(state.messages)
+        request = state.forward_message.content
+        for attempt in range(1, self.max_router_retries + 1):
+            try:
+                agent_response = await self._call_llm(self.agent, {"request": request, "chat_history": chat_history})
+                if self.detailed_logs:
+                    logger.info(AGENT_CALL_LOG_MESSAGE, request, agent_response)
+
+                state.messages += [agent_response]
+
+                # Determine chosen branch based on agent response
+                if state.chosen_branch == "":
+                    for branch in self._branches:
+                        if branch.name.lower() in str(agent_response.content).lower():
+                            state.chosen_branch = branch.name
+                            if self.detailed_logs:
+                                logger.debug("%s Router Agent has chosen branch: %s", AGENT_LOG_PREFIX, branch.name)
+                            return state
+
+                # The agent failed to choose a branch
+                if state.chosen_branch == "":
+                    if attempt == self.max_router_retries:
+                        logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
+                        raise RuntimeError("Router Agent failed to choose a branch")
+                    logger.warning("%s Router Agent failed to choose a branch, retrying %d out of %d",
+                                   AGENT_LOG_PREFIX,
+                                   attempt,
+                                   self.max_router_retries)
+
+            except Exception as ex:
+                logger.error("%s Router Agent failed to call agent_node: %s", AGENT_LOG_PREFIX, ex)
+                raise
+
+        return state
+
+    async def branch_node(self, state: RouterAgentGraphState):
+        """Execute the selected branch with the forwarded message.
+
+        This method calls the tool/branch that was selected by the agent node
+        and processes the response.
+
+        Args:
+            state: The current state containing the chosen branch and message.
+
+        Returns:
+            RouterAgentGraphState: Updated state with the branch response.
+
+        Raises:
+            RuntimeError: If no branch was chosen or branch execution fails.
+            ValueError: If the requested tool is not found in the configuration.
+        """
+        logger.debug("%s Starting Router Agent Tool Node", AGENT_LOG_PREFIX)
+        try:
+            if state.chosen_branch == "":
+                logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
+                raise RuntimeError("Router Agent failed to choose a branch")
+            requested_branch = self._get_branch(state.chosen_branch)
+            if not requested_branch:
+                logger.error("%s Router Agent wants to call tool %s but it is not in the config file",
+                             AGENT_LOG_PREFIX,
+                             state.chosen_branch)
+                raise ValueError("Tool not found in config file")
+
+            branch_input = state.forward_message.content
+            branch_response = await self._call_tool(requested_branch, branch_input)
+            state.messages += [branch_response]
+            if self.detailed_logs:
+                self._log_tool_response(requested_branch.name, branch_input, branch_response.content)
+
+            return state
+
+        except Exception as ex:
+            logger.error("%s Router Agent throws exception during branch node execution: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    async def _build_graph(self, state_schema):
+        logger.debug("%s Building and compiling the Router Agent Graph", AGENT_LOG_PREFIX)
+
+        graph = StateGraph(state_schema)
+        graph.add_node("agent", self.agent_node)
+        graph.add_node("branch", self.branch_node)
+        graph.add_edge("agent", "branch")
+        graph.set_entry_point("agent")
+
+        self.graph = graph.compile()
+        logger.debug("%s Router Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
+
+        return self.graph
+
+    async def build_graph(self):
+        """Build and compile the router agent execution graph.
+
+        Creates a state graph with agent and branch nodes, configures the
+        execution flow, and compiles the graph for execution.
+
+        Returns:
+            The compiled execution graph.
+
+        Raises:
+            Exception: If graph building or compilation fails.
+        """
+        try:
+            await self._build_graph(state_schema=RouterAgentGraphState)
+            return self.graph
+        except Exception as ex:
+            logger.error("%s Router Agent failed to build graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    @staticmethod
+    def validate_system_prompt(system_prompt: str) -> bool:
+        """Validate that the system prompt contains required variables.
+
+        Checks that the system prompt includes necessary template variables
+        for branch information that the router agent needs.
+
+        Args:
+            system_prompt: The system prompt string to validate.
+
+        Returns:
+            True if the prompt is valid, False otherwise.
+        """
+        errors = []
+        required_prompt_variables = {
+            "{branches}": "The system prompt must contain {branches} so the agent knows about configured branches.",
+            "{branch_names}": "The system prompt must contain {branch_names} so the agent knows branch names."
+        }
+        for variable_name, error_message in required_prompt_variables.items():
+            if variable_name not in system_prompt:
+                errors.append(error_message)
+        if errors:
+            error_text = "\n".join(errors)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
+            return False
+        return True
+
+    @staticmethod
+    def validate_user_prompt(user_prompt: str) -> bool:
+        """Validate that the user prompt contains required variables.
+
+        Checks that the user prompt includes necessary template variables
+        for chat history and other required information.
+
+        Args:
+            user_prompt: The user prompt string to validate.
+
+        Returns:
+            True if the prompt is valid, False otherwise.
+        """
+        errors = []
+        if not user_prompt:
+            errors.append("The user prompt cannot be empty.")
+        else:
+            required_prompt_variables = {
+                "{chat_history}":
+                    "The user prompt must contain {chat_history} so the agent knows about the conversation history.",
+                "{request}":
+                    "The user prompt must contain {request} so the agent sees the current request.",
+            }
+            for variable_name, error_message in required_prompt_variables.items():
+                if variable_name not in user_prompt:
+                    errors.append(error_message)
+        if errors:
+            error_text = "\n".join(errors)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
+            return False
+        return True
+
+
+def create_router_agent_prompt(config: "RouterAgentWorkflowConfig") -> ChatPromptTemplate:
+    """Create a Router Agent prompt from the configuration.
+
+    Builds a ChatPromptTemplate using either custom prompts from the config
+    or default system and user prompts. Validates the prompts to ensure they
+    contain required template variables.
+
+    Args:
+        config: The router agent workflow configuration containing prompt settings.
+
+    Returns:
+        A configured ChatPromptTemplate for the router agent.
+
+    Raises:
+        ValueError: If the system_prompt or user_prompt validation fails.
+    """
+    from nat.control_flow.router_agent.prompt import SYSTEM_PROMPT
+    from nat.control_flow.router_agent.prompt import USER_PROMPT
+    # the Router Agent prompt can be customized via config option system_prompt and user_prompt.
+
+    if config.system_prompt:
+        system_prompt = config.system_prompt
+    else:
+        system_prompt = SYSTEM_PROMPT
+
+    if config.user_prompt:
+        user_prompt = config.user_prompt
+    else:
+        user_prompt = USER_PROMPT
+
+    if not RouterAgentGraph.validate_system_prompt(system_prompt):
+        logger.error("%s Invalid system_prompt", AGENT_LOG_PREFIX)
+        raise ValueError("Invalid system_prompt")
+
+    if not RouterAgentGraph.validate_user_prompt(user_prompt):
+        logger.error("%s Invalid user_prompt", AGENT_LOG_PREFIX)
+        raise ValueError("Invalid user_prompt")
+
+    return ChatPromptTemplate([("system", system_prompt), ("user", user_prompt)])

nat/control_flow/router_agent/prompt.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+SYSTEM_PROMPT = """
+You are a Router Agent responsible for analyzing incoming requests and routing them to the most appropriate branch.
+
+Available branches:
+{branches}
+
+CRITICAL INSTRUCTIONS:
+- Analyze the user's request carefully
+- Select exactly ONE branch that best handles the request from: [{branch_names}]
+- Respond with ONLY the exact branch name, nothing else
+- Be decisive - choose the single best match, if the request could fit multiple branches,
+  choose the most specific/specialized one
+- If no branch perfectly fits, choose the closest match
+
+Your response MUST contain ONLY the branch name. Do not include any explanations, reasoning, or additional text.
+
+Examples:
+User: "How do I calculate 15 + 25?"
+Response: calculator_tool
+
+User: "What's the weather like today?"
+Response: weather_service
+
+User: "Send an email to John"
+Response: email_tool"""
+
+USER_PROMPT = """
+Previous conversation history:
+{chat_history}
+
+To respond to the request: {request}, which branch should be chosen?
+
+Respond with only the branch name."""