aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/tool/memory_tools/add_memory_tool.py

@@ -0,0 +1,67 @@
+# 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.
+
+import logging
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import MemoryRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.memory.models import MemoryItem
+
+logger = logging.getLogger(__name__)
+
+
+class AddToolConfig(FunctionBaseConfig, name="add_memory"):
+    """Function to add memory to a hosted memory platform."""
+
+    description: str = Field(default=("Tool to add memory about a user's interactions to a system "
+                                      "for retrieval later."),
+                             description="The description of this function's use for tool calling agents.")
+    memory: MemoryRef = Field(default="saas_memory",
+                              description=("Instance name of the memory client instance from the workflow "
+                                           "configuration object."))
+
+
+@register_function(config_type=AddToolConfig)
+async def add_memory_tool(config: AddToolConfig, builder: Builder):
+    """
+    Function to add memory to a hosted memory platform.
+    """
+
+    from langchain_core.tools import ToolException
+
+    # First, retrieve the memory client
+    memory_editor = builder.get_memory_client(config.memory)
+
+    async def _arun(item: MemoryItem) -> str:
+        """
+        Asynchronous execution of addition of memories.
+        """
+
+        try:
+
+            await memory_editor.add_items([item])
+
+            return "Memory added successfully. You can continue. Please respond to the user."
+
+        except Exception as e:
+
+            raise ToolException(f"Error adding memory: {e}") from e
+
+    yield FunctionInfo.from_fn(_arun, description=config.description)

aiq/tool/memory_tools/delete_memory_tool.py

@@ -0,0 +1,67 @@
+# 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.
+
+import logging
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import MemoryRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.memory.models import DeleteMemoryInput
+
+logger = logging.getLogger(__name__)
+
+
+class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"):
+    """Function to delete memory from a hosted memory platform."""
+
+    description: str = Field(default=("Tool to retrieve memory about a user's "
+                                      "interactions to help answer questions in a personalized way."),
+                             description="The description of this function's use for tool calling agents.")
+    memory: MemoryRef = Field(default="saas_memory",
+                              description=("Instance name of the memory client instance from the workflow "
+                                           "configuration object."))
+
+
+@register_function(config_type=DeleteToolConfig)
+async def delete_memory_tool(config: DeleteToolConfig, builder: Builder):
+    """
+    Function to delete memory from a hosted memory platform.
+    """
+
+    from langchain_core.tools import ToolException
+
+    # First, retrieve the memory client
+    memory_editor = builder.get_memory_client(config.memory)
+
+    async def _arun(user_id: str) -> str:
+        """
+        Asynchronous execution of deletion of memories.
+        """
+
+        try:
+
+            await memory_editor.remove_items(user_id=user_id, )
+
+            return "Memories deleted!"
+
+        except Exception as e:
+
+            raise ToolException(f"Error deleting memory: {e}") from e
+
+    yield FunctionInfo.from_fn(_arun, description=config.description, input_schema=DeleteMemoryInput)

aiq/tool/memory_tools/get_memory_tool.py

@@ -0,0 +1,72 @@
+# 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.
+
+import logging
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import MemoryRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.memory.models import SearchMemoryInput
+
+logger = logging.getLogger(__name__)
+
+
+class GetToolConfig(FunctionBaseConfig, name="get_memory"):
+    """Function to get memory to a hosted memory platform."""
+
+    description: str = Field(default=("Tool to retrieve memory about a user's "
+                                      "interactions to help answer questions in a personalized way."),
+                             description="The description of this function's use for tool calling agents.")
+    memory: MemoryRef = Field(default="saas_memory",
+                              description=("Instance name of the memory client instance from the workflow "
+                                           "configuration object."))
+
+
+@register_function(config_type=GetToolConfig)
+async def get_memory_tool(config: GetToolConfig, builder: Builder):
+    """
+    Function to get memory to a hosted memory platform.
+    """
+
+    import json
+
+    from langchain_core.tools import ToolException
+
+    # First, retrieve the memory client
+    memory_editor = builder.get_memory_client(config.memory)
+
+    async def _arun(search_input: SearchMemoryInput) -> str:
+        """
+        Asynchronous execution of collection of memories.
+        """
+        try:
+            memories = await memory_editor.search(
+                query=search_input.query,
+                top_k=search_input.top_k,
+                user_id=search_input.user_id,
+            )
+
+            memory_str = f"Memories as a JSON: \n{json.dumps([mem.model_dump(mode='json') for mem in memories])}"
+            return memory_str
+
+        except Exception as e:
+
+            raise ToolException(f"Error retreiving memory: {e}") from e
+
+    yield FunctionInfo.from_fn(_arun, description=config.description)

aiq/tool/nvidia_rag.py

