aiqtoolkit 1.2.0rc3__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/tool/mcp/mcp_client.py

@@ -1,255 +0,0 @@
-# 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
-
-from aiq.tool.mcp.exceptions import MCPToolNotFoundError
-from aiq.utils.exception_handlers.mcp import mcp_exception_handler
-
-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", {})
-    required_fields = set(mcp_input_schema.get("required", []))
-    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=item_properties)
-            else:
-                item_type = _type_map.get(item_properties.get("type", "string"), Any)
-            field_type = list[item_type]
-        elif isinstance(json_type, list):
-            field_type = None
-            for t in json_type:
-                mapped = _type_map.get(t, Any)
-                field_type = mapped if field_type is None else field_type | mapped
-
-            return field_type, Field(
-                default=field_properties.get("default", None if "null" in json_type else ...),
-                description=field_properties.get("description", "")
-            )
-        else:
-            field_type = _type_map.get(json_type, Any)
-
-        # Determine the default value based on whether the field is required
-        if field_name in required_fields:
-            # Field is required - use explicit default if provided, otherwise make it required
-            default_value = field_properties.get("default", ...)
-        else:
-            # Field is optional - use explicit default if provided, otherwise None
-            default_value = field_properties.get("default", None)
-            # Make the type optional if no default was provided
-            if "default" not in field_properties:
-                field_type = field_type | None
-
-        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
-
-    @mcp_exception_handler
-    async def get_tools(self):
-        """
-        Retrieve a dictionary of all tools served by the MCP server.
-
-        Returns:
-            Dict of tool name to MCPToolClient
-
-        Raises:
-            MCPError: If connection or tool retrieval fails
-        """
-        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
-        }
-
-    @mcp_exception_handler
-    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.
-
-        Raises:
-            MCPToolNotFoundError: If no tool is available with that name
-            MCPError: If connection fails
-        """
-        if not self._tools:
-            self._tools = await self.get_tools()
-
-        tool = self._tools.get(tool_name)
-        if not tool:
-            raise MCPToolNotFoundError(tool_name, self.url)
-        return tool
-
-    @mcp_exception_handler
-    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
-
-    @mcp_exception_handler
-    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

@@ -1,96 +0,0 @@
-# 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 a NeMo Agent 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
-        """)
-    return_exception: bool = Field(default=True,
-                                   description="""
-        If true, the tool will return the exception message if the tool call fails.
-        If false, raise the exception.
-        """)
-
-
-@register_function(config_type=MCPToolConfig)
-async def mcp_tool(config: MCPToolConfig, builder: Builder):  # pylint: disable=unused-argument
-    """
-    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:
-        # Run the tool, catching any errors and sending to agent for correction
-        try:
-            if tool_input:
-                args = tool_input.model_dump()
-                return await tool.acall(args)
-
-            _ = tool.input_schema.model_validate(kwargs)
-            filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
-            return await tool.acall(filtered_kwargs)
-        except Exception as e:
-            if config.return_exception:
-                if tool_input:
-                    logger.warning("Error calling tool %s with serialized input: %s",
-                                   tool.name,
-                                   tool_input.model_dump(),
-                                   exc_info=True)
-                else:
-                    logger.warning("Error calling tool %s with input: %s", tool.name, kwargs, exc_info=True)
-                return str(e)
-            # If the tool call fails, raise the exception.
-            raise
-
-    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

@@ -1,79 +0,0 @@
-# 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.
-
-        Args:
-            item (MemoryItem): The memory item to add. Must include:
-                - conversation: List of dicts with "role" and "content" keys
-                - user_id: String identifier for the user
-                - metadata: Dict of metadata (can be empty)
-                - tags: Optional list of tags
-                - memory: Optional memory string
-
-        Note: If conversation is not provided, it will be created from the memory field
-        if available, otherwise an error will be raised.
-        """
-        try:
-            # If conversation is not provided but memory is, create a conversation
-            if not item.conversation and item.memory:
-                item.conversation = [{"role": "user", "content": item.memory}]
-            elif not item.conversation:
-                raise ToolException("Either conversation or memory must be provided")
-
-            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

@@ -1,67 +0,0 @@
-# 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

@@ -1,72 +0,0 @@
-# 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

@@ -1,95 +0,0 @@
-# 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"))