nvidia-nat 1.1.0a20251020__py3-none-any.whl

aiq/__init__.py

@@ -0,0 +1,66 @@
+# 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 importlib
+import importlib.abc
+import importlib.util
+import sys
+import warnings
+
+
+class CompatFinder(importlib.abc.MetaPathFinder):
+
+    def __init__(self, alias_prefix, target_prefix):
+        self.alias_prefix = alias_prefix
+        self.target_prefix = target_prefix
+
+    def find_spec(self, fullname, path, target=None):
+        if fullname == self.alias_prefix or fullname.startswith(self.alias_prefix + "."):
+            # Map aiq.something -> nat.something
+            target_name = self.target_prefix + fullname[len(self.alias_prefix):]
+            spec = importlib.util.find_spec(target_name)
+            if spec is None:
+                return None
+            # Wrap the loader so it loads under the alias name
+            return importlib.util.spec_from_loader(fullname, CompatLoader(fullname, target_name))
+        return None
+
+
+class CompatLoader(importlib.abc.Loader):
+
+    def __init__(self, alias_name, target_name):
+        self.alias_name = alias_name
+        self.target_name = target_name
+
+    def create_module(self, spec):
+        # Reuse the actual module so there's only one instance
+        target_module = importlib.import_module(self.target_name)
+        sys.modules[self.alias_name] = target_module
+        return target_module
+
+    def exec_module(self, module):
+        # Nothing to execute since the target is already loaded
+        pass
+
+
+# Register the compatibility finder
+sys.meta_path.insert(0, CompatFinder("aiq", "nat"))
+
+warnings.warn(
+    "!!! The 'aiq' namespace is deprecated and will be removed in a future release. "
+    "Please use the 'nat' namespace instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)

nat/agent/base.py

@@ -0,0 +1,265 @@
+# 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 asyncio
+import json
+import logging
+from abc import ABC
+from abc import abstractmethod
+from enum import Enum
+from typing import Any
+
+from colorama import Fore
+from langchain_core.callbacks import AsyncCallbackHandler
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage
+from langchain_core.messages import BaseMessage
+from langchain_core.messages import ToolMessage
+from langchain_core.runnables import Runnable
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from langgraph.graph.state import CompiledStateGraph
+
+logger = logging.getLogger(__name__)
+
+TOOL_NOT_FOUND_ERROR_MESSAGE = "There is no tool named {tool_name}. Tool must be one of {tools}."
+INPUT_SCHEMA_MESSAGE = ". Arguments must be provided as a valid JSON object following this format: {schema}"
+NO_INPUT_ERROR_MESSAGE = "No human input received to the agent, Please ask a valid question."
+
+AGENT_LOG_PREFIX = "[AGENT]"
+AGENT_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
+                                 AGENT_LOG_PREFIX + "\n" + \
+                                 Fore.YELLOW + \
+                                 "Agent input: %s\n" + \
+                                 Fore.CYAN + \
+                                 "Agent's thoughts: \n%s" + \
+                                 Fore.RESET + \
+                                 f"\n{'-' * 30}"
+
+TOOL_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
+                                 AGENT_LOG_PREFIX + "\n" + \
+                                 Fore.WHITE + \
+                                 "Calling tools: %s\n" + \
+                                 Fore.YELLOW + \
+                                 "Tool's input: %s\n" + \
+                                 Fore.CYAN + \
+                                 "Tool's response: \n%s" + \
+                                 Fore.RESET + \
+                                 f"\n{'-' * 30}"
+
+
+class AgentDecision(Enum):
+    TOOL = "tool"
+    END = "finished"
+
+
+class BaseAgent(ABC):
+
+    def __init__(self,
+                 llm: BaseChatModel,
+                 tools: list[BaseTool],
+                 callbacks: list[AsyncCallbackHandler] | None = None,
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000) -> None:
+        logger.debug("Initializing Agent Graph")
+        self.llm = llm
+        self.tools = tools
+        self.callbacks = callbacks or []
+        self.detailed_logs = detailed_logs
+        self.log_response_max_chars = log_response_max_chars
+        self.graph = None
+
+    async def _stream_llm(self,
+                          runnable: Any,
+                          inputs: dict[str, Any],
+                          config: RunnableConfig | None = None) -> AIMessage:
+        """
+        Stream from LLM runnable. Retry logic is handled automatically by the underlying LLM client.
+
+        Parameters
+        ----------
+        runnable : Any
+            The LLM runnable (prompt | llm or similar)
+        inputs : Dict[str, Any]
+            The inputs to pass to the runnable
+        config : RunnableConfig | None
+            The config to pass to the runnable (should include callbacks)
+
+        Returns
+        -------
+        AIMessage
+            The LLM response
+        """
+        output_message = ""
+        async for event in runnable.astream(inputs, config=config):
+            output_message += event.content
+
+        return AIMessage(content=output_message)
+
+    async def _call_llm(self, llm: Runnable, inputs: dict[str, Any], config: RunnableConfig | None = None) -> AIMessage:
+        """
+        Call the LLM directly. Retry logic is handled automatically by the underlying LLM client.
+
+        Parameters
+        ----------
+        llm : Runnable
+            The LLM runnable (prompt | llm or similar)
+        inputs : dict[str, Any]
+            The inputs to pass to the runnable
+        config : RunnableConfig | None
+            The config to pass to the runnable (should include callbacks)
+
+        Returns
+        -------
+        AIMessage
+            The LLM response
+        """
+        response = await llm.ainvoke(inputs, config=config)
+        return AIMessage(content=str(response.content))
+
+    async def _call_tool(self,
+                         tool: BaseTool,
+                         tool_input: dict[str, Any] | str,
+                         config: RunnableConfig | None = None,
+                         max_retries: int = 3) -> ToolMessage:
+        """
+        Call a tool with retry logic and error handling.
+
+        Parameters
+        ----------
+        tool : BaseTool
+            The tool to call
+        tool_input : Union[Dict[str, Any], str]
+            The input to pass to the tool
+        config : RunnableConfig | None
+            The config to pass to the tool
+        max_retries : int
+            Maximum number of retry attempts (default: 3)
+
+        Returns
+        -------
+        ToolMessage
+            The tool response
+        """
+        last_exception = None
+
+        for attempt in range(1, max_retries + 1):
+            try:
+                response = await tool.ainvoke(tool_input, config=config)
+
+                # Handle empty responses
+                if response is None or (isinstance(response, str) and response == ""):
+                    return ToolMessage(name=tool.name,
+                                       tool_call_id=tool.name,
+                                       content=f"The tool {tool.name} provided an empty response.")
+
+                # ToolMessage only accepts str or list[str | dict] as content.
+                # Convert into list if the response is a dict.
+                if isinstance(response, dict):
+                    response = [response]
+
+                return ToolMessage(name=tool.name, tool_call_id=tool.name, content=response)
+
+            except Exception as e:
+                last_exception = e
+
+                # If this was the last attempt, don't sleep
+                if attempt == max_retries:
+                    break
+
+                logger.warning("%s Tool call attempt %d/%d failed for tool %s: %s",
+                               AGENT_LOG_PREFIX,
+                               attempt,
+                               max_retries,
+                               tool.name,
+                               str(e))
+
+                # Exponential backoff: 2^attempt seconds
+                sleep_time = 2**attempt
+                logger.debug("%s Retrying tool call for %s in %d seconds...", AGENT_LOG_PREFIX, tool.name, sleep_time)
+                await asyncio.sleep(sleep_time)
+
+        # All retries exhausted, return error message
+        error_content = f"Tool call failed after all retry attempts. Last error: {str(last_exception)}"
+        logger.error("%s %s", AGENT_LOG_PREFIX, error_content, exc_info=True)
+        return ToolMessage(name=tool.name, tool_call_id=tool.name, content=error_content, status="error")
+
+    def _log_tool_response(self, tool_name: str, tool_input: Any, tool_response: str) -> None:
+        """
+        Log tool response with consistent formatting and length limits.
+
+        Parameters
+        ----------
+        tool_name : str
+            The name of the tool that was called
+        tool_input : Any
+            The input that was passed to the tool
+        tool_response : str
+            The response from the tool
+        """
+        if self.detailed_logs:
+            # Truncate tool response if too long
+            display_response = tool_response[:self.log_response_max_chars] + "...(rest of response truncated)" if len(
+                tool_response) > self.log_response_max_chars else tool_response
+
+            # Format the tool input for display
+            tool_input_str = str(tool_input)
+
+            tool_response_log_message = TOOL_CALL_LOG_MESSAGE % (tool_name, tool_input_str, display_response)
+            logger.info(tool_response_log_message)
+
+    def _parse_json(self, json_string: str) -> dict[str, Any]:
+        """
+        Safely parse JSON with graceful error handling.
+        If JSON parsing fails, returns an empty dict or error info.
+
+        Parameters
+        ----------
+        json_string : str
+            The JSON string to parse
+
+        Returns
+        -------
+        Dict[str, Any]
+            The parsed JSON or error information
+        """
+        try:
+            return json.loads(json_string)
+        except json.JSONDecodeError as e:
+            logger.warning("%s JSON parsing failed, returning the original string: %s", AGENT_LOG_PREFIX, str(e))
+            return {"error": f"JSON parsing failed: {str(e)}", "original_string": json_string}
+        except Exception as e:
+            logger.warning("%s Unexpected error during JSON parsing: %s", AGENT_LOG_PREFIX, str(e))
+            return {"error": f"Unexpected parsing error: {str(e)}", "original_string": json_string}
+
+    def _get_chat_history(self, messages: list[BaseMessage]) -> str:
+        """
+        Get the chat history excluding the last message.
+
+        Parameters
+        ----------
+        messages : list[BaseMessage]
+            The messages to get the chat history from
+
+        Returns
+        -------
+        str
+            The chat history excluding the last message
+        """
+        return "\n".join([f"{message.type}: {message.content}" for message in messages[:-1]])
+
+    @abstractmethod
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
+        pass

nat/agent/dual_node.py

@@ -0,0 +1,72 @@
+# 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 abc import abstractmethod
+
+from langchain_core.callbacks import AsyncCallbackHandler
+from langchain_core.language_models import BaseChatModel
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from pydantic import BaseModel
+
+from .base import AgentDecision
+from .base import BaseAgent
+
+log = logging.getLogger(__name__)
+
+
+class DualNodeAgent(BaseAgent):
+
+    def __init__(self,
+                 llm: BaseChatModel,
+                 tools: list[BaseTool],
+                 callbacks: list[AsyncCallbackHandler] | None = None,
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000):
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
+
+    @abstractmethod
+    async def agent_node(self, state: BaseModel) -> BaseModel:
+        pass
+
+    @abstractmethod
+    async def tool_node(self, state: BaseModel) -> BaseModel:
+        pass
+
+    @abstractmethod
+    async def conditional_edge(self, state: BaseModel) -> str:
+        pass
+
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
+        log.debug("Building and compiling the Agent Graph")
+
+        graph = StateGraph(state_schema)
+        graph.add_node("agent", self.agent_node)
+        graph.add_node("tool", self.tool_node)
+        graph.add_edge("tool", "agent")
+
+        conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
+        graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
+
+        graph.set_entry_point("agent")
+        self.graph = graph.compile()
+
+        return self.graph

nat/agent/prompt_optimizer/prompt.py

@@ -0,0 +1,68 @@
+# 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 W291
+
+mutator_prompt = """
+
+## CORE DIRECTIVES
+- **Preserve the original objective and task.** Do not change what the prompt is meant to accomplish.  
+- **Keep the intent intact.** The improved prompt must solve the same problem as the original.  
+- **Do not invent new goals.** Only improve clarity, structure, constraints, and usability.  
+- **Do not drop critical instructions.** Everything essential from the original prompt must remain.  
+- **Return only the mutated prompt text.** No rationale, no diffs, no explanations.  
+- **Be Creative within bounds.** You may rephrase, reorganize, and enhance, but not alter meaning.
+- **DO NOT use curly braces in your prompt** for anything other than existing variables in the prompt as the string
+will be treated as an f-string.
+- **Examples are a good idea** if the original prompt lacks them. They help clarify expected output.
+
+---
+
+## IMPROVEMENT HINTS
+When modifying, apply these principles:
+1. **Clarity & Precision** – remove vague language, strengthen directives.  
+2. **Structure & Flow** – order sections as: *Objective → Constraints → Tools → Steps → Output Schema → Examples*.  
+3. **Schema Adherence** – enforce a single canonical output schema (JSON/XML) with `schema_version`.  
+4. **Tool Governance** – clarify when/how tools are used, their inputs/outputs, and fallback behavior.  
+5. **Error Handling** – specify behavior if tools fail or inputs are insufficient.  
+6. **Budget Awareness** – minimize verbosity, respect token/latency limits.  
+7. **Safety** – include refusals for unsafe requests, enforce compliance with rules.  
+8. **Consistency** – avoid format drift; always maintain the same schema.  
+9. **Integrity** – confirm the task, objective, and intent are preserved.  
+
+---
+
+## MUTATION OPERATORS
+You may:
+- **Tighten** (remove fluff, redundancies)  
+- **Reorder** (improve logical flow)  
+- **Constrain** (add explicit rules/limits)  
+- **Harden** (improve error handling/fallbacks)  
+- **Defuse** (replace ambiguous verbs with measurable actions)  
+- **Format-lock** (wrap outputs in JSON/XML fenced blocks)
+- **Example-ify** (add examples if missing or weak)  
+
+---
+
+## INPUT
+Here is the prompt to mutate:
+{original_prompt}
+
+## OBJECTIVE
+The prompt must acheive the following objective:
+{objective}
+
+The modified prompt is: \n
+
+"""

nat/agent/prompt_optimizer/register.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-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.
+
+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.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.profiler.parameter_optimization.prompt_optimizer import PromptOptimizerInputSchema
+
+
+class PromptOptimizerConfig(FunctionBaseConfig, name="prompt_init"):
+
+    optimizer_llm: LLMRef = Field(description="LLM to use for prompt optimization")
+    optimizer_prompt: str = Field(
+        description="Prompt template for the optimizer",
+        default=(
+            "You are an expert at optimizing prompts for LLMs. "
+            "Your task is to take a given prompt and suggest an optimized version of it. "
+            "Note that the prompt might be a template with variables and curly braces. Remember to always keep the "
+            "variables and curly braces in the prompt the same. Only modify the instructions in the prompt that are"
+            "not variables. The system is meant to achieve the following objective\n"
+            "{system_objective}\n Of which, the prompt is one part. The details of the prompt and context as below.\n"))
+    system_objective: str = Field(description="Objective of the workflow")
+
+
+@register_function(config_type=PromptOptimizerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def prompt_optimizer_function(config: PromptOptimizerConfig, builder: Builder):
+    """
+    Function to optimize prompts for LLMs.
+    """
+
+    try:
+        from langchain_core.prompts import PromptTemplate
+
+        from .prompt import mutator_prompt
+    except ImportError as exc:
+        raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
+
+    llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    template = PromptTemplate(template=config.optimizer_prompt,
+                              input_variables=["system_objective"],
+                              validate_template=True)
+
+    base_prompt: str = (await template.ainvoke(input={"system_objective": config.system_objective})).to_string()
+    prompt_extension_template = PromptTemplate(template=mutator_prompt,
+                                               input_variables=["original_prompt", "objective"],
+                                               validate_template=True)
+
+    async def _inner(input_message: PromptOptimizerInputSchema) -> str:
+        """
+        Optimize the prompt using the provided LLM.
+        """
+
+        original_prompt = input_message.original_prompt
+        prompt_objective = input_message.objective
+
+        prompt_extension = (await prompt_extension_template.ainvoke(input={
+            "original_prompt": original_prompt,
+            "objective": prompt_objective,
+        })).to_string()
+
+        prompt = f"{base_prompt}\n\n{prompt_extension}"
+
+        optimized_prompt = await llm.ainvoke(prompt)
+        return optimized_prompt.content
+
+    yield FunctionInfo.from_fn(
+        fn=_inner,
+        description="Optimize prompts for LLMs using a feedback LLM.",
+    )
+
+
+class PromptRecombinerConfig(FunctionBaseConfig, name="prompt_recombiner"):
+
+    optimizer_llm: LLMRef = Field(description="LLM to use for prompt recombination")
+    optimizer_prompt: str = Field(
+        description="Prompt template for the recombiner",
+        default=("You are an expert at combining prompt instructions for LLMs. "
+                 "Your task is to merge two prompts for the same objective into a single, stronger prompt. "
+                 "Do not introduce new variables or modify existing placeholders."),
+    )
+    system_objective: str = Field(description="Objective of the workflow")
+
+
+@register_function(config_type=PromptRecombinerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def prompt_recombiner_function(config: PromptRecombinerConfig, builder: Builder):
+    """
+    Function to recombine two parent prompts into a child prompt using the optimizer LLM.
+    Uses the same base template and objective instructions.
+    """
+
+    try:
+        from langchain_core.prompts import PromptTemplate
+    except ImportError as exc:
+        raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
+
+    llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    template = PromptTemplate(template=config.optimizer_prompt,
+                              input_variables=["system_objective"],
+                              validate_template=True)
+
+    base_prompt: str = (await template.ainvoke(input={"system_objective": config.system_objective})).to_string()
+
+    class RecombineSchema(PromptOptimizerInputSchema):
+        parent_b: str | None = None
+
+    async def _inner(input_message: RecombineSchema) -> str:
+        parent_a = input_message.original_prompt
+        parent_b = input_message.parent_b or ""
+        prompt_objective = input_message.objective
+
+        prompt = (
+            f"{base_prompt}\n\n"
+            "We are performing genetic recombination between two prompts that satisfy the same objective.\n"
+            f"Objective: {prompt_objective}\n\n"
+            f"Parent A:\n{parent_a}\n\n"
+            f"Parent B:\n{parent_b}\n\n"
+            "Combine the strongest instructions and phrasing from both parents to produce a single, coherent child "
+            "prompt.\n"
+            "Maintain variables and placeholders unchanged.\n"
+            "Return only the child prompt text, with no additional commentary.")
+
+        child_prompt = await llm.ainvoke(prompt)
+        return child_prompt.content
+
+    yield FunctionInfo.from_fn(
+        fn=_inner,
+        description="Recombine two prompts into a stronger child prompt.",
+    )