nvidia-nat 1.3.0a20250828__py3-none-any.whl → 1.3.0a20250830__py3-none-any.whl

nat/tool/mcp/{mcp_client.py → mcp_client_base.py}

@@ -16,12 +16,18 @@
 from __future__ import annotations
 
 import logging
+from abc import ABC
+from abc import abstractmethod
+from contextlib import AsyncExitStack
 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.client.stdio import StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import TextContent
 from pydantic import BaseModel
 from pydantic import Field
@@ -29,6 +35,7 @@ from pydantic import create_model
 
 from nat.tool.mcp.exceptions import MCPToolNotFoundError
 from nat.utils.exception_handlers.mcp import mcp_exception_handler
+from nat.utils.type_utils import override
 
 logger = logging.getLogger(__name__)
 
@@ -107,56 +114,78 @@ def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
     return create_model(f"{_generate_valid_classname(name)}InputSchema", **schema_dict)
 
 
-class MCPSSEClient:
+class MCPBaseClient(ABC):
     """
-    Client for creating a session and connecting to an MCP server using SSE
+    Base client for creating a session and connecting to an MCP server
 
     Args:
-      url (str): The url of the MCP server
+        transport (str): The type of client to use ('sse', 'stdio', or 'streamable-http')
     """
 
-    def __init__(self, url: str):
-        self.url = url
+    def __init__(self, transport: str = 'streamable-http'):
+        self._tools = None
+        self._transport = transport.lower()
+        if self._transport not in ['sse', 'stdio', 'streamable-http']:
+            raise ValueError("transport must be either 'sse', 'stdio' or 'streamable-http'")
 
-    @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
+        self._exit_stack: AsyncExitStack | None = None
 
+        self._session: ClientSession | None = None
 
-class MCPBuilder(MCPSSEClient):
-    """
-    Builder class used to connect to an MCP Server and generate ToolClients
+    @property
+    def transport(self) -> str:
+        return self._transport
 
-    Args:
-        url (str): The url of the MCP server
-    """
+    async def __aenter__(self):
+        if self._exit_stack:
+            raise RuntimeError("MCPBaseClient already initialized. Use async with to initialize.")
 
-    def __init__(self, url):
-        super().__init__(url)
-        self._tools = None
+        self._exit_stack = AsyncExitStack()
+
+        self._session = await self._exit_stack.enter_async_context(self.connect_to_server())
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+
+        if not self._exit_stack:
+            raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
+
+        await self._exit_stack.aclose()
+        self._session = None
+        self._exit_stack = None
+
+    @property
+    def server_name(self):
+        """
+        Provide server name for logging
+        """
+        return self._transport
+
+    @abstractmethod
+    @asynccontextmanager
+    async def connect_to_server(self):
+        """
+        Establish a session with an MCP server within an async context
+        """
+        pass
 
-    @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
+        if not self._session:
+            raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
 
-        Raises:
-            MCPError: If connection or tool retrieval fails
-        """
-        async with self.connect_to_sse_server() as session:
-            response = await session.list_tools()
+        response = await self._session.list_tools()
 
         return {
-            tool.name: MCPToolClient(self.url, tool.name, tool.description, tool_input_schema=tool.inputSchema)
+            tool.name:
+                MCPToolClient(session=self._session,
+                              tool_name=tool.name,
+                              tool_description=tool.description,
+                              tool_input_schema=tool.inputSchema)
             for tool in response.tools
         }
 
@@ -172,9 +201,11 @@ class MCPBuilder(MCPSSEClient):
             MCPToolClient for the configured tool.
 
         Raises:
-            MCPToolNotFoundError: If no tool is available with that name
-            MCPError: If connection fails
+            MCPToolNotFoundError: If no tool is available with that name.
         """
+        if not self._exit_stack:
+            raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
+
         if not self._tools:
             self._tools = await self.get_tools()
 
@@ -185,27 +216,143 @@ class MCPBuilder(MCPSSEClient):
 
     @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