@@ -0,0 +1,95 @@
+# 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.
+
+import json
+import logging
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class NVIDIARAGToolConfig(FunctionBaseConfig, name="nvidia_rag"):
+    """
+    Tool used to search the NVIDIA Developer database for documents across a variety of NVIDIA asset types.
+    """
+    base_url: str = Field(description="The base url to the RAG service.")
+    timeout: int = Field(default=60, description="The timeout configuration to use when sending requests.")
+    document_separator: str = Field(default="\n\n", description="The delimiter to use between retrieved documents.")
+    document_prompt: str = Field(default=("-------\n\n" + "Title: {document_title}\n"
+                                          "Text: {page_content}\nSource URL: {document_url}"),
+                                 description="The prompt to use to retrieve documents from the RAG service")
+    top_k: int = Field(default=4, description="The number of results to return from the RAG service.")
+    collection_name: str = Field(default="nvidia_api_catalog",
+                                 description=("The name of the collection to use when retrieving documents."))
+
+
+@register_function(config_type=NVIDIARAGToolConfig)
+async def nvidia_rag_tool(config: NVIDIARAGToolConfig, builder: Builder):
+    import httpx
+    from langchain.prompts import PromptTemplate
+    from langchain_core.documents import Document
+    from langchain_core.prompts import aformat_document
+
+    document_prompt = PromptTemplate.from_template(config.document_prompt)
+
+    async with httpx.AsyncClient(headers={
+            "accept": "application/json", "Content-Type": "application/json"
+    },
+                                 timeout=config.timeout) as client:
+
+        async def runnable(query: str) -> str:
+
+            try:
+                url = f"{config.base_url}/search"
+
+                payload = {"query": query, "top_k": config.top_k, "collection_name": config.collection_name}
+
+                logger.debug("Sending request to the RAG endpoint %s.", url)
+                response = await client.post(url, content=json.dumps(payload))
+
+                response.raise_for_status()
+
+                output = response.json()
+
+                docs = [
+                    Document(
+                        page_content=ret["content"],
+                        metadata={
+                            "document_title": ret["filename"],
+                            "document_url": "nemo_framework",
+                            "document_full_text": ret["content"],
+                            "score_rerank": ret["score"]
+                        },
+                        type="Document",
+                    ) for ret in output["chunks"]
+                ]
+
+                parsed_output = config.document_separator.join(
+                    [await aformat_document(doc, document_prompt) for doc in docs])
+                return parsed_output
+            except Exception as e:
+                logger.exception("Error while running the tool", exc_info=True)
+                return f"Error while running the tool: {e}"
+
+        yield FunctionInfo.from_fn(
+            runnable,
+            description=("Search the NVIDIA Developer database for documents across a variety of "
+                         "NVIDIA asset types"))

aiq/tool/register.py

@@ -0,0 +1,36 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+
+# Import any tools which need to be automatically registered here
+from . import datetime_tools
+from . import document_search
+from . import github_tools
+from . import nvidia_rag
+from . import retriever
+from .code_execution import register
+from .github_tools import create_github_commit
+from .github_tools import create_github_issue
+from .github_tools import create_github_pr
+from .github_tools import get_github_file
+from .github_tools import get_github_issue
+from .github_tools import get_github_pr
+from .github_tools import update_github_issue
+from .mcp import mcp_tool
+from .memory_tools import add_memory_tool
+from .memory_tools import delete_memory_tool
+from .memory_tools import get_memory_tool

aiq/tool/retriever.py

@@ -0,0 +1,89 @@
+# 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 BaseModel
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.component_ref import RetrieverRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.retriever.interface import AIQRetriever
+from aiq.retriever.models import RetrieverError
+from aiq.retriever.models import RetrieverOutput
+
+logger = logging.getLogger(__name__)
+
+
+class AIQRetrieverConfig(FunctionBaseConfig, name="aiq_retriever"):
+    """
+    AIQRetriever tool which provides a common interface for different vectorstores. Its
+    configuration uses clients, which are the vectorstore-specific implementaiton of the retriever interface.
+    """
+    retriever: RetrieverRef = Field(description="The retriever instance name from the workflow configuration object.")
+    raise_errors: bool = Field(
+        default=True,
+        description="If true the tool will raise exceptions, otherwise it will log them as warnings and return []",
+    )
+    topic: str | None = Field(default=None, description="Used to provide a more detailed tool description to the agent")
+    description: str | None = Field(default=None, description="If present it will be used as the tool description")
+
+
+def _get_description_from_config(config: AIQRetrieverConfig) -> str:
+    """
+    Generate a description of what the tool will do based on how it is configured.
+    """
+    description = "Retrieve document chunks{topic} which can be used to answer the provided question."
+
+    _topic = f" related to {config.topic}" if config.topic else ""
+
+    return description.format(topic=_topic) if not config.description else config.description
+
+
+@register_function(config_type=AIQRetrieverConfig)
+async def aiq_retriever_tool(config: AIQRetrieverConfig, builder: Builder):
+    """
+    Configure an AgentIQ Retriever Tool which supports different clients such as Milvus and Nemo Retriever.
+
+    Args:
+        config: A config object with required parameters 'client' and 'client_config'
+        builder: A workflow builder object
+    """
+
+    class RetrieverInputSchema(BaseModel):
+        query: str = Field(description="The query to be searched in the configured data store")
+
+    client: AIQRetriever = await builder.get_retriever(config.retriever)
+
+    async def _retrieve(query: str) -> RetrieverOutput:
+        try:
+            retrieved_context = await client.search(query=query)
+            logger.info("Retrieved %s records for query %s.", len(retrieved_context), query)
+            return retrieved_context
+
+        except RetrieverError as e:
+            if config.raise_errors:
+                raise e
+            logger.warning("Retriever threw an error: %s. Returning an empty response.", e)
+            return RetrieverOutput(results=[])
+
+    yield FunctionInfo.from_fn(
+        fn=_retrieve,
+        input_schema=RetrieverInputSchema,
+        description=_get_description_from_config(config),
+    )

