nvidia-nat 1.3.0a20250909__py3-none-any.whl → 1.3.0a20250917__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

@@ -15,6 +15,8 @@
 
 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
@@ -24,22 +26,20 @@ logger = logging.getLogger(__name__)
 
 
 class SpanTaggingProcessor(Processor[Span, Span]):
-    """Processor that tags spans with key-value metadata attributes.
+    """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 only applied when both tag_key and tag_value are provided. The processor uses
+    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:
-        tag_key: The key name for the tag to add to spans.
-        tag_value: The value for the tag to add to spans.
-        span_prefix: The prefix to use for tag attributes (default: from NAT_SPAN_PREFIX env var or "nat").
+        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, tag_key: str | None = None, tag_value: str | None = None, span_prefix: str | None = None):
-        self.tag_key = tag_key
-        self.tag_value = tag_value
+    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"
@@ -48,14 +48,21 @@ class SpanTaggingProcessor(Processor[Span, Span]):
 
     @override
     async def process(self, item: Span) -> Span:
-        """Tag the span with a tag if both tag_key and tag_value are provided.
+        """Tag the span with all configured tags.
 
         Args:
-            item (Span): The span to tag.
+            item (Span): The span to tag
 
         Returns:
-            Span: The tagged span.
+            Span: The tagged span with all configured tags applied
         """
-        if self.tag_key and self.tag_value:
-            item.set_attribute(f"{self._span_prefix}.{self.tag_key}", self.tag_value)
+        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/profiler/decorators/framework_wrapper.py

@@ -51,16 +51,19 @@ def set_framework_profiler_handler(
         @asynccontextmanager
         async def wrapper(workflow_config, builder):
 
-            if LLMFrameworkEnum.LANGCHAIN in frameworks and not _library_instrumented["langchain"]:
-                from langchain_core.tracers.context import register_configure_hook
-
+            if LLMFrameworkEnum.LANGCHAIN in frameworks:
+                # Always set a fresh handler in the current context so callbacks
+                # route to the active run. Only register the hook once globally.
                 from nat.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
 
                 handler = LangchainProfilerHandler()
                 callback_handler_var.set(handler)
-                register_configure_hook(callback_handler_var, inheritable=True)
-                _library_instrumented["langchain"] = True
-                logger.debug("Langchain callback handler registered")
+
+                if not _library_instrumented["langchain"]:
+                    from langchain_core.tracers.context import register_configure_hook
+                    register_configure_hook(callback_handler_var, inheritable=True)
+                    _library_instrumented["langchain"] = True
+                    logger.debug("LangChain/LangGraph callback hook registered")
 
             if LLMFrameworkEnum.LLAMA_INDEX in frameworks:
                 from llama_index.core import Settings

nat/profiler/parameter_optimization/optimizable_utils.py

@@ -0,0 +1,93 @@
+# 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 get_args
+from typing import get_origin
+
+from pydantic import BaseModel
+
+from nat.data_models.optimizable import SearchSpace
+
+logger = logging.getLogger(__name__)
+
+
+def walk_optimizables(obj: BaseModel, path: str = "") -> dict[str, SearchSpace]:
+    """
+    Recursively build ``{flattened.path: SearchSpace}`` for every optimizable
+    field inside *obj*.
+
+    * Honors ``optimizable_params`` on any model that mixes in
+      ``OptimizableMixin`` – only listed fields are kept.
+    * If a model contains optimizable fields **but** omits
+      ``optimizable_params``, we emit a warning and skip them.
+    """
+    spaces: dict[str, SearchSpace] = {}
+
+    allowed_params_raw = getattr(obj, "optimizable_params", None)
+    allowed_params = set(allowed_params_raw) if allowed_params_raw is not None else None
+    overrides = getattr(obj, "search_space", {}) or {}
+    has_optimizable_flag = False
+
+    for name, fld in obj.model_fields.items():
+        full = f"{path}.{name}" if path else name
+        extra = fld.json_schema_extra or {}
+
+        is_field_optimizable = extra.get("optimizable", False) or name in overrides
+        has_optimizable_flag = has_optimizable_flag or is_field_optimizable
+
+        # honour allow-list
+        if allowed_params is not None and name not in allowed_params:
+            continue
+
+        # 1. plain optimizable field or override from config
+        if is_field_optimizable:
+            space = overrides.get(name, extra.get("search_space"))
+            if space is None:
+                logger.error(
+                    "Field %s is marked optimizable but no search space was provided.",
+                    full,
+                )
+                raise ValueError(f"Field {full} is marked optimizable but no search space was provided")
+            spaces[full] = space
+
+        value = getattr(obj, name, None)
+
+        # 2. nested BaseModel
+        if isinstance(value, BaseModel):
+            spaces.update(walk_optimizables(value, full))
+
+        # 3. dict[str, BaseModel] container
+        elif isinstance(value, dict):
+            for key, subval in value.items():
+                if isinstance(subval, BaseModel):
+                    spaces.update(walk_optimizables(subval, f"{full}.{key}"))
+
+        # 4. static-type fallback for class-level annotations
+        elif isinstance(obj, type):
+            ann = fld.annotation
+            if get_origin(ann) in (dict, dict):
+                _, val_t = get_args(ann) or (None, None)
+                if isinstance(val_t, type) and issubclass(val_t, BaseModel):
+                    if allowed_params is None or name in allowed_params:
+                        spaces[f"{full}.*"] = SearchSpace(low=None, high=None)  # sentinel
+
+    if allowed_params is None and has_optimizable_flag:
+        logger.warning(
+            "Model %s contains optimizable fields but no `optimizable_params` "
+            "were defined; these fields will be ignored.",
+            obj.__class__.__name__,
+        )
+    return spaces