+        if not self._session:
+            raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
+
+        result = await self._session.call_tool(tool_name, tool_args)
+        return result
+
+
+class MCPSSEClient(MCPBaseClient):
+    """
+    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):
+        super().__init__("sse")
+        self._url = url
+
+    @property
+    def url(self) -> str:
+        return self._url
+
+    @property
+    def server_name(self):
+        return f"sse:{self._url}"
+
+    @asynccontextmanager
+    @override
+    async def connect_to_server(self):
+        """
+        Establish a session with an MCP SSE server within an async context
+        """
+        async with sse_client(url=self._url) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+
+class MCPStdioClient(MCPBaseClient):
+    """
+    Client for creating a session and connecting to an MCP server using stdio
+
+    Args:
+      command (str): The command to run
+      args (list[str] | None): Additional arguments for the command
+      env (dict[str, str] | None): Environment variables to set for the process
+    """
+
+    def __init__(self, command: str, args: list[str] | None = None, env: dict[str, str] | None = None):
+        super().__init__("stdio")
+        self._command = command
+        self._args = args
+        self._env = env
+
+    @property
+    def command(self) -> str:
+        return self._command
+
+    @property
+    def server_name(self):
+        return f"stdio:{self._command}"
+
+    @property
+    def args(self) -> list[str] | None:
+        return self._args
+
+    @property
+    def env(self) -> dict[str, str] | None:
+        return self._env
+
+    @asynccontextmanager
+    @override
+    async def connect_to_server(self):
+        """
+        Establish a session with an MCP server via stdio within an async context
+        """
+
+        server_params = StdioServerParameters(command=self._command, args=self._args or [], env=self._env)
+        async with stdio_client(server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+
+class MCPStreamableHTTPClient(MCPBaseClient):
+    """
+    Client for creating a session and connecting to an MCP server using streamable-http
+
+    Args:
+      url (str): The url of the MCP server
+    """
+
+    def __init__(self, url: str):
+        super().__init__("streamable-http")
+
+        self._url = url
+
+    @property
+    def url(self) -> str:
+        return self._url
 
+    @property
+    def server_name(self):
+        return f"streamable-http:{self._url}"
 
-class MCPToolClient(MCPSSEClient):
+    @asynccontextmanager
+    async def connect_to_server(self):
+        """
+        Establish a session with an MCP server via streamable-http within an async context
+        """
+        async with streamablehttp_client(url=self._url) as (read, write, get_session_id):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+
+class MCPToolClient:
     """
     Client wrapper used to call an MCP tool.
 
     Args:
-        url (str): The url of the MCP server
+        connect_fn (callable): Function that returns an async context manager for connecting to the 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)
+    def __init__(self,
+                 session: ClientSession,
+                 tool_name: str,
+                 tool_description: str | None,
+                 tool_input_schema: dict | None = None):
+        self._session = session
         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
+        self._input_schema = (model_from_mcp_schema(self._tool_name, tool_input_schema) if tool_input_schema else None)
 
     @property
     def name(self):
@@ -234,7 +381,6 @@ class MCPToolClient(MCPSSEClient):
         """
         self._tool_description = description
 
-    @mcp_exception_handler
     async def acall(self, tool_args: dict) -> str:
         """
         Call the MCP tool with the provided arguments.
@@ -242,14 +388,19 @@ class MCPToolClient(MCPSSEClient):
         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)
+        result = await self._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)
+        result_str = "\n".join(output)
+
+        if result.isError:
+            raise RuntimeError(result_str)
+
+        return result_str

nat/tool/mcp/mcp_client_impl.py