aiq/utils/data_models/schema_validator.py

@@ -0,0 +1,58 @@
+# 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.
+
+import yaml
+from pydantic import ValidationError
+
+from ..exception_handlers.schemas import schema_exception_handler
+from ..exception_handlers.schemas import yaml_exception_handler
+
+
+@schema_exception_handler
+def validate_schema(metadata, Schema):  # pylint: disable=invalid-name
+
+    try:
+        return Schema(**metadata)
+    except ValidationError as e:
+
+        raise e
+
+
+@yaml_exception_handler
+def validate_yaml(ctx, param, value):  # pylint: disable=unused-argument
+    """
+    Validate that the file is a valid YAML file
+
+    Parameters
+    ----------
+    ctx: Click context
+    param: Click parameter
+    value: Path to YAML file
+
+    Returns
+    -------
+    str: Path to valid YAML file
+
+    Raises
+    ------
+    ValueError: If file is invalid or unreadable
+    """
+    if value is None:
+        return None
+
+    with open(value, 'r', encoding="utf-8") as f:
+        yaml.safe_load(f)
+
+    return value

aiq/utils/debugging_utils.py

@@ -0,0 +1,43 @@
+# 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.
+
+
+def is_debugger_attached() -> bool:
+    """
+    Check if a debugger is attached to the current process.
+
+    Returns
+    -------
+    bool
+        True if a debugger is attached, False otherwise
+    """
+    import sys
+
+    if "debugpy" in sys.modules:
+
+        import debugpy
+
+        return debugpy.is_client_connected()
+
+    trace_func = sys.gettrace()
+
+    # The presence of a trace function and pydevd means a debugger is attached
+    if (trace_func is not None):
+        trace_module = getattr(trace_func, "__module__", None)
+
+        if (trace_module is not None and trace_module.find("pydevd") != -1):
+            return True
+
+    return False

aiq/utils/exception_handlers/schemas.py

@@ -0,0 +1,114 @@
+# 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.
+
+import logging
+
+import yaml
+from pydantic import ValidationError
+
+logger = logging.getLogger(__name__)
+
+
+def schema_exception_handler(func, **kwargs):  # pylint: disable=unused-argument
+    """
+    A decorator that handles `ValidationError` exceptions for schema validation functions.
+
+    This decorator wraps a function that performs schema validation using Pydantic.
+    If a `ValidationError` is raised, it logs detailed error messages and raises a `ValueError` with the combined error
+    messages.
+
+    Parameters
+    ----------
+    func : callable
+        The function to be decorated. This function is expected to perform schema validation.
+
+    kwargs : dict
+        Additional keyword arguments to be passed to the function.
+
+    Returns
+    -------
+    callable
+        The wrapped function that executes `func` with exception handling.
+
+    Raises
+    ------
+    ValueError
+        If a `ValidationError` is caught, this decorator logs the error details and raises a `ValueError` with the
+        combined error messages.
+
+    Notes
+    -----
+    This decorator is particularly useful for functions that validate configurations or data models,
+    ensuring that any validation errors are logged and communicated clearly.
+
+    Examples
+    --------
+    >>> @schema_exception_handler
+    ... def validate_config(config_data):
+    ...     schema = MySchema(**config_data)
+    ...     return schema
+    ...
+    >>> try:
+    ...     validate_config(invalid_config)
+    ... except ValueError as e:
+    ...     logger.error("Caught error: %s", e)
+    Caught error: Invalid configuration: field1: value is not a valid integer; field2: field required
+    """
+
+    def inner_function(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except ValidationError as e:
+            error_messages = "; ".join([f"{error['loc'][0]}: {error['msg']}" for error in e.errors()])
+            log_error_message = f"Invalid configuration: {error_messages}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+    return inner_function
+
+
+def yaml_exception_handler(func):
+    """
+    A decorator that handles YAML parsing exceptions.
+
+    This decorator wraps a function that performs YAML file operations.
+    If a YAML-related error occurs, it logs the error and raises a ValueError
+    with a clear error message.
+
+    Returns
+    -------
+    callable
+        The wrapped function that executes `func` with YAML exception handling.
+
+    Raises
+    ------
+    ValueError
+        If a YAML error is caught, with details about the parsing failure.
+    """
+
+    def inner_function(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except yaml.YAMLError as e:
+            log_error_message = f"Invalid YAML configuration: {str(e)}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+        except Exception as e:
+            log_error_message = f"Error reading YAML file: {str(e)}"
+            logger.error(log_error_message)
+            raise ValueError(log_error_message) from e
+
+    return inner_function