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

nat/registry_handlers/schemas/pull.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/registry_handlers/schemas/remove.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/registry_handlers/schemas/search.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/registry_handlers/schemas/status.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/interface.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/milvus/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/milvus/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -48,6 +48,7 @@ class MilvusRetrieverConfig(RetrieverBaseConfig, name="milvus_retriever"):
     description: str | None = Field(default=None,
                                     description="If present it will be used as the tool description",
                                     alias="collection_description")
+    use_async_client: bool = Field(default=False, description="Use AsyncMilvusClient for async I/O operations. ")
 
 
 @register_retriever_provider(config_type=MilvusRetrieverConfig)
@@ -58,13 +59,20 @@ async def milvus_retriever(retriever_config: MilvusRetrieverConfig, builder: Bui
 
 @register_retriever_client(config_type=MilvusRetrieverConfig, wrapper_type=None)
 async def milvus_retriever_client(config: MilvusRetrieverConfig, builder: Builder):
-    from pymilvus import MilvusClient
-
     from nat.retriever.milvus.retriever import MilvusRetriever
 
     embedder = await builder.get_embedder(embedder_name=config.embedding_model, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
-    milvus_client = MilvusClient(uri=str(config.uri), **config.connection_args)
+    # Create Milvus client based on use_async_client flag
+    if config.use_async_client:
+        from pymilvus import AsyncMilvusClient
+
+        milvus_client = AsyncMilvusClient(uri=str(config.uri), **config.connection_args)
+    else:
+        from pymilvus import MilvusClient
+
+        milvus_client = MilvusClient(uri=str(config.uri), **config.connection_args)
+
     retriever = MilvusRetriever(
         client=milvus_client,
         embedder=embedder,

nat/retriever/milvus/retriever.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,13 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import logging
 from functools import partial
+from typing import TYPE_CHECKING
 
 from langchain_core.embeddings import Embeddings
-from pymilvus import MilvusClient
 from pymilvus.client.abstract import Hit
 
+if TYPE_CHECKING:
+    from pymilvus import AsyncMilvusClient
+    from pymilvus import MilvusClient
+
 from nat.retriever.interface import Retriever
 from nat.retriever.models import Document
 from nat.retriever.models import RetrieverError
@@ -39,20 +44,27 @@ class MilvusRetriever(Retriever):
 
     def __init__(
         self,
-        client: MilvusClient,
+        client: "MilvusClient | AsyncMilvusClient",
         embedder: Embeddings,
         content_field: str = "text",
         use_iterator: bool = False,
     ) -> None:
         """
-        Initialize the Milvus Retriever using a preconfigured MilvusClient
+        Initialize the Milvus Retriever using a preconfigured MilvusClient or AsyncMilvusClient
 
         Args:
-           client (MilvusClient): Preinstantiate pymilvus.MilvusClient object.
         """
-        self._client = client
+        self._client: MilvusClient | AsyncMilvusClient = client
         self._embedder = embedder
 
+        # Detect if client is async by inspecting method capabilities
+        search_method = getattr(client, "search", None)
+        list_collections_method = getattr(client, "list_collections", None)
+        self._is_async = any(
+            inspect.iscoroutinefunction(method) for method in (search_method, list_collections_method)
+            if method is not None)
+        logger.info("Initialized Milvus Retriever with %s client", "async" if self._is_async else "sync")
+
         if use_iterator and "search_iterator" not in dir(self._client):
             raise ValueError("This version of the pymilvus.MilvusClient does not support the search iterator.")
 
@@ -60,7 +72,7 @@ class MilvusRetriever(Retriever):
         self._default_params = None
         self._bound_params = []
         self.content_field = content_field
-        logger.info("Mivlus Retriever using %s for search.", self._search_func.__name__)
+        logger.info("Milvus Retriever using %s for search.", self._search_func.__name__)
 
     def bind(self, **kwargs) -> None:
         """
@@ -81,8 +93,13 @@ class MilvusRetriever(Retriever):
         """
         return [param for param in ["query", "collection_name", "top_k", "filters"] if param not in self._bound_params]
 
-    def _validate_collection(self, collection_name: str) -> bool:
-        return collection_name in self._client.list_collections()
+    async def _validate_collection(self, collection_name: str) -> bool:
+        """Validate that a collection exists."""
+        if self._is_async:
+            collections = await self._client.list_collections()
+        else:
+            collections = self._client.list_collections()
+        return collection_name in collections
 
     async def search(self, query: str, **kwargs):
         return await self._search_func(query=query, **kwargs)
@@ -108,39 +125,64 @@ class MilvusRetriever(Retriever):
                      collection_name,
                      top_k)
 
-        if not self._validate_collection(collection_name):
+        if not await self._validate_collection(collection_name):
             raise CollectionNotFoundError(f"Collection: {collection_name} does not exist")
 
         # If no output fields are specified, return all of them
         if not output_fields:
-            collection_schema = self._client.describe_collection(collection_name)
+            if self._is_async:
+                collection_schema = await self._client.describe_collection(collection_name)
+            else:
+                collection_schema = self._client.describe_collection(collection_name)
             output_fields = [
                 field["name"] for field in collection_schema.get("fields") if field["name"] != vector_field_name
             ]
 
-        search_vector = self._embedder.embed_query(query)
-
-        search_iterator = self._client.search_iterator(
-            collection_name=collection_name,
-            data=[search_vector],
-            batch_size=kwargs.get("batch_size", 1000),
-            filter=filters,
-            limit=top_k,
-            output_fields=output_fields,
-            search_params=search_params if search_params else {"metric_type": "L2"},
-            timeout=timeout,
-            anns_field=vector_field_name,
-            round_decimal=kwargs.get("round_decimal", -1),
-            partition_names=kwargs.get("partition_names", None),
-        )
+        search_vector = await self._embedder.aembed_query(query)
+
+        # Create search iterator
+        if self._is_async:
+            search_iterator = await self._client.search_iterator(
+                collection_name=collection_name,
+                data=[search_vector],
+                batch_size=kwargs.get("batch_size", 1000),
+                filter=filters,
+                limit=top_k,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                round_decimal=kwargs.get("round_decimal", -1),
+                partition_names=kwargs.get("partition_names", None),
+            )
+        else:
+            search_iterator = self._client.search_iterator(
+                collection_name=collection_name,
+                data=[search_vector],
+                batch_size=kwargs.get("batch_size", 1000),
+                filter=filters,
+                limit=top_k,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                round_decimal=kwargs.get("round_decimal", -1),
+                partition_names=kwargs.get("partition_names", None),
+            )
 
         results = []
         try:
             while True:
-                _res = search_iterator.next()
+                if self._is_async:
+                    _res = await search_iterator.next()
+                else:
+                    _res = search_iterator.next()
                 res = _res.get_res()
                 if len(_res) == 0:
-                    search_iterator.close()
+                    if self._is_async:
+                        await search_iterator.close()
+                    else:
+                        search_iterator.close()
                     break
 
                 if distance_cutoff and res[0][-1].distance > distance_cutoff:
@@ -176,10 +218,16 @@ class MilvusRetriever(Retriever):
                      collection_name,
                      top_k)
 
-        if not self._validate_collection(collection_name):
+        if not await self._validate_collection(collection_name):
             raise CollectionNotFoundError(f"Collection: {collection_name} does not exist")
 
-        available_fields = [v.get("name") for v in self._client.describe_collection(collection_name).get("fields", {})]
+        # Get collection schema
+        if self._is_async:
+            collection_schema = await self._client.describe_collection(collection_name)
+        else:
+            collection_schema = self._client.describe_collection(collection_name)
+
+        available_fields = [v.get("name") for v in collection_schema.get("fields", [])]
 
         if self.content_field not in available_fields:
             raise ValueError(f"The specified content field: {self.content_field} is not part of the schema.")
@@ -194,17 +242,31 @@ class MilvusRetriever(Retriever):
         if self.content_field not in output_fields:
             output_fields.append(self.content_field)
 
-        search_vector = self._embedder.embed_query(query)
-        res = self._client.search(
-            collection_name=collection_name,
-            data=[search_vector],
-            filter=filters,
-            output_fields=output_fields,
-            search_params=search_params if search_params else {"metric_type": "L2"},
-            timeout=timeout,
-            anns_field=vector_field_name,
-            limit=top_k,
-        )
+        search_vector = await self._embedder.aembed_query(query)
+
+        # Perform search
+        if self._is_async:
+            res = await self._client.search(
+                collection_name=collection_name,
+                data=[search_vector],
+                filter=filters,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                limit=top_k,
+            )
+        else:
+            res = self._client.search(
+                collection_name=collection_name,
+                data=[search_vector],
+                filter=filters,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                limit=top_k,
+            )
 
         return _wrap_milvus_results(res[0], content_field=self.content_field)
 

nat/retriever/models.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/nemo_retriever/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/nemo_retriever/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/retriever/nemo_retriever/retriever.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -68,10 +68,10 @@ class NemoRetriever(Retriever):
 
     def bind(self, **kwargs) -> None:
         """
-        Bind default values to the search method. Cannot bind the 'query' parameter.
+        Bind default values to the search method. Cannot bind the ``query`` parameter.
 
         Args:
-          kwargs (dict): Key value pairs corresponding to the default values of search parameters.
+            kwargs: Key value pairs corresponding to the default values of search parameters.
         """
         if "query" in kwargs:
             kwargs = {k: v for k, v in kwargs.items() if k != "query"}
@@ -87,7 +87,7 @@ class NemoRetriever(Retriever):
 
     async def get_collections(self, client) -> list[Collection]:
         """
-        Get a list of all available collections as pydantic `Collection` objects
+        Get a list of all available collections as pydantic Collection objects.
         """
         collection_response = await client.get(urljoin(self.base_url, "/v1/collections"))
         collection_response.raise_for_status()
@@ -102,7 +102,7 @@ class NemoRetriever(Retriever):
 
     async def get_collection_by_name(self, collection_name, client) -> Collection:
         """
-        Retrieve a collection using it's name. Will return the first collection found if the name is ambiguous.
+        Retrieve a collection using its name. Will return the first collection found if the name is ambiguous.
         """
         collections = await self.get_collections(client)
         if (collection := next((c for c in collections if c.name == collection_name), None)) is None:

nat/retriever/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/runtime/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/runtime/loader.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -112,9 +112,16 @@ async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
     config = load_config(config_file)
 
     # Must yield the workflow function otherwise it cleans up
-    async with WorkflowBuilder.from_config(config=config) as workflow:
+    async with WorkflowBuilder.from_config(config=config) as builder:
 
-        yield SessionManager(await workflow.build(), max_concurrency=max_concurrency)
+        session_manager = await SessionManager.create(config=config,
+                                                      shared_builder=builder,
+                                                      max_concurrency=max_concurrency)
+
+        try:
+            yield session_manager
+        finally:
+            await session_manager.shutdown()
 
 
 @lru_cache

nat/runtime/metrics.py

@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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.
+"""Per-user workflow resource usage monitoring models and collector."""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+from pydantic import Field
+
+if TYPE_CHECKING:
+    from nat.runtime.session import SessionManager
+
+logger = logging.getLogger(__name__)
+
+
+class PerUserSessionMetrics(BaseModel):
+    """Session lifecycle metrics for a per-user workflow."""
+
+    created_at: datetime = Field(description="When the per-user workflow was created")
+    last_activity: datetime = Field(description="Last time the workflow was accessed")
+    ref_count: int = Field(ge=0, description="Current number of active references (in-flight requests)")
+    is_active: bool = Field(description="Whether the workflow is currently being used")
+
+
+class PerUserRequestMetrics(BaseModel):
+    """Request-level metrics for a per-user workflow."""
+
+    total_requests: int = Field(ge=0, default=0, description="Total number of requests processed")
+    active_requests: int = Field(ge=0, default=0, description="Number of currently active requests")
+    avg_latency_ms: float = Field(ge=0, default=0.0, description="Average request latency in milliseconds")
+    error_count: int = Field(ge=0, default=0, description="Total number of failed requests")
+
+
+class PerUserMemoryMetrics(BaseModel):
+    """Memory/resource count metrics for a per-user workflow."""
+
+    per_user_functions_count: int = Field(ge=0, default=0, description="Number of per-user functions built")
+    per_user_function_groups_count: int = Field(ge=0, default=0, description="Number of per-user function groups built")
+    exit_stack_size: int = Field(ge=0, default=0, description="Number of resources in the async exit stack")
+
+
+class PerUserResourceUsage(BaseModel):
+    """Combined resource usage metrics for a single per-user workflow."""
+
+    user_id: str = Field(description="The user identifier")
+    session: PerUserSessionMetrics = Field(description="Session lifecycle metrics")
+    requests: PerUserRequestMetrics = Field(description="Request-level metrics")
+    memory: PerUserMemoryMetrics = Field(description="Memory/resource count metrics")
+
+
+class PerUserMonitorResponse(BaseModel):
+    """Response model for the /monitor/users endpoint."""
+
+    timestamp: datetime = Field(default_factory=datetime.now, description="When the metrics were collected")
+    total_active_users: int = Field(ge=0, description="Number of users with active per-user workflows")
+    users: list[PerUserResourceUsage] = Field(default_factory=list, description="Per-user resource usage details")
+
+
+class PerUserMetricsCollector:
+    """Collector for per-user workflow metrics.
+
+    This class aggregates metrics from SessionManager's per-user builders
+    and provides methods to collect metrics for individual users or all users.
+    """
+
+    def __init__(self, session_manager: SessionManager):
+        """Initialize the collector with a SessionManager reference.
+
+        Args:
+            session_manager: The SessionManager instance to collect metrics from
+        """
+        self._session_manager = session_manager
+
+    async def collect_user_metrics(self, user_id: str) -> PerUserResourceUsage | None:
+        """Collect metrics for a specific user.
+
+        Args:
+            user_id: The user identifier to collect metrics for
+
+        Returns:
+            PerUserResourceUsage if user exists, None otherwise
+        """
+        async with self._session_manager._per_user_builders_lock:
+            if user_id not in self._session_manager._per_user_builders:
+                return None
+
+            builder_info = self._session_manager._per_user_builders[user_id]
+            return self._build_user_metrics(user_id, builder_info)
+
+    async def collect_all_metrics(self) -> PerUserMonitorResponse:
+        """Collect metrics for all active per-user workflows.
+
+        Returns:
+            PerUserMonitorResponse with all user metrics
+        """
+        users: list[PerUserResourceUsage] = []
+
+        async with self._session_manager._per_user_builders_lock:
+            for user_id, builder_info in self._session_manager._per_user_builders.items():
+                try:
+                    user_metrics = self._build_user_metrics(user_id, builder_info)
+                    users.append(user_metrics)
+                except Exception:
+                    logger.exception("Failed to collect metrics for user %s", user_id)
+
+        return PerUserMonitorResponse(
+            timestamp=datetime.now(),
+            total_active_users=len(users),
+            users=users,
+        )
+
+    def _build_user_metrics(self, user_id: str, builder_info) -> PerUserResourceUsage:
+        """Build metrics for a single user from builder info.
+
+        Args:
+            user_id: The user identifier
+            builder_info: The PerUserBuilderInfo instance
+
+        Returns:
+            PerUserResourceUsage with all metrics
+        """
+        # Session metrics
+        session_metrics = PerUserSessionMetrics(
+            created_at=builder_info.created_at,
+            last_activity=builder_info.last_activity,
+            ref_count=builder_info.ref_count,
+            is_active=builder_info.ref_count > 0,
+        )
+
+        # Request metrics
+        avg_latency = (builder_info.total_latency_ms /
+                       builder_info.total_requests if builder_info.total_requests > 0 else 0.0)
+
+        request_metrics = PerUserRequestMetrics(
+            total_requests=builder_info.total_requests,
+            active_requests=builder_info.ref_count,
+            avg_latency_ms=round(avg_latency, 2),
+            error_count=builder_info.error_count,
+        )
+
+        # Memory/resource count metrics from the builder
+        builder = builder_info.builder
+        per_user_functions_count = len(builder._per_user_functions)
+        per_user_function_groups_count = len(builder._per_user_function_groups)
+
+        # Count resources in exit stack (if accessible)
+        exit_stack = getattr(builder, '_exit_stack', None)
+        if exit_stack and hasattr(exit_stack, '_exit_callbacks'):
+            exit_stack_size = len(exit_stack._exit_callbacks)
+        else:
+            exit_stack_size = 0
+
+        memory_metrics = PerUserMemoryMetrics(
+            per_user_functions_count=per_user_functions_count,
+            per_user_function_groups_count=per_user_function_groups_count,
+            exit_stack_size=exit_stack_size,
+        )
+
+        return PerUserResourceUsage(
+            user_id=user_id,
+            session=session_metrics,
+            requests=request_metrics,
+            memory=memory_metrics,
+        )

nat/runtime/runner.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,16 +26,13 @@ from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.invocation_node import InvocationNode
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 from nat.observability.exporter_manager import ExporterManager
 from nat.utils.reactive.subject import Subject
 
 logger = logging.getLogger(__name__)
 
 
-class UserManagerBase:
-    pass
-
-
 class RunnerState(Enum):
     UNINITIALIZED = 0
     INITIALIZED = 1
@@ -53,7 +50,8 @@ class Runner:
                  input_message: typing.Any,
                  entry_fn: Function,
                  context_state: ContextState,
-                 exporter_manager: ExporterManager):
+                 exporter_manager: ExporterManager,
+                 runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
         """
         The Runner class is used to run a workflow. It handles converting input and output data types and running the
         workflow with the specified concurrency.
@@ -68,6 +66,8 @@ class Runner:
             The context state to use
         exporter_manager : ExporterManager
             The exporter manager to use
+        runtime_type : RuntimeTypeEnum
+            The runtime type (RUN_OR_SERVE, EVALUATE, OTHER)
         """
 
         if (entry_fn is None):
@@ -86,6 +86,9 @@ class Runner:
 
         self._exporter_manager = exporter_manager
 
+        self._runtime_type = runtime_type
+        self._runtime_type_token = None
+
     @property
     def context(self) -> Context:
         return self._context
@@ -105,6 +108,8 @@ class Runner:
             function_id="root",
         ))
 
+        self._runtime_type_token = self._context_state.runtime_type.set(self._runtime_type)
+
         if (self._state == RunnerState.UNINITIALIZED):
             self._state = RunnerState.INITIALIZED
         else:
@@ -119,6 +124,8 @@ class Runner:
 
         self._context_state.input_message.reset(self._input_message_token)
 
+        self._context_state.runtime_type.reset(self._runtime_type_token)
+
         if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
             raise ValueError("Cannot exit the context without completing the workflow")