nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/agent/react_agent/register.py

@@ -0,0 +1,168 @@
+# 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
+
+from pydantic import AliasChoices
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.agent import AgentBaseConfig
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatRequestOrMessage
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
+from nat.data_models.component_ref import FunctionGroupRef
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
+from nat.utils.type_converter import GlobalTypeConverter
+
+logger = logging.getLogger(__name__)
+
+
+class ReActAgentWorkflowConfig(AgentBaseConfig, OptimizableMixin, name="react_agent"):
+    """
+    Defines a NAT function that uses a ReAct Agent performs reasoning inbetween tool calls, and utilizes the
+    tool names and descriptions to select the optimal tool.
+    """
+    description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        default_factory=list, description="The list of tools to provide to the react agent.")
+    retry_agent_response_parsing_errors: bool = Field(
+        default=True,
+        validation_alias=AliasChoices("retry_agent_response_parsing_errors", "retry_parsing_errors"),
+        description="Whether to retry when encountering parsing errors in the agent's response.")
+    parse_agent_response_max_retries: int = Field(
+        default=1,
+        validation_alias=AliasChoices("parse_agent_response_max_retries", "max_retries"),
+        description="Maximum number of times the Agent may retry parsing errors. "
+        "Prevents the Agent from getting into infinite hallucination loops.")
+    tool_call_max_retries: int = Field(default=1, description="The number of retries before raising a tool call error.")
+    max_tool_calls: int = Field(default=15,
+                                validation_alias=AliasChoices("max_tool_calls", "max_iterations"),
+                                description="Maximum number of tool calls before stopping the agent.")
+    pass_tool_call_errors_to_agent: bool = Field(
+        default=True,
+        description="Whether to pass tool call errors to agent. If False, failed tool calls will raise an exception.")
+    include_tool_input_schema_in_tool_description: bool = Field(
+        default=True, description="Specify inclusion of tool input schemas in the prompt.")
+    normalize_tool_input_quotes: bool = Field(
+        default=True,
+        description="Whether to replace single quotes with double quotes in the tool input. "
+        "This is useful for tools that expect structured json input.")
+    system_prompt: str | None = Field(
+        default=None,
+        description="Provides the SYSTEM_PROMPT to use with the agent")  # defaults to SYSTEM_PROMPT in prompt.py
+    max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+    additional_instructions: str | None = OptimizableField(
+        default=None,
+        description="Additional instructions to provide to the agent in addition to the base prompt.",
+        space=SearchSpace(
+            is_prompt=True,
+            prompt="No additional instructions.",
+            prompt_purpose="Additional instructions to provide to the agent in addition to the base prompt.",
+        ))
+
+
+@register_function(config_type=ReActAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
+    from langchain.schema import BaseMessage
+    from langchain_core.messages import trim_messages
+    from langgraph.graph.state import CompiledStateGraph
+
+    from nat.agent.base import AGENT_LOG_PREFIX
+    from nat.agent.react_agent.agent import ReActAgentGraph
+    from nat.agent.react_agent.agent import ReActGraphState
+    from nat.agent.react_agent.agent import create_react_agent_prompt
+
+    prompt = create_react_agent_prompt(config)
+
+    # we can choose an LLM for the ReAct agent in the config file
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    # the agent can run any installed tool, simply install the tool and add it to the config file
+    # the sample tool provided can easily be copied or changed
+    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    if not tools:
+        raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'")
+    # configure callbacks, for sending intermediate steps
+    # construct the ReAct Agent Graph from the configured llm, prompt, and tools
+    graph: CompiledStateGraph = await ReActAgentGraph(
+        llm=llm,
+        prompt=prompt,
+        tools=tools,
+        use_tool_schema=config.include_tool_input_schema_in_tool_description,
+        detailed_logs=config.verbose,
+        log_response_max_chars=config.log_response_max_chars,
+        retry_agent_response_parsing_errors=config.retry_agent_response_parsing_errors,
+        parse_agent_response_max_retries=config.parse_agent_response_max_retries,
+        tool_call_max_retries=config.tool_call_max_retries,
+        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent,
+        normalize_tool_input_quotes=config.normalize_tool_input_quotes).build_graph()
+
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatResponse | str:
+        """
+        Main workflow entry function for the ReAct Agent.
+
+        This function invokes the ReAct Agent Graph and returns the response.
+
+        Args:
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
+
+        Returns:
+            ChatResponse | str: The response from the agent or error message
+        """
+        try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+
+            # initialize the starting state with the user query
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
+                                                        max_tokens=config.max_history,
+                                                        strategy="last",
+                                                        token_counter=len,
+                                                        start_on="human",
+                                                        include_system=True)
+
+            state = ReActGraphState(messages=messages)
+
+            # run the ReAct Agent Graph
+            state = await graph.ainvoke(state, config={'recursion_limit': (config.max_tool_calls + 1) * 2})
+            # setting recursion_limit: 4 allows 1 tool call
+            #   - allows the ReAct Agent to perform 1 cycle / call 1 single tool,
+            #   - but stops the agent when it tries to call a tool a second time
+
+            # get and return the output from the state
+            state = ReActGraphState(**state)
+            output_message = state.messages[-1]
+            content = str(output_message.content)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+            completion_tokens = len(content.split()) if content else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            response = ChatResponse.from_string(content, usage=usage)
+            if chat_request_or_message.is_string:
+                return GlobalTypeConverter.get().convert(response, to_type=str)
+            return response
+        except Exception as ex:
+            logger.error("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
+            raise
+
+    yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/reasoning_agent/reasoning_agent.py

@@ -0,0 +1,227 @@
+# 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 re
+from collections.abc import AsyncGenerator
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.agent import AgentBaseConfig
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.component_ref import FunctionRef
+
+logger = logging.getLogger(__name__)
+
+
+class ReasoningFunctionConfig(AgentBaseConfig, name="reasoning_agent"):
+    """
+    Defines a NAT function that performs reasoning on the input data.
+    Output is passed to the next function in the workflow.
+
+    Designed to be used with an InterceptingFunction.
+    """
+    description: str = Field(default="Reasoning Agent", description="The description of this function's use.")
+    augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
+    reasoning_prompt_template: str = Field(
+        default=("You are an expert reasoning model task with creating a detailed execution plan"
+                 " for a system that has the following description:\n\n"
+                 "**Description:** \n{augmented_function_desc}\n\n"
+                 "Given the following input and a list of available tools, please provide a detailed step-by-step plan"
+                 " that an instruction following system can use to address the input. Ensure the plan includes:\n\n"
+                 "1. Identifying the key components of the input.\n"
+                 "2. Determining the most suitable tools for each task.\n"
+                 "3. Outlining the sequence of actions to be taken.\n\n"
+                 "**Input:** \n{input_text}\n\n"
+                 "**Tools and description of the tool:** \n{tools}\n\n"
+                 "An example plan could look like this:\n\n"
+                 "1. Call tool A with input X\n"
+                 "2. Call tool B with input Y\n"
+                 "3. Interpret the output of tool A and B\n"
+                 "4. Return the final result"
+                 "\n\n **PLAN:**\n"),
+        description="The reasoning model prompt template.")
+
+    instruction_prompt_template: str = Field(
+        default=("Answer the following question based on message history: {input_text}"
+                 "\n\nHere is a plan for execution that you could use to guide you if you wanted to:"
+                 "\n\n{reasoning_output}"
+                 "\n\nNOTE: Remember to follow your guidance on how to format output, etc."
+                 "\n\n You must respond with the answer to the original question directly to the user."),
+        description="The instruction prompt template.")
+
+
+@register_function(config_type=ReasoningFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Builder):
+    """
+    Build a ReasoningFunction from the provided config.
+
+    Args:
+        config (ReasoningFunctionConfig): The config for the ReasoningFunction.
+        builder (Builder): The Builder instance to use for building the function.
+
+    Returns:
+        ReasoningFunction: The built ReasoningFunction.
+    """
+    from langchain_core.language_models import BaseChatModel
+    from langchain_core.prompts import PromptTemplate
+
+    from nat.agent.base import AGENT_LOG_PREFIX
+
+    def remove_r1_think_tags(text: str):
+        pattern = r'(<think>)?.*?</think>\s*(.*)'
+
+        # Add re.DOTALL flag to make . match newlines
+        match = re.match(pattern, text, re.DOTALL)
+
+        if match:
+            return match.group(2)
+
+        return text
+
+    # Get the LLM to use for reasoning
+    llm: BaseChatModel = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    # Get the augmented function's description
+    augmented_function = await builder.get_function(config.augmented_fn)
+
+    # For now, we rely on runtime checking for type conversion
+
+    if augmented_function.description and augmented_function.description != "":
+        augmented_function_desc = augmented_function.description
+    else:
+        raise ValueError(f"Function {config.augmented_fn} does not have a description. Cannot augment "
+                         f"function without a description.")
+
+    # Get the function dependencies of the augmented function
+    function_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set()
+    function_used_tools.update(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
+    tool_names_with_desc: list[tuple[str, str]] = []
+
+    for tool in function_used_tools:
+        tool_impl = await builder.get_function(tool)
+        tool_names_with_desc.append((tool, tool_impl.description if hasattr(tool_impl, "description") else ""))
+
+    # Draft the reasoning prompt for the augmented function
+    template = PromptTemplate(template=config.reasoning_prompt_template,
+                              input_variables=["augmented_function_desc", "input_text", "tools"],
+                              validate_template=True)
+
+    downstream_template = PromptTemplate(template=config.instruction_prompt_template,
+                                         input_variables=["input_text", "reasoning_output"],
+                                         validate_template=True)
+
+    streaming_inner_fn = None
+    single_inner_fn = None
+
+    if augmented_function.has_streaming_output:
+
+        async def streaming_inner(
+                input_message: ChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (ChatRequest): The input text to reason on.
+            """
+
+            input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
+
+            prompt = await template.ainvoke(
+                input={
+                    "augmented_function_desc": augmented_function_desc,
+                    "input_text": input_text,
+                    "tools": "\n".join([f"- {tool[0]}: {tool[1]}" for tool in tool_names_with_desc])
+                })
+
+            prompt = prompt.to_string()
+
+            # Get the reasoning output from the LLM
+            reasoning_output = ""
+
+            async for chunk in llm.astream(prompt):
+                reasoning_output += chunk.content
+
+            reasoning_output = remove_r1_think_tags(reasoning_output)
+
+            output = await downstream_template.ainvoke(input={
+                "input_text": input_text, "reasoning_output": reasoning_output
+            })
+
+            output = output.to_string()
+
+            if config.verbose:
+                logger.info("%s Reasoning plan and input to agent: \n\n%s", AGENT_LOG_PREFIX, output)
+
+            async for chunk in augmented_function.acall_stream(output):
+                yield chunk
+
+        streaming_inner_fn = streaming_inner
+
+    if augmented_function.has_single_output:
+
+        async def single_inner(input_message: ChatRequest) -> augmented_function.single_output_type:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (ChatRequest): The input text to reason on.
+            """
+
+            input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
+
+            prompt = await template.ainvoke(
+                input={
+                    "augmented_function_desc": augmented_function_desc,
+                    "input_text": input_text,
+                    "tools": "\n".join([f"- {tool[0]}: {tool[1]}" for tool in tool_names_with_desc])
+                })
+
+            prompt = prompt.to_string()
+
+            # Get the reasoning output from the LLM
+            reasoning_output = ""
+
+            async for chunk in llm.astream(prompt):
+                reasoning_output += chunk.content
+
+            reasoning_output = remove_r1_think_tags(reasoning_output)
+
+            output = await downstream_template.ainvoke(input={
+                "input_text": input_text, "reasoning_output": reasoning_output
+            })
+
+            output = output.to_string()
+
+            if config.verbose:
+                logger.info("%s Reasoning plan and input to agent: \n\n%s", AGENT_LOG_PREFIX, output)
+
+            return await augmented_function.acall_invoke(output)
+
+        single_inner_fn = single_inner
+
+    yield FunctionInfo.create(
+        single_fn=single_inner_fn,
+        stream_fn=streaming_inner_fn,
+        description=("Reasoning function that generates a detailed execution plan for a system based on the input."),
+        converters=augmented_function.converter_list)

nat/agent/register.py

@@ -0,0 +1,23 @@
+# 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.
+
+# flake8: noqa
+
+# Import any workflows which need to be automatically registered here
+from .prompt_optimizer import register as prompt_optimizer
+from .react_agent import register as react_agent
+from .reasoning_agent import reasoning_agent
+from .rewoo_agent import register as rewoo_agent
+from .tool_calling_agent import register as tool_calling_agent