nvidia-nat 1.1.0a20251020__py3-none-any.whl

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

nat/observability/register.py

@@ -0,0 +1,114 @@
+# 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 pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_logging_method
+from nat.cli.register_workflow import register_telemetry_exporter
+from nat.data_models.logging import LoggingBaseConfig
+from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+from nat.observability.mixin.file_mode import FileMode
+
+logger = logging.getLogger(__name__)
+
+
+class FileTelemetryExporterConfig(TelemetryExporterBaseConfig, name="file"):
+    """A telemetry exporter that writes runtime traces to local files with optional rolling."""
+
+    output_path: str = Field(description="Output path for logs. When rolling is disabled: exact file path. "
+                             "When rolling is enabled: directory path or file path (directory + base name).")
+    project: str = Field(description="Name to affiliate with this application.")
+    mode: FileMode = Field(
+        default=FileMode.APPEND,
+        description="File write mode: 'append' to add to existing file or 'overwrite' to start fresh.")
+    enable_rolling: bool = Field(default=False, description="Enable rolling log files based on size limits.")
+    max_file_size: int = Field(
+        default=10 * 1024 * 1024,  # 10MB
+        description="Maximum file size in bytes before rolling to a new file.")
+    max_files: int = Field(default=5, description="Maximum number of rolled files to keep.")
+    cleanup_on_init: bool = Field(default=False, description="Clean up old files during initialization.")
+
+
+@register_telemetry_exporter(config_type=FileTelemetryExporterConfig)
+async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):
+    """
+    Build and return a FileExporter for file-based telemetry export with optional rolling.
+    """
+
+    from nat.observability.exporter.file_exporter import FileExporter
+
+    yield FileExporter(output_path=config.output_path,
+                       project=config.project,
+                       mode=config.mode,
+                       enable_rolling=config.enable_rolling,
+                       max_file_size=config.max_file_size,
+                       max_files=config.max_files,
+                       cleanup_on_init=config.cleanup_on_init)
+
+
+class ConsoleLoggingMethodConfig(LoggingBaseConfig, name="console"):
+    """A logger to write runtime logs to the console."""
+
+    level: str = Field(description="The logging level of console logger.")
+
+
+@register_logging_method(config_type=ConsoleLoggingMethodConfig)
+async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):
+    """
+    Build and return a StreamHandler for console-based logging.
+    """
+    import sys
+
+    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
+
+
+class FileLoggingMethod(LoggingBaseConfig, name="file"):
+    """A logger to write runtime logs to a file."""
+
+    path: str = Field(description="The file path to save the logging output.")
+    level: str = Field(description="The logging level of file logger.")
+
+
+@register_logging_method(config_type=FileLoggingMethod)
+async def file_logging_method(config: FileLoggingMethod, builder: Builder):
+    """
+    Build and return a FileHandler for file-based logging.
+    """
+    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/observability/utils/__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.