nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/observability/mixin/file_mode.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 enum import StrEnum
+
+
+class FileMode(StrEnum):
+    """File write modes for FileExportMixin."""
+
+    APPEND = "append"
+    OVERWRITE = "overwrite"

aiq/observability/mixin/resource_conflict_mixin.py

@@ -0,0 +1,134 @@
+# 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
+import weakref
+from abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class ResourceConflictError(ValueError):
+    """Raised when multiple exporter instances would conflict over the same resource."""
+    pass
+
+
+class ResourceConflictMixin(ABC):
+    """Abstract mixin for detecting resource conflicts between exporter instances.
+
+    This mixin provides a framework for exporters to detect when multiple instances
+    would conflict over the same resources (files, database tables, API endpoints, etc.).
+    Each concrete implementation defines what constitutes a resource conflict for that
+    exporter type.
+
+    The mixin maintains class-level registries using weakrefs for automatic cleanup
+    when instances are garbage collected.
+    """
+
+    # Each subclass gets its own registry - prevents cross-contamination
+    _registries: dict[type, dict[str, weakref.ref]] = {}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Register this instance's resources and check for conflicts
+        self._register_resources()
+
+    @abstractmethod
+    def _get_resource_identifiers(self) -> dict[str, Any]:
+        """Return dict of resource_type -> identifier that this instance will use.
+
+        Examples:
+            Files: {"file_path": "/logs/app.log", "cleanup_pattern": "app_*.log"}
+            Phoenix: {"project_endpoint": "my_project@http://localhost:6006"}
+            Database: {"table_name": "events", "connection": "postgresql://..."}
+
+        Returns:
+            dict[str, Any]: Dict mapping resource type names to unique identifiers for those resources.
+        """
+        pass
+
+    @abstractmethod
+    def _format_conflict_error(self, resource_type: str, identifier: Any, existing_instance: Any) -> str:
+        """Format a user-friendly error message for a resource conflict.
+
+        Args:
+            resource_type (str): The type of resource that conflicts (e.g., "file_path", "project_endpoint")
+            identifier (Any): The identifier for this resource
+            existing_instance (Any): The existing instance that conflicts with this one
+
+        Returns:
+            A clear error message explaining the conflict and how to resolve it.
+        """
+        pass
+
+    def _register_resources(self):
+        """Register this instance's resources and check for conflicts.
+
+        Raises:
+            ResourceConflictError: If any resource conflicts with an existing instance.
+        """
+        # Get our class-specific registry
+        cls = type(self)
+        if cls not in self._registries:
+            self._registries[cls] = {}
+        registry = self._registries[cls]
+
+        # Clean up dead references first
+        self._cleanup_dead_references(registry)
+
+        # Check each resource for conflicts
+        resources = self._get_resource_identifiers()
+        for resource_type, identifier in resources.items():
+            resource_key = f"{resource_type}:{identifier}"
+
+            # Check for existing instance using this resource
+            if resource_key in registry:
+                existing_ref = registry[resource_key]
+                existing_instance = existing_ref()
+                if existing_instance is not None:
+                    error_msg = self._format_conflict_error(resource_type, identifier, existing_instance)
+                    raise ResourceConflictError(error_msg)
+
+            # Register this instance for this resource
+            registry[resource_key] = weakref.ref(self, lambda ref, key=resource_key: registry.pop(key, None))
+
+        logger.debug("Registered %d resources for %s", len(resources), self.__class__.__name__)
+
+    def _cleanup_dead_references(self, registry: dict[str, weakref.ref]):
+        """Remove dead weakref entries from the registry.
+
+        Args:
+            registry (dict[str, weakref.ref]): The registry to clean up.
+        """
+        dead_keys = [key for key, ref in registry.items() if ref() is None]
+        for key in dead_keys:
+            registry.pop(key, None)
+
+    @classmethod
+    def get_active_resource_count(cls) -> int:
+        """Get the number of active resources registered for this class.
+
+        Returns:
+            int: Number of active resource registrations.
+        """
+        if cls not in cls._registries:
+            return 0
+
+        registry = cls._registries[cls]
+        # Clean up and count live references
+        live_count = sum(1 for ref in registry.values() if ref() is not None)
+        return live_count

aiq/observability/mixin/serialize_mixin.py

@@ -0,0 +1,61 @@
+# 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
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import TypeAdapter
+
+
+class SerializeMixin:
+
+    def _process_streaming_output(self, input_value: Any) -> Any:
+        """
+        Serialize a list of values to a JSON string.
+        """
+        if isinstance(input_value, BaseModel):
+            return json.loads(TypeAdapter(type(input_value)).dump_json(input_value).decode('utf-8'))
+        if isinstance(input_value, dict):
+            return input_value
+        return input_value
+
+    def _serialize_payload(self, input_value: Any) -> tuple[str, bool]:
+        """
+        Serialize the input value to a string. Returns a tuple with the serialized value and a boolean indicating if the
+        serialization is JSON or a string.
+
+        Args:
+            input_value (Any): The input value to serialize.
+
+        Returns:
+            tuple[str, bool]: A tuple with the serialized value and a boolean indicating if the serialization is
+                JSON or a string.
+        """
+        try:
+            if isinstance(input_value, BaseModel):
+                return TypeAdapter(type(input_value)).dump_json(input_value).decode('utf-8'), True
+            if isinstance(input_value, dict):
+                return json.dumps(input_value), True
+            if isinstance(input_value, list):
+                serialized_list = []
+                for value in input_value:
+                    serialized_value = self._process_streaming_output(value)
+                    serialized_list.append(serialized_value)
+                return json.dumps(serialized_list), True
+            return str(input_value), False
+        except Exception:
+            # Fallback to string representation if we can't serialize using pydantic
+            return str(input_value), False

