nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -1,362 +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 abc import ABC
-from abc import abstractmethod
-from collections.abc import Mapping
-from typing import Any
-
-from mcp.server.fastmcp import FastMCP
-from starlette.exceptions import HTTPException
-from starlette.requests import Request
-
-from nat.builder.function import Function
-from nat.builder.function_base import FunctionBase
-from nat.builder.workflow import Workflow
-from nat.builder.workflow_builder import WorkflowBuilder
-from nat.data_models.config import Config
-from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
-from nat.front_ends.mcp.memory_profiler import MemoryProfiler
-
-logger = logging.getLogger(__name__)
-
-
-class MCPFrontEndPluginWorkerBase(ABC):
-    """Base class for MCP front end plugin workers.
-
-    This abstract base class provides shared utilities and defines the contract
-    for MCP worker implementations. Most users should inherit from
-    MCPFrontEndPluginWorker instead of this class directly.
-    """
-
-    def __init__(self, config: Config):
-        """Initialize the MCP worker with configuration.
-
-        Args:
-            config: The full NAT configuration
-        """
-        self.full_config = config
-        self.front_end_config: MCPFrontEndConfig = config.general.front_end
-
-        # Initialize memory profiler if enabled
-        self.memory_profiler = MemoryProfiler(enabled=self.front_end_config.enable_memory_profiling,
-                                              log_interval=self.front_end_config.memory_profile_interval,
-                                              top_n=self.front_end_config.memory_profile_top_n,
-                                              log_level=self.front_end_config.memory_profile_log_level)
-
-    def _setup_health_endpoint(self, mcp: FastMCP):
-        """Set up the HTTP health endpoint that exercises MCP ping handler."""
-
-        @mcp.custom_route("/health", methods=["GET"])
-        async def health_check(_request: Request):
-            """HTTP health check using server's internal ping handler"""
-            from starlette.responses import JSONResponse
-
-            try:
-                from mcp.types import PingRequest
-
-                # Create a ping request
-                ping_request = PingRequest(method="ping")
-
-                # Call the ping handler directly (same one that responds to MCP pings)
-                await mcp._mcp_server.request_handlers[PingRequest](ping_request)
-
-                return JSONResponse({
-                    "status": "healthy",
-                    "error": None,
-                    "server_name": mcp.name,
-                })
-
-            except Exception as e:
-                return JSONResponse({
-                    "status": "unhealthy",
-                    "error": str(e),
-                    "server_name": mcp.name,
-                },
-                                    status_code=503)
-
-    @abstractmethod
-    async def create_mcp_server(self) -> FastMCP:
-        """Create and configure the MCP server instance.
-
-        This is the main extension point. Plugins can return FastMCP or any subclass
-        to customize server behavior (for example, add authentication, custom transports).
-
-        Returns:
-            FastMCP instance or a subclass with custom behavior
-        """
-        ...
-
-    @abstractmethod
-    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
-        """Add routes to the MCP server.
-
-        Plugins must implement this method. Most plugins can call
-        _default_add_routes() for standard behavior and then add
-        custom enhancements.
-
-        Args:
-            mcp: The FastMCP server instance
-            builder: The workflow builder instance
-        """
-        ...
-
-    async def _default_add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
-        """Default route registration logic - reusable by subclasses.
-
-        This is a protected helper method that plugins can call to get
-        standard route registration behavior. Plugins typically call this
-        from their add_routes() implementation and then add custom features.
-
-        This method:
-        - Sets up the health endpoint
-        - Builds the workflow and extracts all functions
-        - Filters functions based on tool_names config
-        - Registers each function as an MCP tool
-        - Sets up debug endpoints for tool introspection
-
-        Args:
-            mcp: The FastMCP server instance
-            builder: The workflow builder instance
-        """
-        from nat.front_ends.mcp.tool_converter import register_function_with_mcp
-
-        # Set up the health endpoint
-        self._setup_health_endpoint(mcp)
-
-        # Build the workflow and register all functions with MCP
-        workflow = await builder.build()
-
-        # Get all functions from the workflow
-        functions = await self._get_all_functions(workflow)
-
-        # Filter functions based on tool_names if provided
-        if self.front_end_config.tool_names:
-            logger.info("Filtering functions based on tool_names: %s", self.front_end_config.tool_names)
-            filtered_functions: dict[str, Function] = {}
-            for function_name, function in functions.items():
-                if function_name in self.front_end_config.tool_names:
-                    # Treat current tool_names as function names, so check if the function name is in the list
-                    filtered_functions[function_name] = function
-                elif any(function_name.startswith(f"{group_name}.") for group_name in self.front_end_config.tool_names):
-                    # Treat tool_names as function group names, so check if the function name starts with the group name
-                    filtered_functions[function_name] = function
-                else:
-                    logger.debug("Skipping function %s as it's not in tool_names", function_name)
-            functions = filtered_functions
-
-        # Register each function with MCP, passing workflow context for observability
-        for function_name, function in functions.items():
-            register_function_with_mcp(mcp, function_name, function, workflow, self.memory_profiler)
-
-        # Add a simple fallback function if no functions were found
-        if not functions:
-            raise RuntimeError("No functions found in workflow. Please check your configuration.")
-
-        # After registration, expose debug endpoints for tool/schema inspection
-        self._setup_debug_endpoints(mcp, functions)
-
-    async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
-        """Get all functions from the workflow.
-
-        Args:
-            workflow: The NAT workflow.
-
-        Returns:
-            Dict mapping function names to Function objects.
-        """
-        functions: dict[str, Function] = {}
-
-        # Extract all functions from the workflow
-        functions.update(workflow.functions)
-        for function_group in workflow.function_groups.values():
-            functions.update(await function_group.get_accessible_functions())
-
-        if workflow.config.workflow.workflow_alias:
-            functions[workflow.config.workflow.workflow_alias] = workflow
-        else:
-            functions[workflow.config.workflow.type] = workflow
-
-        return functions
-
-    def _setup_debug_endpoints(self, mcp: FastMCP, functions: Mapping[str, FunctionBase]) -> None:
-        """Set up HTTP debug endpoints for introspecting tools and schemas.
-
-        Exposes:
-          - GET /debug/tools/list: List tools. Optional query param `name` (one or more, repeatable or comma separated)
-            selects a subset and returns details for those tools.
-          - GET /debug/memory/stats: Get current memory profiling statistics (read-only)
-        """
-
-        @mcp.custom_route("/debug/tools/list", methods=["GET"])
-        async def list_tools(request: Request):
-            """HTTP list tools endpoint."""
-
-            from starlette.responses import JSONResponse
-
-            from nat.front_ends.mcp.tool_converter import get_function_description
-
-            # Query params
-            # Support repeated names and comma-separated lists
-            names_param_list = set(request.query_params.getlist("name"))
-            names: list[str] = []
-            for raw in names_param_list:
-                # if p.strip() is empty, it won't be included in the list!
-                parts = [p.strip() for p in raw.split(",") if p.strip()]
-                names.extend(parts)
-            detail_raw = request.query_params.get("detail")
-
-            def _parse_detail_param(detail_param: str | None, has_names: bool) -> bool:
-                if detail_param is None:
-                    if has_names:
-                        return True
-                    return False
-                v = detail_param.strip().lower()
-                if v in ("0", "false", "no", "off"):
-                    return False
-                if v in ("1", "true", "yes", "on"):
-                    return True
-                # For invalid values, default based on whether names are present
-                return has_names
-
-            # Helper function to build the input schema info
-            def _build_schema_info(fn: FunctionBase) -> dict[str, Any] | None:
-                schema = getattr(fn, "input_schema", None)
-                if schema is None:
-                    return None
-
-                # check if schema is a ChatRequest
-                schema_name = getattr(schema, "__name__", "")
-                schema_qualname = getattr(schema, "__qualname__", "")
-                if "ChatRequest" in schema_name or "ChatRequest" in schema_qualname:
-                    # Simplified interface used by MCP wrapper for ChatRequest
-                    return {
-                        "type": "object",
-                        "properties": {
-                            "query": {
-                                "type": "string", "description": "User query string"
-                            }
-                        },
-                        "required": ["query"],
-                        "title": "ChatRequestQuery",
-                    }
-
-                # Pydantic models provide model_json_schema
-                if schema is not None and hasattr(schema, "model_json_schema"):
-                    return schema.model_json_schema()
-
-                return None
-
-            def _build_final_json(functions_to_include: Mapping[str, FunctionBase],
-                                  include_schemas: bool = False) -> dict[str, Any]:
-                tools = []
-                for name, fn in functions_to_include.items():
-                    list_entry: dict[str, Any] = {
-                        "name": name, "description": get_function_description(fn), "is_workflow": hasattr(fn, "run")
-                    }
-                    if include_schemas:
-                        list_entry["schema"] = _build_schema_info(fn)
-                    tools.append(list_entry)
-
-                return {
-                    "count": len(tools),
-                    "tools": tools,
-                    "server_name": mcp.name,
-                }
-
-            if names:
-                # Return selected tools
-                try:
-                    functions_to_include = {n: functions[n] for n in names}
-                except KeyError as e:
-                    raise HTTPException(status_code=404, detail=f"Tool \"{e.args[0]}\" not found.") from e
-            else:
-                functions_to_include = functions
-
-            # Default for listing all: detail defaults to False unless explicitly set true
-            return JSONResponse(
-                _build_final_json(functions_to_include, _parse_detail_param(detail_raw, has_names=bool(names))))
-
-        # Memory profiling endpoint (read-only)
-        @mcp.custom_route("/debug/memory/stats", methods=["GET"])
-        async def get_memory_stats(_request: Request):
-            """Get current memory profiling statistics."""
-            from starlette.responses import JSONResponse
-
-            stats = self.memory_profiler.get_stats()
-            return JSONResponse(stats)
-
-
-class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
-    """Default MCP server worker implementation.
-
-    Inherit from this class to create custom MCP workers that extend or modify
-    server behavior. Override create_mcp_server() to use a different server type,
-    and override add_routes() to add custom functionality.
-
-    Example:
-        class CustomWorker(MCPFrontEndPluginWorker):
-            async def create_mcp_server(self):
-                # Return custom MCP server instance
-                return MyCustomFastMCP(...)
-
-            async def add_routes(self, mcp, builder):
-                # Get default routes
-                await super().add_routes(mcp, builder)
-                # Add custom features
-                self._add_my_custom_features(mcp)
-    """
-
-    async def create_mcp_server(self) -> FastMCP:
-        """Create default MCP server with optional authentication.
-
-        Returns:
-            FastMCP instance configured with settings from NAT config
-        """
-        # Handle auth if configured
-        auth_settings = None
-        token_verifier = None
-
-        if self.front_end_config.server_auth:
-            from mcp.server.auth.settings import AuthSettings
-            from pydantic import AnyHttpUrl
-
-            server_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}"
-            auth_settings = AuthSettings(issuer_url=AnyHttpUrl(self.front_end_config.server_auth.issuer_url),
-                                         required_scopes=self.front_end_config.server_auth.scopes,
-                                         resource_server_url=AnyHttpUrl(server_url))
-
-            # Create token verifier
-            from nat.front_ends.mcp.introspection_token_verifier import IntrospectionTokenVerifier
-
-            token_verifier = IntrospectionTokenVerifier(self.front_end_config.server_auth)
-
-        return FastMCP(name=self.front_end_config.name,
-                       host=self.front_end_config.host,
-                       port=self.front_end_config.port,
-                       debug=self.front_end_config.debug,
-                       auth=auth_settings,
-                       token_verifier=token_verifier)
-
-    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
-        """Add default routes to the MCP server.
-
-        Args:
-            mcp: The FastMCP server instance
-            builder: The workflow builder instance
-        """
-        # Use the default implementation from base class to add the tools to the MCP server
-        await self._default_add_routes(mcp, builder)

