nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/observability/processor/falsy_batch_filter_processor.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import TypeVar
+
+from nat.observability.processor.processor import Processor
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+FalsyT = TypeVar("FalsyT")
+
+
+class FalsyBatchFilterProcessor(Processor[list[FalsyT], list[FalsyT]]):
+    """Processor that filters out falsy items from a batch."""
+
+    @override
+    async def process(self, item: list[FalsyT]) -> list[FalsyT]:
+        """Filter out falsy items from a batch.
+
+        Args:
+            item (list[FalsyT]): The batch of items to filter.
+
+        Returns:
+            list[FalsyT]: The filtered batch.
+        """
+        return [i for i in item if i]
+
+
+class DictBatchFilterProcessor(FalsyBatchFilterProcessor[dict]):
+    """Processor that filters out empty dict items from a batch."""
+    pass
+
+
+class ListBatchFilterProcessor(FalsyBatchFilterProcessor[list]):
+    """Processor that filters out empty list items from a batch."""
+    pass
+
+
+class SetBatchFilterProcessor(FalsyBatchFilterProcessor[set]):
+    """Processor that filters out empty set items from a batch."""
+    pass

nat/observability/processor/processor.py

@@ -58,6 +58,9 @@ class Processor(Generic[InputT, OutputT], TypeIntrospectionMixin, ABC):
         chained processors.
     """
 
+    # All processors automatically use this for signature checking
+    _signature_method = 'process'
+
     @abstractmethod
     async def process(self, item: InputT) -> OutputT:
         """Process an item and return a potentially different type.

nat/observability/processor/processor_factory.py

