nvidia-nat 1.3.0a20250829__py3-none-any.whl → 1.3.0a20250831__py3-none-any.whl

nat/front_ends/mcp/tool_converter.py

@@ -17,13 +17,17 @@ import json
 import logging
 from inspect import Parameter
 from inspect import Signature
+from typing import TYPE_CHECKING
 
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel
 
+from nat.builder.context import ContextState
 from nat.builder.function import Function
 from nat.builder.function_base import FunctionBase
-from nat.builder.workflow import Workflow
+
+if TYPE_CHECKING:
+    from nat.builder.workflow import Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -33,14 +37,16 @@ def create_function_wrapper(
     function: FunctionBase,
     schema: type[BaseModel],
     is_workflow: bool = False,
+    workflow: 'Workflow | None' = None,
 ):
     """Create a wrapper function that exposes the actual parameters of a NAT Function as an MCP tool.
 
     Args:
-        function_name: The name of the function/tool
-        function: The NAT Function object
-        schema: The input schema of the function
-        is_workflow: Whether the function is a Workflow
+        function_name (str): The name of the function/tool
+        function (FunctionBase): The NAT Function object
+        schema (type[BaseModel]): The input schema of the function
+        is_workflow (bool): Whether the function is a Workflow
+        workflow (Workflow | None): The parent workflow for observability context
 
     Returns:
         A wrapper function suitable for registration with MCP
@@ -101,6 +107,19 @@ def create_function_wrapper(
                 await ctx.report_progress(0, 100)
 
             try:
+                # Helper function to wrap function calls with observability
+                async def call_with_observability(func_call):
+                    # Use workflow's observability context (workflow should always be available)
+                    if not workflow:
+                        logger.error("Missing workflow context for function %s - observability will not be available",
+                                     function_name)
+                        raise RuntimeError("Workflow context is required for observability")
+
+                    logger.debug("Starting observability context for function %s", function_name)
+                    context_state = ContextState.get()
+                    async with workflow.exporter_manager.start(context_state=context_state):
+                        return await func_call()
+
                 # Special handling for ChatRequest
                 if is_chat_request:
                     from nat.data_models.api_server import ChatRequest
@@ -118,7 +137,7 @@ def create_function_wrapper(
                             result = await runner.result(to_type=str)
                     else:
                         # Regular functions use ainvoke
-                        result = await function.ainvoke(chat_request, to_type=str)
+                        result = await call_with_observability(lambda: function.ainvoke(chat_request, to_type=str))
                 else:
                     # Regular handling
                     # Handle complex input schema - if we extracted fields from a nested schema,
@@ -129,7 +148,7 @@ def create_function_wrapper(
                         field_type = schema.model_fields[field_name].annotation
 
                         # If it's a pydantic model, we need to create an instance
-                        if hasattr(field_type, "model_validate"):
+                        if field_type and hasattr(field_type, "model_validate"):
                             # Create the nested object
                             nested_obj = field_type.model_validate(kwargs)
                             # Call with the nested object
@@ -147,7 +166,7 @@ def create_function_wrapper(
                             result = await runner.result(to_type=str)
                     else:
                         # Regular function call
-                        result = await function.acall_invoke(**kwargs)
+                        result = await call_with_observability(lambda: function.acall_invoke(**kwargs))
 
                 # Report completion
                 if ctx:
@@ -170,7 +189,7 @@ def create_function_wrapper(
     wrapper = create_wrapper()
 
     # Set the signature on the wrapper function (WITHOUT ctx)
-    wrapper.__signature__ = sig
+    wrapper.__signature__ = sig  # type: ignore
     wrapper.__name__ = function_name
 
     # Return the wrapper with proper signature
@@ -183,8 +202,8 @@ def get_function_description(function: FunctionBase) -> str:
 
     The description is determined using the following precedence:
        1. If the function is a Workflow and has a 'description' attribute, use it.
-       2. If the Workflow's config has a 'topic', use it.
-       3. If the Workflow's config has a 'description', use it.
+       2. If the Workflow's config has a 'description', use it.
+       3. If the Workflow's config has a 'topic', use it.
        4. If the function is a regular Function, use its 'description' attribute.
 
     Args:
@@ -195,6 +214,9 @@ def get_function_description(function: FunctionBase) -> str:
     """
     function_description = ""
 
+    # Import here to avoid circular imports
+    from nat.builder.workflow import Workflow
+
     if isinstance(function, Workflow):
         config = function.config
 
@@ -214,13 +236,17 @@ def get_function_description(function: FunctionBase) -> str:
     return function_description
 
 
-def register_function_with_mcp(mcp: FastMCP, function_name: str, function: FunctionBase) -> None:
+def register_function_with_mcp(mcp: FastMCP,
+                               function_name: str,
+                               function: FunctionBase,
+                               workflow: 'Workflow | None' = None) -> None:
     """Register a NAT Function as an MCP tool.
 
     Args:
         mcp: The FastMCP instance
         function_name: The name to register the function under
         function: The NAT Function to register
+        workflow: The parent workflow for observability context (if available)
     """
     logger.info("Registering function %s with MCP", function_name)
 
@@ -229,6 +255,7 @@ def register_function_with_mcp(mcp: FastMCP, function_name: str, function: Funct
     logger.info("Function %s has input schema: %s", function_name, input_schema)
 
     # Check if we're dealing with a Workflow
+    from nat.builder.workflow import Workflow
     is_workflow = isinstance(function, Workflow)
     if is_workflow:
         logger.info("Function %s is a Workflow", function_name)
@@ -237,5 +264,5 @@ def register_function_with_mcp(mcp: FastMCP, function_name: str, function: Funct
     function_description = get_function_description(function)
 
     # Create and register the wrapper function with MCP
-    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow)
+    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow, workflow)
     mcp.tool(name=function_name, description=function_description)(wrapper_func)

nat/observability/register.py

@@ -72,8 +72,10 @@ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Bu
     """
     Build and return a StreamHandler for console-based logging.
     """
+    import sys
+
     level = getattr(logging, config.level.upper(), logging.INFO)
-    handler = logging.StreamHandler()
+    handler = logging.StreamHandler(stream=sys.stdout)
     handler.setLevel(level)
     yield handler
 

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