nat/front_ends/mcp/memory_profiler.py

@@ -1,320 +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.
-"""Memory profiling utilities for MCP frontend."""
-
-import gc
-import logging
-import tracemalloc
-from typing import Any
-
-logger = logging.getLogger(__name__)
-
-
-class MemoryProfiler:
-    """Memory profiler for tracking memory usage and potential leaks."""
-
-    def __init__(self, enabled: bool = False, log_interval: int = 50, top_n: int = 10, log_level: str = "DEBUG"):
-        """Initialize the memory profiler.
-
-        Args:
-            enabled: Whether memory profiling is enabled
-            log_interval: Log stats every N requests
-            top_n: Number of top allocations to log
-            log_level: Log level for memory profiling output (e.g., "DEBUG", "INFO")
-        """
-        self.enabled = enabled
-        # normalize interval to avoid modulo-by-zero
-        self.log_interval = max(1, int(log_interval))
-        self.top_n = top_n
-        self.log_level = getattr(logging, log_level.upper(), logging.DEBUG)
-        self.request_count = 0
-        self.baseline_snapshot = None
-
-        # Track whether this instance started tracemalloc (to avoid resetting external tracing)
-        self._we_started_tracemalloc = False
-
-        if self.enabled:
-            logger.info("Memory profiling ENABLED (interval=%d, top_n=%d, log_level=%s)",
-                        self.log_interval,
-                        top_n,
-                        log_level)
-            try:
-                if not tracemalloc.is_tracing():
-                    tracemalloc.start()
-                    self._we_started_tracemalloc = True
-                # Take baseline snapshot
-                gc.collect()
-                self.baseline_snapshot = tracemalloc.take_snapshot()
-            except RuntimeError as e:
-                logger.warning("tracemalloc unavailable or not tracing: %s", e)
-        else:
-            logger.info("Memory profiling DISABLED")
-
-    def _log(self, message: str, *args: Any) -> None:
-        """Log a message at the configured log level.
-
-        Args:
-            message: Log message format string
-            args: Arguments for the format string
-        """
-        logger.log(self.log_level, message, *args)
-
-    def on_request_complete(self) -> None:
-        """Called after each request completes."""
-        if not self.enabled:
-            return
-        self.request_count += 1
-        if self.request_count % self.log_interval == 0:
-            self.log_memory_stats()
-
-    def _ensure_tracing(self) -> bool:
-        """Ensure tracemalloc is running if we started it originally.
-
-        Returns:
-            True if tracemalloc is active, False otherwise
-        """
-        if tracemalloc.is_tracing():
-            return True
-
-        # Only restart if we started it originally (respect external control)
-        if not self._we_started_tracemalloc:
-            return False
-
-        # Attempt to restart
-        try:
-            logger.warning("tracemalloc was stopped externally; restarting (we started it originally)")
-            tracemalloc.start()
-
-            # Reset baseline since old tracking data is lost
-            gc.collect()
-            self.baseline_snapshot = tracemalloc.take_snapshot()
-            logger.info("Baseline snapshot reset after tracemalloc restart")
-
-            return True
-        except RuntimeError as e:
-            logger.error("Failed to restart tracemalloc: %s", e)
-            return False
-
-    def _safe_traced_memory(self) -> tuple[float, float] | None:
-        """Return (current, peak usage in MB) if tracemalloc is active, else None."""
-        if not self._ensure_tracing():
-            return None
-
-        try:
-            current, peak = tracemalloc.get_traced_memory()
-            megabyte = (1 << 20)
-            return (current / megabyte, peak / megabyte)
-        except RuntimeError:
-            return None
-
-    def _safe_snapshot(self) -> tracemalloc.Snapshot | None:
-        """Return a tracemalloc Snapshot if available, else None."""
-        if not self._ensure_tracing():
-            return None
-
-        try:
-            return tracemalloc.take_snapshot()
-        except RuntimeError:
-            return None
-
-    def log_memory_stats(self) -> dict[str, Any]:
-        """Log current memory statistics and return them."""
-        if not self.enabled:
-            return {}
-
-        # Force garbage collection first
-        gc.collect()
-
-        # Get current memory usage
-        mem = self._safe_traced_memory()
-        if mem is None:
-            logger.info("tracemalloc is not active; cannot collect memory stats.")
-            # still return structural fields
-            stats = {
-                "request_count": self.request_count,
-                "current_memory_mb": None,
-                "peak_memory_mb": None,
-                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
-                "outstanding_steps": self._safe_outstanding_step_count(),
-                "active_exporters": self._safe_exporter_count(),
-                "isolated_exporters": self._safe_isolated_exporter_count(),
-                "subject_instances": self._count_instances_of_type("Subject"),
-            }
-            return stats
-
-        current_mb, peak_mb = mem
-
-        # Take snapshot and compare to baseline
-        snapshot = self._safe_snapshot()
-
-        # Track BaseExporter instances (observability layer)
-        exporter_count = self._safe_exporter_count()
-        isolated_exporter_count = self._safe_isolated_exporter_count()
-
-        # Track Subject instances (event streams)
-        subject_count = self._count_instances_of_type("Subject")
-
-        stats = {
-            "request_count": self.request_count,
-            "current_memory_mb": round(current_mb, 2),
-            "peak_memory_mb": round(peak_mb, 2),
-            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
-            "outstanding_steps": self._safe_outstanding_step_count(),
-            "active_exporters": exporter_count,
-            "isolated_exporters": isolated_exporter_count,
-            "subject_instances": subject_count,
-        }
-
-        self._log("=" * 80)
-        self._log("MEMORY PROFILE AFTER %d REQUESTS:", self.request_count)
-        self._log("  Current Memory: %.2f MB", current_mb)
-        self._log("  Peak Memory: %.2f MB", peak_mb)
-        self._log("")
-        self._log("NAT COMPONENT INSTANCES:")
-        self._log("  IntermediateStepManagers: %d active (%d outstanding steps)",
-                  stats["active_intermediate_managers"],
-                  stats["outstanding_steps"])
-        self._log("  BaseExporters: %d active (%d isolated)", stats["active_exporters"], stats["isolated_exporters"])
-        self._log("  Subject (event streams): %d instances", stats["subject_instances"])
-
-        # Show top allocations
-        if snapshot is None:
-            self._log("tracemalloc snapshot unavailable.")
-        else:
-            if self.baseline_snapshot:
-                self._log("TOP %d MEMORY GROWTH SINCE BASELINE:", self.top_n)
-                top_stats = snapshot.compare_to(self.baseline_snapshot, 'lineno')
-            else:
-                self._log("TOP %d MEMORY ALLOCATIONS:", self.top_n)
-                top_stats = snapshot.statistics('lineno')
-
-            for i, stat in enumerate(top_stats[:self.top_n], 1):
-                self._log("  #%d: %s", i, stat)
-
-        self._log("=" * 80)
-
-        return stats
-
-    def _count_instances_of_type(self, type_name: str) -> int:
-        """Count instances of a specific type in memory."""
-        count = 0
-        for obj in gc.get_objects():
-            try:
-                if type(obj).__name__ == type_name:
-                    count += 1
-            except Exception:
-                pass
-        return count
-
-    def _safe_exporter_count(self) -> int:
-        try:
-            from nat.observability.exporter.base_exporter import BaseExporter
-            return BaseExporter.get_active_instance_count()
-        except Exception as e:
-            logger.debug("Could not get BaseExporter stats: %s", e)
-            return 0
-
-    def _safe_isolated_exporter_count(self) -> int:
-        try:
-            from nat.observability.exporter.base_exporter import BaseExporter
-            return BaseExporter.get_isolated_instance_count()
-        except Exception:
-            return 0
-
-    def _safe_intermediate_step_manager_count(self) -> int:
-        try:
-            from nat.builder.intermediate_step_manager import IntermediateStepManager
-            # len() is atomic in CPython, but catch RuntimeError just in case
-            try:
-                return IntermediateStepManager.get_active_instance_count()
-            except RuntimeError:
-                # Set was modified during len() - very rare
-                logger.debug("Set changed during count, returning 0")
-                return 0
-        except Exception as e:
-            logger.debug("Could not get IntermediateStepManager stats: %s", e)
-            return 0
-
-    def _safe_outstanding_step_count(self) -> int:
-        """Get total outstanding steps across all active IntermediateStepManager instances."""
-        try:
-            from nat.builder.intermediate_step_manager import IntermediateStepManager
-
-            # Make a snapshot to avoid "Set changed size during iteration" if GC runs
-            try:
-                instances_snapshot = list(IntermediateStepManager._active_instances)
-            except RuntimeError:
-                # Set changed during list() call - rare but possible
-                logger.debug("Set changed during snapshot, returning 0 for outstanding steps")
-                return 0
-
-            total_outstanding = 0
-            # Iterate through snapshot safely
-            for ref in instances_snapshot:
-                try:
-                    manager = ref()
-                    if manager is not None:
-                        total_outstanding += manager.get_outstanding_step_count()
-                except (ReferenceError, AttributeError):
-                    # Manager was GC'd or in invalid state - skip it
-                    continue
-            return total_outstanding
-        except Exception as e:
-            logger.debug("Could not get outstanding step count: %s", e)
-            return 0
-
-    def get_stats(self) -> dict[str, Any]:
-        """Get current memory statistics without logging."""
-        if not self.enabled:
-            return {"enabled": False}
-
-        mem = self._safe_traced_memory()
-        if mem is None:
-            return {
-                "enabled": True,
-                "request_count": self.request_count,
-                "current_memory_mb": None,
-                "peak_memory_mb": None,
-                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
-                "outstanding_steps": self._safe_outstanding_step_count(),
-                "active_exporters": self._safe_exporter_count(),
-                "isolated_exporters": self._safe_isolated_exporter_count(),
-                "subject_instances": self._count_instances_of_type("Subject"),
-            }
-
-        current_mb, peak_mb = mem
-        return {
-            "enabled": True,
-            "request_count": self.request_count,
-            "current_memory_mb": round(current_mb, 2),
-            "peak_memory_mb": round(peak_mb, 2),
-            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
-            "outstanding_steps": self._safe_outstanding_step_count(),
-            "active_exporters": self._safe_exporter_count(),
-            "isolated_exporters": self._safe_isolated_exporter_count(),
-            "subject_instances": self._count_instances_of_type("Subject"),
-        }
-
-    def reset_baseline(self) -> None:
-        """Reset the baseline snapshot to current state."""
-        if not self.enabled:
-            return
-        gc.collect()
-        snap = self._safe_snapshot()
-        if snap is None:
-            logger.info("Cannot reset baseline: tracemalloc is not active.")
-            return
-        self.baseline_snapshot = snap
-        logger.info("Memory profiling baseline reset at request %d", self.request_count)