@@ -0,0 +1,229 @@
+# 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 typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+from pydantic import model_validator
+
+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
+from nat.experimental.decorators.experimental_warning_decorator import experimental
+from nat.tool.mcp.mcp_client_base import MCPBaseClient
+
+logger = logging.getLogger(__name__)
+
+# All functions in this file are experimental
+
+
+class ToolOverrideConfig(BaseModel):
+    """
+    Configuration for overriding tool properties when exposing from MCP server.
+    """
+    alias: str | None = Field(default=None, description="Override the tool name (function name in the workflow)")
+    description: str | None = Field(default=None, description="Override the tool description")
+
+
+class MCPServerConfig(BaseModel):
+    """
+    Server connection details for MCP client.
+    Supports stdio, sse, and streamable-http transports.
+    streamable-http is the recommended default for HTTP-based connections.
+    """
+    transport: Literal["stdio", "sse", "streamable-http"] = Field(
+        ..., description="Transport type to connect to the MCP server (stdio, sse, or streamable-http)")
+    url: HttpUrl | None = Field(default=None,
+                                description="URL of the MCP server (for sse or streamable-http transport)")
+    command: str | None = Field(default=None,
+                                description="Command to run for stdio transport (e.g. 'python' or 'docker')")
+    args: list[str] | None = Field(default=None, description="Arguments for the stdio command")
+    env: dict[str, str] | None = Field(default=None, description="Environment variables for the stdio process")
+
+    @model_validator(mode="after")
+    def validate_model(self):
+        """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
+        if self.transport == "stdio":
+            if self.url is not None:
+                raise ValueError("url should not be set when using stdio transport")
+            if not self.command:
+                raise ValueError("command is required when using stdio transport")
+        elif self.transport in ("sse", "streamable-http"):
+            if self.command is not None or self.args is not None or self.env is not None:
+                raise ValueError("command, args, and env should not be set when using sse or streamable-http transport")
+            if not self.url:
+                raise ValueError("url is required when using sse or streamable-http transport")
+        return self
+
+
+class MCPClientConfig(FunctionBaseConfig, name="mcp_client"):
+    """
+    Configuration for connecting to an MCP server as a client and exposing selected tools.
+    """
+    server: MCPServerConfig = Field(..., description="Server connection details (transport, url/command, etc.)")
+    tool_filter: dict[str, ToolOverrideConfig] | list[str] | None = Field(
+        default=None,
+        description="""Filter or map tools to expose from the server (list or dict).
+        Can be:
+        - A list of tool names to expose: ['tool1', 'tool2']
+        - A dict mapping tool names to override configs:
+          {'tool1': {'alias': 'new_name', 'description': 'New desc'}}
+          {'tool2': {'description': 'Override description only'}}  # alias defaults to 'tool2'
+        """)
+
+
+class MCPSingleToolConfig(FunctionBaseConfig, name="mcp_single_tool"):
+    """
+    Configuration for wrapping a single tool from an MCP server as a NeMo Agent toolkit function.
+    """
+    client: MCPBaseClient = Field(..., description="MCP client to use for the tool")
+    tool_name: str = Field(..., description="Name of the tool to use")
+    tool_description: str | None = Field(default=None, description="Description of the tool")
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+def _get_server_name_safe(client: MCPBaseClient) -> str:
+
+    # Avoid leaking env secrets from stdio client in logs.
+    if client.transport == "stdio":
+        safe_server = f"stdio: {client.command}"
+    else:
+        safe_server = f"{client.transport}: {client.url}"
+
+    return safe_server
+
+
+@register_function(config_type=MCPSingleToolConfig)
+async def mcp_single_tool(config: MCPSingleToolConfig, builder: Builder):
+    """
+    Wrap a single tool from an MCP server as a NeMo Agent toolkit function.
+    """
+    tool = await config.client.get_tool(config.tool_name)
+    if config.tool_description:
+        tool.set_description(description=config.tool_description)
+    input_schema = tool.input_schema
+
+    logger.info("Configured to use tool: %s from MCP server at %s", tool.name, _get_server_name_safe(config.client))
+
+    def _convert_from_str(input_str: str) -> BaseModel:
+        return input_schema.model_validate_json(input_str)
+
+    @experimental(feature_name="mcp_client")
+    async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
+        try:
+            if tool_input:
+                return await tool.acall(tool_input.model_dump())
+            _ = input_schema.model_validate(kwargs)
+            return await tool.acall(kwargs)
+        except Exception as e:
+            return str(e)
+
+    fn = FunctionInfo.create(single_fn=_response_fn,
+                             description=tool.description,
+                             input_schema=input_schema,
+                             converters=[_convert_from_str])
+    yield fn
+
+
+@register_function(MCPClientConfig)
+async def mcp_client_function_handler(config: MCPClientConfig, builder: Builder):
+    """
+    Connect to an MCP server, discover tools, and register them as functions in the workflow.
+
+    Note:
+    - Uses builder's exit stack to manage client lifecycle
+    - Applies tool filters if provided
+    """
+    from nat.tool.mcp.mcp_client_base import MCPSSEClient
+    from nat.tool.mcp.mcp_client_base import MCPStdioClient
+    from nat.tool.mcp.mcp_client_base import MCPStreamableHTTPClient
+
+    # Build the appropriate client
+    client_cls = {
+        "stdio": lambda: MCPStdioClient(config.server.command, config.server.args, config.server.env),
+        "sse": lambda: MCPSSEClient(str(config.server.url)),
+        "streamable-http": lambda: MCPStreamableHTTPClient(str(config.server.url)),
+    }.get(config.server.transport)
+
+    if not client_cls:
+        raise ValueError(f"Unsupported transport: {config.server.transport}")
+
+    client = client_cls()
+    logger.info("Configured to use MCP server at %s", _get_server_name_safe(client))
+
+    # client aenter connects to the server and stores the client in the exit stack
+    # so it's cleaned up when the workflow is done
+    async with client:
+        all_tools = await client.get_tools()
+        tool_configs = _filter_and_configure_tools(all_tools, config.tool_filter)
+
+        for tool_name, tool_cfg in tool_configs.items():
+            await builder.add_function(
+                tool_cfg["function_name"],
+                MCPSingleToolConfig(
+                    client=client,
+                    tool_name=tool_name,
+                    tool_description=tool_cfg["description"],
+                ))
+
+        @experimental(feature_name="mcp_client")
+        async def idle_fn(text: str) -> str:
+            # This function is a placeholder and will be removed when function groups are used
+            return f"MCP client connected: {text}"
+
+        yield FunctionInfo.create(single_fn=idle_fn, description="MCP client")
+
+
+def _filter_and_configure_tools(all_tools: dict, tool_filter) -> dict[str, dict]:
+    """
+    Apply tool filtering and optional aliasing/description overrides.
+
+    Returns:
+        Dict[str, dict] where each value has:
+            - function_name
+            - description
+    """
+    if tool_filter is None:
+        return {name: {"function_name": name, "description": tool.description} for name, tool in all_tools.items()}
+
+    if isinstance(tool_filter, list):
+        return {
+            name: {
+                "function_name": name, "description": all_tools[name].description
+            }
+            for name in tool_filter if name in all_tools
+        }
+
+    if isinstance(tool_filter, dict):
+        result = {}
+        for name, override in tool_filter.items():
+            tool = all_tools.get(name)
+            if not tool:
+                logger.warning("Tool '%s' specified in tool_filter not found in MCP server", name)
+                continue
+
+            if isinstance(override, ToolOverrideConfig):
+                result[name] = {
+                    "function_name": override.alias or name, "description": override.description or tool.description
+                }
+            else:
+                logger.warning("Unsupported override type for '%s': %s", name, type(override))
+                result[name] = {"function_name": name, "description": tool.description}
+        return result