nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__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/observability/register.py

@@ -77,6 +77,14 @@ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Bu
     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
 
 
@@ -95,4 +103,12 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     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/profiler/callbacks/langchain_callback_handler.py

@@ -33,6 +33,7 @@ from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import ServerToolUseSchema
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import ToolSchema
 from nat.data_models.intermediate_step import TraceMetadata
@@ -48,7 +49,14 @@ def _extract_tools_schema(invocation_params: dict) -> list:
     tools_schema = []
     if invocation_params is not None:
         for tool in invocation_params.get("tools", []):
-            tools_schema.append(ToolSchema(**tool))
+            try:
+                tools_schema.append(ToolSchema(**tool))
+            except Exception:
+                logger.debug(
+                    "Failed to parse tool schema from invocation params: %s. \n This "
+                    "can occur when the LLM server has native tools and can be ignored if "
+                    "using the responses API.",
+                    tool)
 
     return tools_schema
 
@@ -93,11 +101,15 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
             completion_tokens = usage_metadata.get("output_tokens", 0)
             total_tokens = usage_metadata.get("total_tokens", 0)
 
-            return TokenUsageBaseModel(
-                prompt_tokens=prompt_tokens,
-                completion_tokens=completion_tokens,
-                total_tokens=total_tokens,
-            )
+            cache_tokens = usage_metadata.get("input_token_details", {}).get("cache_read", 0)
+
+            reasoning_tokens = usage_metadata.get("output_token_details", {}).get("reasoning", 0)
+
+            return TokenUsageBaseModel(prompt_tokens=prompt_tokens,
+                                       completion_tokens=completion_tokens,
+                                       total_tokens=total_tokens,
+                                       cached_tokens=cache_tokens,
+                                       reasoning_tokens=reasoning_tokens)
         return TokenUsageBaseModel()
 
     async def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
@@ -213,6 +225,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
         except IndexError:
             generation = None
 
+        message = None
         if isinstance(generation, ChatGeneration):
             try:
                 message = generation.message
@@ -232,6 +245,17 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
         else:
             llm_text_output = ""
 
+        tool_outputs_list = []
+        # Check if message.additional_kwargs as tool_outputs indicative of server side tool calling
+        if message and message.additional_kwargs and "tool_outputs" in message.additional_kwargs:
+            tools_outputs = message.additional_kwargs["tool_outputs"]
+            if isinstance(tools_outputs, list):
+                for tool in tools_outputs:
+                    try:
+                        tool_outputs_list.append(ServerToolUseSchema(**tool))
+                    except Exception:
+                        pass
+
         # update shared state behind lock
         with self._lock:
             usage_stat = IntermediateStepPayload(
@@ -243,7 +267,8 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
                 data=StreamEventData(input=self._run_id_to_llm_input.get(str(kwargs.get("run_id", "")), ""),
                                      output=llm_text_output),
                 usage_info=UsageInfo(token_usage=self._extract_token_base_model(usage_metadata)),
-                metadata=TraceMetadata(chat_responses=[generation] if generation else []))
+                metadata=TraceMetadata(chat_responses=[generation] if generation else [],
+                                       tool_outputs=tool_outputs_list if tool_outputs_list else []))
 
             self.step_manager.push_intermediate_step(usage_stat)