nvidia-nat 1.4.0a20251008__py3-none-any.whl → 1.4.0a20251011__py3-none-any.whl

nat/front_ends/mcp/memory_profiler.py

@@ -0,0 +1,320 @@
+# 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/tool_converter.py

@@ -28,6 +28,7 @@ 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__)
 
@@ -38,6 +39,7 @@ def create_function_wrapper(
     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.
 
@@ -47,6 +49,7 @@ def create_function_wrapper(
         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
@@ -172,6 +175,10 @@ def create_function_wrapper(
                 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
@@ -181,6 +188,11 @@ def create_function_wrapper(
             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
@@ -242,7 +254,8 @@ def get_function_description(function: FunctionBase) -> str:
 def register_function_with_mcp(mcp: FastMCP,
                                function_name: str,
                                function: FunctionBase,
-                               workflow: 'Workflow | None' = None) -> None:
+                               workflow: 'Workflow | None' = None,
+                               memory_profiler: 'MemoryProfiler | None' = None) -> None:
     """Register a NAT Function as an MCP tool.
 
     Args:
@@ -250,6 +263,7 @@ def register_function_with_mcp(mcp: FastMCP,
         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)
 
@@ -267,5 +281,10 @@ def register_function_with_mcp(mcp: FastMCP,
     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)
+    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)

nat/observability/register.py

@@ -77,6 +77,14 @@ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Bu
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.StreamHandler(stream=sys.stdout)
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler
 
 
@@ -95,4 +103,12 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler

nat/runtime/runner.py

@@ -196,8 +196,7 @@ class Runner:
 
             return result
         except Exception as e:
-            err_msg = f": {e}" if str(e).strip() else "."
-            logger.error("Error running workflow%s", err_msg)
+            logger.error("Error running workflow: %s", e)
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()

nat/runtime/session.py

@@ -192,7 +192,7 @@ class SessionManager:
                                     user_message_id: str | None,
                                     conversation_id: str | None) -> None:
         """
-        Extracts and sets user metadata for Websocket connections.
+        Extracts and sets user metadata for WebSocket connections.
         """
 
         # Extract cookies from WebSocket headers (similar to HTTP request)

nat/tool/memory_tools/add_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class AddToolConfig(FunctionBaseConfig, name="add_memory"):
     """Function to add memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to add memory about a user's interactions to a system "
+    description: str = Field(default=("Tool to add a memory about a user's interactions to a system "
                                       "for retrieval later."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -46,7 +46,7 @@ async def add_memory_tool(config: AddToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(item: MemoryItem) -> str:
         """

nat/tool/memory_tools/delete_memory_tool.py

@@ -30,10 +30,9 @@ logger = logging.getLogger(__name__)
 class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"):
     """Function to delete memory from a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
-                                      "interactions to help answer questions in a personalized way."),
+    description: str = Field(default="Tool to delete a memory from a hosted memory platform.",
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -47,7 +46,7 @@ async def delete_memory_tool(config: DeleteToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(user_id: str) -> str:
         """

nat/tool/memory_tools/get_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class GetToolConfig(FunctionBaseConfig, name="get_memory"):
     """Function to get memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
+    description: str = Field(default=("Tool to retrieve a memory about a user's "
                                       "interactions to help answer questions in a personalized way."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -49,7 +49,7 @@ async def get_memory_tool(config: GetToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(search_input: SearchMemoryInput) -> str:
         """

nat/utils/type_converter.py

@@ -93,6 +93,14 @@ class TypeConverter:
         if to_type is None or decomposed.is_instance(data):
             return data
 
+        # 2) If data is a union type, try to convert to each type in the union
+        if decomposed.is_union:
+            for union_type in decomposed.args:
+                result = self._convert(data, union_type)
+                if result is not None:
+                    return result
+            return None
+
         root = decomposed.root
 
         # 2) Attempt direct in *this* converter