aiqtoolkit 1.1.0a20250503__py3-none-any.whl

aiq/tool/mcp/mcp_client.py

@@ -0,0 +1,220 @@
+# 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.
+
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from enum import Enum
+from typing import Any
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+from mcp.types import TextContent
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import create_model
+
+logger = logging.getLogger(__name__)
+
+
+def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
+    """
+    Create a pydantic model from the input schema of the MCP tool
+    """
+    _type_map = {
+        "string": str,
+        "number": float,
+        "integer": int,
+        "boolean": bool,
+        "array": list,
+        "null": None,
+        "object": dict,
+    }
+
+    properties = mcp_input_schema.get("properties", {})
+    schema_dict = {}
+
+    def _generate_valid_classname(class_name: str):
+        return class_name.replace('_', ' ').replace('-', ' ').title().replace(' ', '')
+
+    def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
+        json_type = field_properties.get("type", "string")
+        enum_vals = field_properties.get("enum")
+
+        if enum_vals:
+            enum_name = f"{field_name.capitalize()}Enum"
+            field_type = Enum(enum_name, {item: item for item in enum_vals})
+
+        elif json_type == "object" and "properties" in field_properties:
+            field_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
+        elif json_type == "array" and "items" in field_properties:
+            item_properties = field_properties.get("items", {})
+            if item_properties.get("type") == "object":
+                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
+            else:
+                item_type = _type_map.get(json_type, Any)
+            field_type = list[item_type]
+        else:
+            field_type = _type_map.get(json_type, Any)
+
+        default_value = field_properties.get("default", ...)
+        nullable = field_properties.get("nullable", False)
+        description = field_properties.get("description", "")
+
+        field_type = field_type | None if nullable else field_type
+
+        return field_type, Field(default=default_value, description=description)
+
+    for field_name, field_props in properties.items():
+        schema_dict[field_name] = _generate_field(field_name=field_name, field_properties=field_props)
+    return create_model(f"{_generate_valid_classname(name)}InputSchema", **schema_dict)
+
+
+class MCPSSEClient:
+    """
+    Client for creating a session and connecting to an MCP server using SSE
+
+    Args:
+      url (str): The url of the MCP server
+    """
+
+    def __init__(self, url: str):
+        self.url = url
+
+    @asynccontextmanager
+    async def connect_to_sse_server(self):
+        """
+        Establish a session with an MCP SSE server within an aync context
+        """
+        async with sse_client(url=self.url) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+
+class MCPBuilder(MCPSSEClient):
+    """
+    Builder class used to connect to an MCP Server and generate ToolClients
+
+    Args:
+        url (str): The url of the MCP server
+    """
+
+    def __init__(self, url):
+        super().__init__(url)
+        self._tools = None
+
+    async def get_tools(self):
+        """
+        Retrieve a dictionary of all tools served by the MCP server.
+        """
+        async with self.connect_to_sse_server() as session:
+            response = await session.list_tools()
+
+        return {
+            tool.name: MCPToolClient(self.url, tool.name, tool.description, tool_input_schema=tool.inputSchema)
+            for tool in response.tools
+        }
+
+    async def get_tool(self, tool_name: str) -> MCPToolClient:
+        """
+        Get an MCP Tool by name.
+
+        Args:
+            tool_name (str): Name of the tool to load.
+
+        Returns:
+            MCPToolClient for the configured tool.
+
+        Raise:
+            ValueError if no tool is available with that name.
+        """
+        if not self._tools:
+            self._tools = await self.get_tools()
+
+        tool = self._tools.get(tool_name)
+        if not tool:
+            raise ValueError(f"Tool {tool_name} not available at {self.url}")
+        return tool
+
+    async def call_tool(self, tool_name: str, tool_args: dict | None):
+        async with self.connect_to_sse_server() as session:
+            result = await session.call_tool(tool_name, tool_args)
+            return result
+
+
+class MCPToolClient(MCPSSEClient):
+    """
+    Client wrapper used to call an MCP tool.
+
+    Args:
+        url (str): The url of the MCP server
+        tool_name (str): The name of the tool to wrap
+        tool_description (str): The description of the tool provided by the MCP server.
+        tool_input_schema (dict): The input schema for the tool.
+    """
+
+    def __init__(self, url: str, tool_name: str, tool_description: str | None, tool_input_schema: dict | None = None):
+        super().__init__(url)
+        self._tool_name = tool_name
+        self._tool_description = tool_description
+        self._input_schema = model_from_mcp_schema(self._tool_name, tool_input_schema) if tool_input_schema else None
+
+    @property
+    def name(self):
+        """Returns the name of the tool."""
+        return self._tool_name
+
+    @property
+    def description(self):
+        """
+        Returns the tool's description. If none was provided. Provides a simple description using the tool's name
+        """
+        if not self._tool_description:
+            return f"MCP Tool {self._tool_name}"
+        return self._tool_description
+
+    @property
+    def input_schema(self):
+        """
+        Returns the tool's input_schema.
+        """
+        return self._input_schema
+
+    def set_description(self, description: str):
+        """
+        Manually define the tool's description using the provided string.
+        """
+        self._tool_description = description
+
+    async def acall(self, tool_args: dict) -> str:
+        """
+        Call the MCP tool with the provided arguments.
+
+        Args:
+            tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
+        """
+        async with self.connect_to_sse_server() as session:
+            result = await session.call_tool(self._tool_name, tool_args)
+
+        output = []
+        for res in result.content:
+            if isinstance(res, TextContent):
+                output.append(res.text)
+            else:
+                # Log non-text content for now
+                logger.warning("Got not-text output from %s of type %s", self.name, type(res))
+        return "\n".join(output)

aiq/tool/mcp/mcp_tool.py

@@ -0,0 +1,76 @@
+# 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 pydantic import HttpUrl
+
+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 MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
+    """
+    Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as an AIQ Toolkit
+    function.
+    """
+    # Add your custom configuration parameters here
+    url: HttpUrl = Field(description="The URL of the MCP server")
+    mcp_tool_name: str = Field(description="The name of the tool served by the MCP Server that you want to use")
+    description: str | None = Field(default=None,
+                                    description="""
+        Description for the tool that will override the description provided by the MCP server. Should only be used if
+        the description provided by the server is poor or nonexistent
+        """)
+
+
+@register_function(config_type=MCPToolConfig)
+async def mcp_tool(config: MCPToolConfig, builder: Builder):
+    """
+    Generate an AIQ Toolkit Function that wraps a tool provided by the MCP server.
+    """
+
+    from aiq.tool.mcp.mcp_client import MCPBuilder
+    from aiq.tool.mcp.mcp_client import MCPToolClient
+
+    client = MCPBuilder(url=str(config.url))
+
+    tool: MCPToolClient = await client.get_tool(config.mcp_tool_name)
+    if config.description:
+        tool.set_description(description=config.description)
+
+    logger.info("Configured to use tool: %s from MCP server at %s", tool.name, str(config.url))
+
+    def _convert_from_str(input_str: str) -> tool.input_schema:
+        return tool.input_schema.model_validate_json(input_str)
+
+    async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
+        if tool_input:
+            args = tool_input.model_dump()
+            return await tool.acall(args)
+
+        _ = tool.input_schema.model_validate(kwargs)
+        return await tool.acall(kwargs)
+
+    yield FunctionInfo.create(single_fn=_response_fn,
+                              description=tool.description,
+                              input_schema=tool.input_schema,
+                              converters=[_convert_from_str])

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