aiq/observability/mixin/type_introspection_mixin.py

@@ -0,0 +1,183 @@
+# 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 functools import lru_cache
+from typing import Any
+from typing import get_args
+from typing import get_origin
+
+
+class TypeIntrospectionMixin:
+    """Mixin class providing type introspection capabilities for generic classes.
+
+    This mixin extracts type information from generic class definitions,
+    allowing classes to determine their InputT and OutputT types at runtime.
+    """
+
+    def _find_generic_types(self) -> tuple[type[Any], type[Any]] | None:
+        """
+        Recursively search through the inheritance hierarchy to find generic type parameters.
+
+        This method handles cases where a class inherits from a generic parent class,
+        resolving the concrete types through the inheritance chain.
+
+        Returns:
+            tuple[type[Any], type[Any]] | None: (input_type, output_type) if found, None otherwise
+        """
+        # First, try to find types directly in this class's __orig_bases__
+        for base_cls in getattr(self.__class__, '__orig_bases__', []):
+            base_cls_args = get_args(base_cls)
+
+            # Direct case: MyClass[InputT, OutputT]
+            if len(base_cls_args) >= 2:
+                return base_cls_args[0], base_cls_args[1]
+
+            # Indirect case: MyClass[SomeGeneric[ConcreteType]]
+            # Need to resolve the generic parent's types
+            if len(base_cls_args) == 1:
+                base_origin = get_origin(base_cls)
+                if base_origin and hasattr(base_origin, '__orig_bases__'):
+                    # Look at the parent's generic definition
+                    for parent_base in getattr(base_origin, '__orig_bases__', []):
+                        parent_args = get_args(parent_base)
+                        if len(parent_args) >= 2:
+                            # Found the pattern: ParentClass[T, list[T]]
+                            # Substitute T with our concrete type
+                            concrete_type = base_cls_args[0]
+                            input_type = self._substitute_type_var(parent_args[0], concrete_type)
+                            output_type = self._substitute_type_var(parent_args[1], concrete_type)
+                            return input_type, output_type
+
+        return None
+
+    def _substitute_type_var(self, type_expr: Any, concrete_type: type) -> type[Any]:
+        """
+        Substitute TypeVar in a type expression with a concrete type.
+
+        Args:
+            type_expr: The type expression potentially containing TypeVars
+            concrete_type: The concrete type to substitute
+
+        Returns:
+            The type expression with TypeVars substituted
+        """
+        from typing import TypeVar
+
+        # If it's a TypeVar, substitute it
+        if isinstance(type_expr, TypeVar):
+            return concrete_type
+
+        # If it's a generic type like list[T], substitute the args
+        origin = get_origin(type_expr)
+        args = get_args(type_expr)
+
+        if origin and args:
+            # Recursively substitute in the arguments
+            new_args = tuple(self._substitute_type_var(arg, concrete_type) for arg in args)
+            # Reconstruct the generic type
+            return origin[new_args]
+
+        # Otherwise, return as-is
+        return type_expr
+
+    @property
+    @lru_cache
+    def input_type(self) -> type[Any]:
+        """
+        Get the input type of the class. The input type is determined by the generic parameters of the class.
+
+        For example, if a class is defined as `MyClass[list[int], str]`, the `input_type` is `list[int]`.
+
+        Returns
+        -------
+        type[Any]
+            The input type specified in the generic parameters
+
+        Raises
+        ------
+        ValueError
+            If the input type cannot be determined from the class definition
+        """
+        types = self._find_generic_types()
+        if types:
+            return types[0]
+
+        raise ValueError(f"Could not find input type for {self.__class__.__name__}")
+
+    @property
+    @lru_cache
+    def output_type(self) -> type[Any]:
+        """
+        Get the output type of the class. The output type is determined by the generic parameters of the class.
+
+        For example, if a class is defined as `MyClass[list[int], str]`, the `output_type` is `str`.
+
+        Returns
+        -------
+        type[Any]
+            The output type specified in the generic parameters
+
+        Raises
+        ------
+        ValueError
+            If the output type cannot be determined from the class definition
+        """
+        types = self._find_generic_types()
+        if types:
+            return types[1]
+
+        raise ValueError(f"Could not find output type for {self.__class__.__name__}")
+
+    @property
+    @lru_cache
+    def input_class(self) -> type:
+        """
+        Get the python class of the input type. This is the class that can be used to check if a value is an
+        instance of the input type. It removes any generic or annotation information from the input type.
+
+        For example, if the input type is `list[int]`, the `input_class` is `list`.
+
+        Returns
+        -------
+        type
+            The python type of the input type
+        """
+        input_origin = get_origin(self.input_type)
+
+        if input_origin is None:
+            return self.input_type
+
+        return input_origin
+
+    @property
+    @lru_cache
+    def output_class(self) -> type:
+        """
+        Get the python class of the output type. This is the class that can be used to check if a value is an
+        instance of the output type. It removes any generic or annotation information from the output type.
+
+        For example, if the output type is `list[int]`, the `output_class` is `list`.
+
+        Returns
+        -------
+        type
+            The python type of the output type
+        """
+        output_origin = get_origin(self.output_type)
+
+        if output_origin is None:
+            return self.output_type
+
+        return output_origin

aiq/observability/processor/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.