@@ -0,0 +1,70 @@
+# 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 typing import Any
+
+from nat.observability.processor.processor import Processor
+
+
+def processor_factory(processor_class: type, from_type: type[Any], to_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and types.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        from_type (type[Any]): The type of the input data
+        to_type (type[Any]): The type of the output data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[from_type, to_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor
+
+
+def processor_factory_from_type(processor_class: type, from_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and input type.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        from_type (type[Any]): The type of the input data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[from_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor
+
+
+def processor_factory_to_type(processor_class: type, to_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and output type.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        to_type (type[Any]): The type of the output data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[to_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor

nat/observability/processor/redaction/__init__.py

@@ -0,0 +1,24 @@
+# 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 .redaction_processor import RedactionContext
+from .redaction_processor import RedactionContextState
+from .span_header_redaction_processor import SpanHeaderRedactionProcessor
+
+__all__ = [
+    "SpanHeaderRedactionProcessor",
+    "RedactionContext",
+    "RedactionContextState",
+]

nat/observability/processor/redaction/contextual_redaction_processor.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Any
+from typing import TypeVar
+
+from nat.observability.processor.redaction.redaction_processor import RedactionContext
+from nat.observability.processor.redaction.redaction_processor import RedactionContextState
+from nat.observability.processor.redaction.redaction_processor import RedactionInputT
+from nat.observability.processor.redaction.redaction_processor import RedactionProcessor
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+# Type variable for the data type extracted from context
+RedactionDataT = TypeVar('RedactionDataT')
+
+
+class ContextualRedactionProcessor(RedactionProcessor[RedactionInputT, RedactionDataT]):
+    """Generic processor with context-aware caching for any data type.
+
+    Provides a framework for redaction processors that need to:
+    - Extract data from the request context (headers, cookies, query params, etc.)
+    - Execute callbacks to determine redaction decisions
+    - Cache results within the request context to avoid redundant callback executions
+    - Handle race conditions with atomic operations
+
+    This class handles all the generic caching, context management, and callback
+    execution logic. Subclasses only need to implement data extraction and validation.
+
+    Args:
+        callback: Callable that determines if redaction should occur based on extracted data
+        enabled: Whether the processor is enabled
+        force_redact: If True, always redact regardless of data checks
+        redaction_value: The value to replace redacted attributes with
+    """
+
+    def __init__(
+        self,
+        callback: Callable[..., Any],
+        enabled: bool,
+        force_redact: bool,
+        redaction_value: str,
+    ):
+        self.callback = callback
+        self.enabled = enabled
+        self.force_redact = force_redact
+        self.redaction_value = redaction_value
+        self._redaction_context = RedactionContext(RedactionContextState())
+
+    @abstractmethod
+    def extract_data_from_context(self) -> RedactionDataT | None:
+        """Extract the relevant data from the context for redaction decision.
+
+        This method must be implemented by subclasses to extract their specific
+        data type (headers, cookies, query params, etc.) from the request context
+
+        Returns:
+            RedactionDataT | None: The extracted data, or None if no relevant data found
+        """
+        pass
+
+    @abstractmethod
+    def validate_data(self, data: RedactionDataT) -> bool:
+        """Validate that the extracted data is suitable for callback execution.
+
+        This method allows subclasses to implement their own validation logic
+        (e.g., checking if headers exist, if cookies are not empty, etc.).
+
+        Args:
+            data (RedactionDataT): The extracted data to validate
+
+        Returns:
+            bool: True if the data is valid for callback execution, False otherwise
+        """
+        pass
+
+    @override
+    async def should_redact(self, item: RedactionInputT) -> bool:
+        """Determine if this span should be redacted based on extracted data.
+
+        Extracts the relevant data from the context, validates it, and passes it to the
+        callback function to determine if redaction should occur. Results are cached
+        within the request context to avoid redundant callback executions.
+
+        Args:
+            item (RedactionInputT): The item to check
+
+        Returns:
+            bool: True if the span should be redacted, False otherwise
+        """
+        # If force_redact is enabled, always redact regardless of other conditions
+        if self.force_redact:
+            return True
+
+        if not self.enabled:
+            return False
+
+        # Extract data using subclass implementation
+        data = self.extract_data_from_context()
+        if data is None:
+            return False
+
+        # Validate data using subclass implementation
+        if not self.validate_data(data):
+            return False
+
+        # Use the generic caching framework for callback execution
+        async with self._redaction_context.redaction_manager() as manager:
+            return await manager.redaction_check(self.callback, data)

nat/observability/processor/redaction/contextual_span_redaction_processor.py

@@ -0,0 +1,66 @@
+# 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 Callable
+from typing import Any
+
+from nat.data_models.span import Span
+from nat.observability.processor.redaction.contextual_redaction_processor import ContextualRedactionProcessor
+from nat.observability.processor.redaction.redaction_processor import RedactionDataT
+from nat.utils.type_utils import override
+
+
+class ContextualSpanRedactionProcessor(ContextualRedactionProcessor[Span, RedactionDataT]):
+    """Processor that redacts the Span based on the Span attributes.
+
+    Args:
+        attributes: List of span attribute keys to redact
+        callback: Callable that determines if redaction should occur
+        enabled: Whether the processor is enabled
+        force_redact: If True, always redact regardless of callback
+        redaction_value: The value to replace redacted attributes with
+    """
+
+    def __init__(self,
+                 attributes: list[str],
+                 callback: Callable[..., Any],
+                 enabled: bool,
+                 force_redact: bool,
+                 redaction_value: str,
+                 redaction_tag: str | None = None):
+        super().__init__(callback=callback, enabled=enabled, force_redact=force_redact, redaction_value=redaction_value)
+        self.attributes = attributes
+        self.redaction_tag = redaction_tag
+
+    @override
+    async def redact_item(self, item: Span) -> Span:
+        """Redact specified attributes in the span.
+
+        Replaces the values of configured attributes with the redaction value.
+
+        Args:
+            item (Span): The span to redact
+
+        Returns:
+            Span: The span with redacted attributes
+        """
+        for key in self.attributes:
+            if key in item.attributes:
+                item.set_attribute(key, self.redaction_value)
+
+        if self.redaction_tag:
+            item.set_attribute(self.redaction_tag, True)
+
+        return item

nat/observability/processor/redaction/redaction_processor.py

@@ -0,0 +1,177 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import abstractmethod
+from collections.abc import AsyncGenerator
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from nat.observability.processor.processor import Processor
+from nat.utils.callable_utils import ainvoke_any
+from nat.utils.type_utils import override
+
+RedactionInputT = TypeVar('RedactionInputT')
+RedactionDataT = TypeVar('RedactionDataT')
+
+logger = logging.getLogger(__name__)
+
+
+class RedactionProcessor(Processor[RedactionInputT, RedactionInputT], Generic[RedactionInputT, RedactionDataT]):
+    """Abstract base class for redaction processors."""
+
+    @abstractmethod
+    async def should_redact(self, item: RedactionInputT) -> bool:
+        """Determine if this item should be redacted.
+
+        Args:
+            item (RedactionInputT): The item to check.
+
+        Returns:
+            bool: True if the item should be redacted, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    async def redact_item(self, item: RedactionInputT) -> RedactionInputT:
+        """Redact the item.
+
+        Args:
+            item (RedactionInputT): The item to redact.
+
+        Returns:
+            RedactionInputT: The redacted item.
+        """
+        pass
+
+    @override
+    async def process(self, item: RedactionInputT) -> RedactionInputT:
+        """Perform redaction on the item if it should be redacted.
+
+        Args:
+            item (RedactionInputT): The item to process.
+
+        Returns:
+            RedactionInputT: The processed item.
+        """
+        if await self.should_redact(item):
+            return await self.redact_item(item)
+        return item
+
+
+@dataclass
+class RedactionContextState:
+    """Generic context state for redaction results.
+
+    Stores the redaction result in a context variable to avoid redundant
+    callback executions within the same request context.
+    """
+
+    redaction_result: ContextVar[bool
+                                 | None] = field(default_factory=lambda: ContextVar("redaction_result", default=None))
+
+
+class RedactionManager(Generic[RedactionDataT]):
+    """Generic manager for atomic redaction operations.
+
+    Handles state mutations and ensures atomic callback execution
+    with proper result caching within a request context.
+
+    Args:
+        RedactionDataT: The type of data being processed for redaction decisions.
+    """
+
+    def __init__(self, context_state: RedactionContextState):
+        self._context_state = context_state
+
+    def set_redaction_result(self, result: bool) -> None:
+        """Set the redaction result in the context.
+
+        Args:
+            result (bool): The redaction result to cache.
+        """
+        self._context_state.redaction_result.set(result)
+
+    def clear_redaction_result(self) -> None:
+        """Clear the cached redaction result from the context."""
+        self._context_state.redaction_result.set(None)
+
+    async def redaction_check(self, callback: Callable[..., Any], data: RedactionDataT) -> bool:
+        """Execute redaction callback with atomic result caching.
+
+        Checks for existing cached results first, then executes the callback
+        and caches the result atomically. Since data is static per request,
+        subsequent calls within the same context return the cached result.
+
+        Supports sync/async functions, generators, and async generators.
+
+        Args:
+            callback (Callable[..., Any]): The callback to execute (sync/async function, generator, etc.).
+            data (RedactionDataT): The data to pass to the callback for redaction decision.
+
+        Returns:
+            bool: True if the item should be redacted, False otherwise.
+        """
+        # Check if we already have a result for this context
+        existing_result = self._context_state.redaction_result.get()
+        if existing_result is not None:
+            return existing_result
+
+        # Execute callback and cache result
+        result_value = await ainvoke_any(callback, data)
+        result = bool(result_value)
+        self.set_redaction_result(result)
+        return result
+
+
+class RedactionContext(Generic[RedactionDataT]):
+    """Generic context provider for redaction operations.
+
+    Provides read-only access to redaction state and manages the
+    RedactionManager lifecycle through async context managers.
+
+    Args:
+        RedactionDataT: The type of data being processed for redaction decisions.
+    """
+
+    def __init__(self, context: RedactionContextState):
+        self._context_state: RedactionContextState = context
+
+    @property
+    def redaction_result(self) -> bool | None:
+        """Get the current redaction result from context.
+
+        Returns:
+            bool | None: The cached redaction result, or None if not set.
+        """
+        return self._context_state.redaction_result.get()
+
+    @asynccontextmanager
+    async def redaction_manager(self) -> AsyncGenerator[RedactionManager[RedactionDataT], None]:
+        """Provide a redaction manager within an async context.
+
+        Creates and yields a RedactionManager instance for atomic
+        redaction operations within the current context.
+
+        Yields:
+            RedactionManager[RedactionDataT]: Manager instance for redaction operations.
+        """
+        yield RedactionManager(self._context_state)

nat/observability/processor/redaction/span_header_redaction_processor.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from starlette.datastructures import Headers
+
+from nat.builder.context import Context
+from nat.observability.processor.redaction.contextual_span_redaction_processor import ContextualSpanRedactionProcessor
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+
+class SpanHeaderRedactionProcessor(ContextualSpanRedactionProcessor[dict[str, Any]]):
+    """Processor that redacts the Span based on multiple headers and callback logic.
+
+    Uses context-scoped atomic updates to avoid redundant callback executions within a single context.
+    Since headers are static per request, the callback result is cached for the entire context using
+    an asynccontextmanager to ensure atomic operations.
+
+    Args:
+        headers: List of header keys to extract and pass to the callback
+        attributes: List of Span attribute keys to redact
+        callback: Callable that determines if redaction should occur
+        enabled: Whether the processor is enabled (default: True)
+        force_redact: If True, always redact regardless of header checks (default: False)
+        redaction_value: The value to replace redacted attributes with (default: "[REDACTED]")
+    """
+
+    def __init__(self,
+                 headers: list[str],
+                 attributes: list[str],
+                 callback: Callable[..., Any],
+                 enabled: bool = True,
+                 force_redact: bool = False,
+                 redaction_value: str = "[REDACTED]",
+                 redaction_tag: str | None = None):
+        # Initialize the base class with common parameters
+        super().__init__(attributes=attributes,
+                         callback=callback,
+                         enabled=enabled,
+                         force_redact=force_redact,
+                         redaction_value=redaction_value,
+                         redaction_tag=redaction_tag)
+        # Store header-specific configuration
+        self.headers = headers
+
+    @override
+    def extract_data_from_context(self) -> dict[str, Any] | None:
+        """Extract header data from the context.
+
+        Returns:
+            dict[str, Any] | None: Dictionary of header names to values, or None if no headers.
+        """
+
+        context = Context.get()
+        headers: Headers | None = context.metadata.headers
+
+        if headers is None or not self.headers:
+            return None
+
+        header_map: dict[str, Any] = {header: headers.get(header, None) for header in self.headers}
+
+        return header_map
+
+    @override
+    def validate_data(self, data: dict[str, Any]) -> bool:
+        """Validate that the extracted headers are suitable for callback execution.
+
+        Args:
+            data (dict[str, Any]): The extracted header dictionary.
+
+        Returns:
+            bool: True if headers exist and are not all None, False otherwise.
+        """
+        # Skip callback if no headers were found (all None values)
+        return bool(data) and not all(value is None for value in data.values())

nat/observability/processor/span_tagging_processor.py

@@ -0,0 +1,68 @@
+# 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 os
+from collections.abc import Mapping
+from enum import Enum
+
+from nat.data_models.span import Span
+from nat.observability.processor.processor import Processor
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+
+class SpanTaggingProcessor(Processor[Span, Span]):
+    """Processor that tags spans with multiple key-value metadata attributes.
+
+    This processor adds custom tags to spans by setting attributes with a configurable prefix.
+    Tags are applied for each key-value pair in the tags dictionary. The processor uses
+    a span prefix (configurable via NAT_SPAN_PREFIX environment variable) to namespace
+    the tag attributes.
+
+        Args:
+            tags: Mapping of tag keys to their values. Values can be enums (converted to strings) or strings
+            span_prefix: The prefix to use for tag attributes (default: from NAT_SPAN_PREFIX env var or "nat")
+    """
+
+    def __init__(self, tags: Mapping[str, Enum | str] | None = None, span_prefix: str | None = None):
+        self.tags = tags or {}
+
+        if span_prefix is None:
+            span_prefix = os.getenv("NAT_SPAN_PREFIX", "nat").strip() or "nat"
+
+        self._span_prefix = span_prefix
+
+    @override
+    async def process(self, item: Span) -> Span:
+        """Tag the span with all configured tags.
+
+        Args:
+            item (Span): The span to tag
+
+        Returns:
+            Span: The tagged span with all configured tags applied
+        """
+        for tag_key, tag_value in self.tags.items():
+            key = str(tag_key).strip()
+            if not key:
+                continue
+            value_str = str(tag_value.value) if isinstance(tag_value, Enum) else str(tag_value)
+            if value_str == "":
+                continue
+            item.set_attribute(f"{self._span_prefix}.{key}", value_str)
+
+        return item