nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/agent/react_agent/output_parser.py

@@ -0,0 +1,104 @@
+# 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 re
+
+from langchain.agents.agent import AgentOutputParser
+from langchain_core.agents import AgentAction
+from langchain_core.agents import AgentFinish
+from langchain_core.exceptions import LangChainException
+
+from .prompt import SYSTEM_PROMPT
+
+FINAL_ANSWER_ACTION = "Final Answer:"
+MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE = "Invalid Format: Missing 'Action:' after 'Thought:'"
+MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE = "Invalid Format: Missing 'Action Input:' after 'Action:'"
+FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE = ("Parsing LLM output produced both a final answer and a parse-able "
+                                                  "action:")
+
+
+class ReActOutputParserException(ValueError, LangChainException):
+
+    def __init__(self,
+                 observation=None,
+                 missing_action=False,
+                 missing_action_input=False,
+                 final_answer_and_action=False):
+        self.observation = observation
+        self.missing_action = missing_action
+        self.missing_action_input = missing_action_input
+        self.final_answer_and_action = final_answer_and_action
+
+
+class ReActOutputParser(AgentOutputParser):
+    """Parses ReAct-style LLM calls that have a single tool input.
+
+    Expects output to be in one of two formats.
+
+    If the output signals that an action should be taken,
+    should be in the below format. This will result in an AgentAction
+    being returned.
+
+    ```
+    Thought: agent thought here
+    Action: search
+    Action Input: what is the temperature in SF?
+    Observation: Waiting for the tool response...
+    ```
+
+    If the output signals that a final answer should be given,
+    should be in the below format. This will result in an AgentFinish
+    being returned.
+
+    ```
+    Thought: agent thought here
+    Final Answer: The temperature is 100 degrees
+    ```
+
+    """
+
+    def get_format_instructions(self) -> str:
+        return SYSTEM_PROMPT
+
+    def parse(self, text: str) -> AgentAction | AgentFinish:
+        includes_answer = FINAL_ANSWER_ACTION in text
+        regex = r"Action\s*\d*\s*:[\s]*(.*?)\s*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*?)(?=\s*[\n|\s]\s*Observation\b|$)"
+        action_match = re.search(regex, text, re.DOTALL)
+        if action_match:
+            if includes_answer:
+                raise ReActOutputParserException(
+                    final_answer_and_action=True,
+                    observation=f"{FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE}: {text}")
+            action = action_match.group(1).strip()
+            action_input = action_match.group(2)
+            tool_input = action_input.strip(" ")
+            tool_input = tool_input.strip('"')
+
+            return AgentAction(action, tool_input, text)
+
+        if includes_answer:
+            return AgentFinish({"output": text.split(FINAL_ANSWER_ACTION)[-1].strip()}, text)
+
+        if not re.search(r"Action\s*\d*\s*:[\s]*(.*?)", text, re.DOTALL):
+            raise ReActOutputParserException(observation=MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE,
+                                             missing_action=True)
+        if not re.search(r"[\s]*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)", text, re.DOTALL):
+            raise ReActOutputParserException(observation=MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE,
+                                             missing_action_input=True)
+        raise ReActOutputParserException(f"Could not parse LLM output: `{text}`")
+
+    @property
+    def _type(self) -> str:
+        return "react-input"

aiq/agent/react_agent/prompt.py

@@ -0,0 +1,41 @@
+# 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
+
+SYSTEM_PROMPT = """
+Answer the following questions as best you can. You may ask the human to use the following tools:
+
+{tools}
+
+You may respond in one of two formats.
+Use the following format exactly to ask the human to use a tool:
+
+Question: the input question you must answer
+Thought: you should always think about what to do
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action (if there is no required input, include "Action Input: None")  
+Observation: wait for the human to respond with the result from the tool, do not assume the response
+
+... (this Thought/Action/Action Input/Observation can repeat N times. If you do not need to use a tool, or after asking the human to use any tools and waiting for the human to respond, you might know the final answer.)
+Use the following format once you have the final answer:
+
+Thought: I now know the final answer
+Final Answer: the final answer to the original input question
+"""
+
+USER_PROMPT = """
+Question: {question}
+"""

aiq/agent/react_agent/register.py

@@ -0,0 +1,149 @@
+# 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 aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.api_server import AIQChatRequest
+from aiq.data_models.api_server import AIQChatResponse
+from aiq.data_models.component_ref import FunctionRef
+from aiq.data_models.component_ref import LLMRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.utils.type_converter import GlobalTypeConverter
+
+logger = logging.getLogger(__name__)
+
+
+class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
+    """
+    Defines an AIQ Toolkit function that uses a ReAct Agent performs reasoning inbetween tool calls, and utilizes the
+    tool names and descriptions to select the optimal tool.
+    """
+
+    tool_names: list[FunctionRef] = Field(default_factory=list,
+                                          description="The list of tools to provide to the react agent.")
+    llm_name: LLMRef = Field(description="The LLM model to use with the react agent.")
+    verbose: bool = Field(default=False, description="Set the verbosity of the react agent's logging.")
+    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.")
+    description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
+    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.")
+    use_openai_api: bool = Field(default=False,
+                                 description=("Use OpenAI API for the input/output types to the function. "
+                                              "If False, strings will be used."))
+    additional_instructions: str | None = Field(
+        default=None, description="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.graph import CompiledGraph
+
+    from aiq.agent.base import AGENT_LOG_PREFIX
+    from aiq.agent.react_agent.agent import ReActAgentGraph
+    from aiq.agent.react_agent.agent import ReActGraphState
+    from aiq.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 = 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: CompiledGraph = await ReActAgentGraph(
+        llm=llm,
+        prompt=prompt,
+        tools=tools,
+        use_tool_schema=config.include_tool_input_schema_in_tool_description,
+        detailed_logs=config.verbose,
+        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).build_graph()
+
+    async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
+        try:
+            # initialize the starting state with the user query
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_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]  # pylint: disable=E1136
+            return AIQChatResponse.from_string(str(output_message.content))
+
+        except Exception as ex:
+            logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
+            # here, we can implement custom error messages
+            if config.verbose:
+                return AIQChatResponse.from_string(str(ex))
+            return AIQChatResponse.from_string("I seem to be having a problem.")
+
+    if (config.use_openai_api):
+        yield FunctionInfo.from_fn(_response_fn, description=config.description)
+    else:
+
+        async def _str_api_fn(input_message: str) -> str:
+            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
+
+            oai_output = await _response_fn(oai_input)
+
+            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
+
+        yield FunctionInfo.from_fn(_str_api_fn, description=config.description)

aiq/agent/reasoning_agent/reasoning_agent.py

@@ -0,0 +1,225 @@
+# 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 aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.api_server import AIQChatRequest
+from aiq.data_models.component_ref import FunctionRef
+from aiq.data_models.component_ref import LLMRef
+from aiq.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ReasoningFunctionConfig(FunctionBaseConfig, name="reasoning_agent"):
+    """
+    Defines an AIQ Toolkit 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.
+    """
+
+    llm_name: LLMRef = Field(description="The name of the LLM to use for reasoning.")
+    augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
+    verbose: bool = Field(default=False, description="Whether to log detailed information.")
+    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 aiq.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 = 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_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    tool_names_with_desc: list[tuple[str, str]] = []
+
+    for tool in function_used_tools:
+        tool_impl = 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: AIQChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (AIQChatRequest): 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: AIQChatRequest) -> augmented_function.single_output_type:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (AIQChatRequest): 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)

aiq/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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+
+# Import any workflows which need to be automatically registered here
+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