nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/tool/mcp/exceptions.py

@@ -1,142 +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 enum import Enum
-
-
-class MCPErrorCategory(str, Enum):
-    """Categories of MCP errors for structured handling."""
-    CONNECTION = "connection"
-    TIMEOUT = "timeout"
-    SSL = "ssl"
-    AUTHENTICATION = "authentication"
-    TOOL_NOT_FOUND = "tool_not_found"
-    PROTOCOL = "protocol"
-    UNKNOWN = "unknown"
-
-
-class MCPError(Exception):
-    """Base exception for MCP-related errors."""
-
-    def __init__(self,
-                 message: str,
-                 url: str,
-                 category: MCPErrorCategory = MCPErrorCategory.UNKNOWN,
-                 suggestions: list[str] | None = None,
-                 original_exception: Exception | None = None):
-        super().__init__(message)
-        self.url = url
-        self.category = category
-        self.suggestions = suggestions or []
-        self.original_exception = original_exception
-
-
-class MCPConnectionError(MCPError):
-    """Exception for MCP connection failures."""
-
-    def __init__(self, url: str, original_exception: Exception | None = None):
-        super().__init__(f"Unable to connect to MCP server at {url}",
-                         url=url,
-                         category=MCPErrorCategory.CONNECTION,
-                         suggestions=[
-                             "Please ensure the MCP server is running and accessible",
-                             "Check if the URL and port are correct"
-                         ],
-                         original_exception=original_exception)
-
-
-class MCPTimeoutError(MCPError):
-    """Exception for MCP timeout errors."""
-
-    def __init__(self, url: str, original_exception: Exception | None = None):
-        super().__init__(f"Connection timed out to MCP server at {url}",
-                         url=url,
-                         category=MCPErrorCategory.TIMEOUT,
-                         suggestions=[
-                             "The server may be overloaded or network is slow",
-                             "Try again in a moment or check network connectivity"
-                         ],
-                         original_exception=original_exception)
-
-
-class MCPSSLError(MCPError):
-    """Exception for MCP SSL/TLS errors."""
-
-    def __init__(self, url: str, original_exception: Exception | None = None):
-        super().__init__(f"SSL/TLS error connecting to {url}",
-                         url=url,
-                         category=MCPErrorCategory.SSL,
-                         suggestions=[
-                             "Check if the server requires HTTPS or has valid certificates",
-                             "Try using HTTP instead of HTTPS if appropriate"
-                         ],
-                         original_exception=original_exception)
-
-
-class MCPRequestError(MCPError):
-    """Exception for MCP request errors."""
-
-    def __init__(self, url: str, original_exception: Exception | None = None):
-        message = f"Request failed to MCP server at {url}"
-        if original_exception:
-            message += f": {original_exception}"
-
-        super().__init__(message,
-                         url=url,
-                         category=MCPErrorCategory.PROTOCOL,
-                         suggestions=["Check the server URL format and network settings"],
-                         original_exception=original_exception)
-
-
-class MCPToolNotFoundError(MCPError):
-    """Exception for when a specific MCP tool is not found."""
-
-    def __init__(self, tool_name: str, url: str, original_exception: Exception | None = None):
-        super().__init__(f"Tool '{tool_name}' not available at {url}",
-                         url=url,
-                         category=MCPErrorCategory.TOOL_NOT_FOUND,
-                         suggestions=[
-                             "Use 'nat info mcp --detail' to see available tools",
-                             "Check that the tool name is spelled correctly"
-                         ],
-                         original_exception=original_exception)
-
-
-class MCPAuthenticationError(MCPError):
-    """Exception for MCP authentication failures."""
-
-    def __init__(self, url: str, original_exception: Exception | None = None):
-        super().__init__(f"Authentication failed when connecting to MCP server at {url}",
-                         url=url,
-                         category=MCPErrorCategory.AUTHENTICATION,
-                         suggestions=[
-                             "Check if the server requires authentication credentials",
-                             "Verify that your credentials are correct and not expired"
-                         ],
-                         original_exception=original_exception)
-
-
-class MCPProtocolError(MCPError):
-    """Exception for MCP protocol-related errors."""
-
-    def __init__(self, url: str, message: str = "Protocol error", original_exception: Exception | None = None):
-        super().__init__(f"{message} (MCP server at {url})",
-                         url=url,
-                         category=MCPErrorCategory.PROTOCOL,
-                         suggestions=[
-                             "Check that the MCP server is running and accessible at this URL",
-                             "Verify the server supports the expected MCP protocol version"
-                         ],
-                         original_exception=original_exception)

nat/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 nat.tool.mcp.exceptions import MCPToolNotFoundError
-from nat.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)

nat/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 nat.builder.builder import Builder
-from nat.builder.function_info import FunctionInfo
-from nat.cli.register_workflow import register_function
-from nat.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 a NAT Function that wraps a tool provided by the MCP server.
-    """
-
-    from nat.tool.mcp.mcp_client import MCPBuilder
-    from nat.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])