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

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)

nat/front_ends/mcp/register.py

@@ -1,27 +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 collections.abc import AsyncIterator
-
-from nat.cli.register_workflow import register_front_end
-from nat.data_models.config import Config
-from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
-
-
-@register_front_end(config_type=MCPFrontEndConfig)
-async def register_mcp_front_end(config: MCPFrontEndConfig, full_config: Config) -> AsyncIterator:
-    from nat.front_ends.mcp.mcp_front_end_plugin import MCPFrontEndPlugin
-
-    yield MCPFrontEndPlugin(full_config=full_config)

nat/front_ends/mcp/tool_converter.py

@@ -1,321 +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 json
-import logging
-from inspect import Parameter
-from inspect import Signature
-from typing import TYPE_CHECKING
-from typing import Any
-
-from mcp.server.fastmcp import FastMCP
-from pydantic import BaseModel
-from pydantic.fields import FieldInfo
-from pydantic_core import PydanticUndefined
-
-from nat.builder.context import ContextState
-from nat.builder.function import Function
-from nat.builder.function_base import FunctionBase
-
-if TYPE_CHECKING:
-    from nat.builder.workflow import Workflow
-    from nat.front_ends.mcp.memory_profiler import MemoryProfiler
-
-logger = logging.getLogger(__name__)
-
-# Sentinel: marks "optional; let Pydantic supply default/factory"
-_USE_PYDANTIC_DEFAULT = object()
-
-
-def is_field_optional(field: FieldInfo) -> tuple[bool, Any]:
-    """Determine if a Pydantic field is optional and extract its default value for MCP signatures.
-
-    For MCP tool signatures, we need to distinguish:
-    - Required fields: marked with Parameter.empty
-    - Optional with concrete default: use that default
-    - Optional with factory: use sentinel so Pydantic can apply the factory later
-
-    Args:
-        field: The Pydantic FieldInfo to check
-
-    Returns:
-        A tuple of (is_optional, default_value):
-        - (False, Parameter.empty) for required fields
-        - (True, actual_default) for optional fields with explicit defaults
-        - (True, _USE_PYDANTIC_DEFAULT) for optional fields with default_factory
-    """
-    if field.is_required():
-        return False, Parameter.empty
-
-    # Field is optional - has either default or factory
-    if field.default is not PydanticUndefined:
-        return True, field.default
-
-    # Factory case: mark optional in signature but don't fabricate a value
-    if field.default_factory is not None:
-        return True, _USE_PYDANTIC_DEFAULT
-
-    # Rare corner case: non-required yet no default surfaced
-    return True, _USE_PYDANTIC_DEFAULT
-
-
-def create_function_wrapper(
-    function_name: str,
-    function: FunctionBase,
-    schema: type[BaseModel],
-    is_workflow: bool = False,
-    workflow: 'Workflow | None' = None,
-    memory_profiler: 'MemoryProfiler | None' = None,
-):
-    """Create a wrapper function that exposes the actual parameters of a NAT Function as an MCP tool.
-
-    Args:
-        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
-        memory_profiler: Optional memory profiler to track requests
-
-    Returns:
-        A wrapper function suitable for registration with MCP
-    """
-    # Check if we're dealing with ChatRequest - special case
-    is_chat_request = False
-
-    # Check if the schema name is ChatRequest
-    if schema.__name__ == "ChatRequest" or (hasattr(schema, "__qualname__") and "ChatRequest" in schema.__qualname__):
-        is_chat_request = True
-        logger.info("Function %s uses ChatRequest - creating simplified interface", function_name)
-
-        # For ChatRequest, we'll create a simple wrapper with just a query parameter
-        parameters = [Parameter(
-            name="query",
-            kind=Parameter.KEYWORD_ONLY,
-            default=Parameter.empty,
-            annotation=str,
-        )]
-    else:
-        # Regular case - extract parameter information from the input schema
-        # Extract parameter information from the input schema
-        param_fields = schema.model_fields
-
-        parameters = []
-        for name, field in param_fields.items():
-            # Get the field type and convert to appropriate Python type
-            field_type = field.annotation
-
-            # Check if field is optional and get its default value
-            _is_optional, param_default = is_field_optional(field)
-
-            # Add the parameter to our list
-            parameters.append(
-                Parameter(
-                    name=name,
-                    kind=Parameter.KEYWORD_ONLY,
-                    default=param_default,
-                    annotation=field_type,
-                ))
-
-    # Create the function signature WITHOUT the ctx parameter
-    # We'll handle this in the wrapper function internally
-    sig = Signature(parameters=parameters, return_annotation=str)
-
-    # Define the actual wrapper function that accepts ctx but doesn't expose it
-    def create_wrapper():
-
-        async def wrapper_with_ctx(**kwargs):
-            """Internal wrapper that will be called by MCP."""
-            # MCP will add a ctx parameter, extract it
-            ctx = kwargs.get("ctx")
-
-            # Remove ctx if present
-            if "ctx" in kwargs:
-                del kwargs["ctx"]
-
-            # Process the function call
-            if ctx:
-                ctx.info("Calling function %s with args: %s", function_name, json.dumps(kwargs, default=str))
-                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
-
-                    # Create a chat request from the query string
-                    query = kwargs.get("query", "")
-                    chat_request = ChatRequest.from_string(query)
-
-                    # Special handling for Workflow objects
-                    if is_workflow:
-                        # Workflows have a run method that is an async context manager
-                        # that returns a Runner
-                        async with function.run(chat_request) as runner:
-                            # Get the result from the runner
-                            result = await runner.result(to_type=str)
-                    else:
-                        # Regular functions use ainvoke
-                        result = await call_with_observability(lambda: function.ainvoke(chat_request, to_type=str))
-                else:
-                    # Regular handling
-                    # Strip sentinel values so Pydantic can apply defaults/factories
-                    cleaned_kwargs = {k: v for k, v in kwargs.items() if v is not _USE_PYDANTIC_DEFAULT}
-
-                    # Always validate with the declared schema
-                    # This handles defaults, factories, nested models, validators, etc.
-                    model_input = schema.model_validate(cleaned_kwargs)
-
-                    # Call the NAT function with the parameters - special handling for Workflow
-                    if is_workflow:
-                        # Workflows expect the model instance directly
-                        async with function.run(model_input) as runner:
-                            # Get the result from the runner
-                            result = await runner.result(to_type=str)
-                    else:
-                        # Regular function call - unpack the validated model
-                        result = await call_with_observability(lambda: function.acall_invoke(**model_input.model_dump())
-                                                               )
-
-                # Report completion
-                if ctx:
-                    await ctx.report_progress(100, 100)
-
-                # Track request completion for memory profiling
-                if memory_profiler:
-                    memory_profiler.on_request_complete()
-
-                # Handle different result types for proper formatting
-                if isinstance(result, str):
-                    return result
-                if isinstance(result, dict | list):
-                    return json.dumps(result, default=str)
-                return str(result)
-            except Exception as e:
-                if ctx:
-                    ctx.error("Error calling function %s: %s", function_name, str(e))
-
-                # Track request completion even on error
-                if memory_profiler:
-                    memory_profiler.on_request_complete()
-
-                raise
-
-        return wrapper_with_ctx
-
-    # Create the wrapper function
-    wrapper = create_wrapper()
-
-    # Set the signature on the wrapper function (WITHOUT ctx)
-    wrapper.__signature__ = sig  # type: ignore
-    wrapper.__name__ = function_name
-
-    # Return the wrapper with proper signature
-    return wrapper
-
-
-def get_function_description(function: FunctionBase) -> str:
-    """
-    Retrieve a human-readable description for a NAT function or workflow.
-
-    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 '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:
-        function: The NAT FunctionBase instance (Function or Workflow).
-
-    Returns:
-        The best available description string for the function.
-    """
-    function_description = ""
-
-    # Import here to avoid circular imports
-    from nat.builder.workflow import Workflow
-
-    if isinstance(function, Workflow):
-        config = function.config
-
-        # Workflow doesn't have a description, but probably should
-        if hasattr(function, "description") and function.description:
-            function_description = function.description
-        # Try to get description from config
-        elif hasattr(config, "description") and config.description:
-            function_description = config.description
-        # Try to get anything that might be a description
-        elif hasattr(config, "topic") and config.topic:
-            function_description = config.topic
-        # Try to get description from the workflow config
-        elif hasattr(config, "workflow") and hasattr(config.workflow, "description") and config.workflow.description:
-            function_description = config.workflow.description
-
-    elif isinstance(function, Function):
-        function_description = function.description
-
-    return function_description
-
-
-def register_function_with_mcp(mcp: FastMCP,
-                               function_name: str,
-                               function: FunctionBase,
-                               workflow: 'Workflow | None' = None,
-                               memory_profiler: 'MemoryProfiler | 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)
-        memory_profiler: Optional memory profiler to track requests
-    """
-    logger.info("Registering function %s with MCP", function_name)
-
-    # Get the input schema from the function
-    input_schema = function.input_schema
-    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)
-
-    # Get function description
-    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,
-                                           workflow,
-                                           memory_profiler)
-    mcp.tool(name=function_name, description=function_description)(wrapper_func)