aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/agent/react_agent/register.py

@@ -0,0 +1,148 @@
+# 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 Field
+
+from aiq.agent.base import AGENT_LOG_PREFIX
+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 AgentIQ 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_parsing_errors: bool = Field(default=True, description="Specify retrying when encountering parsing errors.")
+    max_retries: int = Field(default=1, description="Sent the number of retries before raising a parsing error.")
+    include_tool_input_schema_in_tool_description: bool = Field(
+        default=True, description="Specify inclusion of tool input schemas in the prompt.")
+    max_iterations: int = Field(default=15, description="Number of tool calls before stoping the react agent.")
+    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 langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.prompts import MessagesPlaceholder
+    from langgraph.graph.graph import CompiledGraph
+
+    from aiq.agent.react_agent.prompt import USER_PROMPT
+
+    from .agent import ReActAgentGraph
+    from .agent import ReActGraphState
+    from .prompt import react_agent_prompt
+
+    # the ReAct Agent prompt comes from prompt.py, and can be customized there or via config option system_prompt.
+    if config.system_prompt:
+        _prompt_str = config.system_prompt
+        if config.additional_instructions:
+            _prompt_str += f" {config.additional_instructions}"
+        valid_prompt = ReActAgentGraph.validate_system_prompt(config.system_prompt)
+        if not valid_prompt:
+            logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
+            raise ValueError("Invalid system_prompt")
+        prompt = ChatPromptTemplate([("system", config.system_prompt), ("user", USER_PROMPT),
+                                     MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
+    else:
+        prompt = react_agent_prompt
+
+    # 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_parsing_errors=config.retry_parsing_errors,
+                                                 max_retries=config.max_retries).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_iterations + 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(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().convert(input_message, to_type=AIQChatRequest)
+
+            oai_output = await _response_fn(oai_input)
+
+            return GlobalTypeConverter.get().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,224 @@
+# 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.agent.base import AGENT_LOG_PREFIX
+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 AgentIQ 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
+
+    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