nvidia-nat 1.3.0a20250909__py3-none-any.whl → 1.3.0a20250917__py3-none-any.whl

nat/front_ends/fastapi/main.py

@@ -13,19 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 import logging
 import os
+import typing
 
 from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorkerBase
+from nat.front_ends.fastapi.utils import get_config_file_path
+from nat.front_ends.fastapi.utils import import_class_from_string
 from nat.runtime.loader import load_config
 
+if typing.TYPE_CHECKING:
+    from fastapi import FastAPI
+
 logger = logging.getLogger(__name__)
 
 
-def get_app():
+def get_app() -> "FastAPI":
 
-    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    config_file_path = get_config_file_path()
     front_end_worker_full_name = os.getenv("NAT_FRONT_END_WORKER")
 
     if (not config_file_path):
@@ -36,28 +41,15 @@ def get_app():
 
     # Try to import the front end worker class
     try:
-        # Split the package from the class
-        front_end_worker_parts = front_end_worker_full_name.split(".")
-
-        front_end_worker_module_name = ".".join(front_end_worker_parts[:-1])
-        front_end_worker_class_name = front_end_worker_parts[-1]
-
-        front_end_worker_module = importlib.import_module(front_end_worker_module_name)
-
-        if not hasattr(front_end_worker_module, front_end_worker_class_name):
-            raise ValueError(f"Front end worker {front_end_worker_full_name} not found.")
-
-        front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = getattr(front_end_worker_module,
-                                                                                front_end_worker_class_name)
+        front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = import_class_from_string(
+            front_end_worker_full_name)
 
         if (not issubclass(front_end_worker_class, FastApiFrontEndPluginWorkerBase)):
             raise ValueError(
                 f"Front end worker {front_end_worker_full_name} is not a subclass of FastApiFrontEndPluginWorker.")
 
         # Load the config
-        abs_config_file_path = os.path.abspath(config_file_path)
-
-        config = load_config(abs_config_file_path)
+        config = load_config(config_file_path)
 
         # Create an instance of the front end worker class
         front_end_worker = front_end_worker_class(config)

nat/front_ends/fastapi/utils.py

@@ -0,0 +1,57 @@
+# 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 importlib
+import os
+
+
+def get_config_file_path() -> str:
+    """
+    Get the path to the NAT configuration file from the environment variable NAT_CONFIG_FILE.
+    Raises ValueError if the environment variable is not set.
+    """
+    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    if (not config_file_path):
+        raise ValueError("Config file not found in environment variable NAT_CONFIG_FILE.")
+
+    return os.path.abspath(config_file_path)
+
+
+def import_class_from_string(class_full_name: str) -> type:
+    """
+    Import a class from a string in the format 'module.submodule.ClassName'.
+    Raises ImportError if the class cannot be imported.
+    """
+    try:
+        class_name_parts = class_full_name.split(".")
+
+        module_name = ".".join(class_name_parts[:-1])
+        class_name = class_name_parts[-1]
+
+        module = importlib.import_module(module_name)
+
+        if not hasattr(module, class_name):
+            raise ValueError(f"Class '{class_full_name}' not found.")
+
+        return getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ImportError(f"Could not import {class_full_name}.") from e
+
+
+def get_class_name(cls: type) -> str:
+    """
+    Get the full class name including the module.
+    """
+    return f"{cls.__module__}.{cls.__qualname__}"

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -94,8 +94,9 @@ class MCPFrontEndPluginWorkerBase(ABC):
         functions: dict[str, Function] = {}
 
         # Extract all functions from the workflow
-        for function_name, function in workflow.functions.items():
-            functions[function_name] = function
+        functions.update(workflow.functions)
+        for function_group in workflow.function_groups.values():
+            functions.update(function_group.get_accessible_functions())
 
         functions[workflow.config.workflow.type] = workflow
 

nat/llm/aws_bedrock_llm.py

@@ -21,27 +21,38 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig,
+                            RetryMixin,
+                            OptimizableMixin,
+                            TemperatureMixin,
+                            TopPMixin,
+                            ThinkingMixin,
+                            name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     # Completion parameters
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted AWS Bedrock.")
-    max_tokens: int | None = Field(default=1024, gt=0, description="Maximum number of tokens to generate.")
+    max_tokens: int = OptimizableField(default=300,
+                                       description="Maximum number of tokens to generate.",
+                                       space=SearchSpace(high=2176, low=128, step=512))
     context_size: int | None = Field(
         default=1024,
         gt=0,
         description="The maximum number of tokens available for input. This is only required for LlamaIndex. "
-        "This field is ignored for Langchain.",
+        "This field is ignored for LangChain/LangGraph.",
     )
 
     # Client parameters

nat/llm/nim_llm.py

@@ -22,23 +22,34 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="nim"):
+class NIMModelConfig(LLMBaseConfig,
+                     RetryMixin,
+                     OptimizableMixin,
+                     TemperatureMixin,
+                     TopPMixin,
+                     ThinkingMixin,
+                     name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
     base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted NIM.")
-    max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
+    max_tokens: PositiveInt = OptimizableField(default=300,
+                                               description="Maximum number of tokens to generate.",
+                                               space=SearchSpace(high=2176, low=128, step=512))
 
 
 @register_llm_provider(config_type=NIMModelConfig)

nat/llm/openai_llm.py

@@ -21,13 +21,20 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableMixin
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig,
+                        RetryMixin,
+                        OptimizableMixin,
+                        TemperatureMixin,
+                        TopPMixin,
+                        ThinkingMixin,
+                        name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")

nat/observability/exporter/processing_exporter.py

@@ -53,6 +53,8 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
     - Configurable None filtering: processors returning None can drop items from pipeline
     - Automatic type validation before export
     """
+    # All ProcessingExporter instances automatically use this for signature checking
+    _signature_method = '_process_pipeline'
 
     def __init__(self, context_state: ContextState | None = None, drop_nones: bool = True):
         """Initialize the processing exporter.
@@ -294,8 +296,6 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             self._check_processor_compatibility(predecessor,
                                                 processor,
                                                 "predecessor",
-                                                predecessor.output_class,
-                                                processor.input_class,
                                                 str(predecessor.output_type),
                                                 str(processor.input_type))
 
@@ -304,8 +304,6 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             self._check_processor_compatibility(processor,
                                                 successor,
                                                 "successor",
-                                                processor.output_class,
-                                                successor.input_class,
                                                 str(processor.output_type),
                                                 str(successor.input_type))
 
@@ -313,34 +311,22 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
                                        source_processor: Processor,
                                        target_processor: Processor,
                                        relationship: str,
-                                       source_class: type,
-                                       target_class: type,
                                        source_type: str,
                                        target_type: str) -> None:
-        """Check type compatibility between two processors.
+        """Check type compatibility between two processors using Pydantic validation.
 
         Args:
             source_processor (Processor): The processor providing output
             target_processor (Processor): The processor receiving input
             relationship (str): Description of relationship ("predecessor" or "successor")
-            source_class (type): The output class of source processor
-            target_class (type): The input class of target processor
             source_type (str): String representation of source type
             target_type (str): String representation of target type
         """
-        try:
-            if not issubclass(source_class, target_class):
-                raise ValueError(f"Processor {target_processor.__class__.__name__} input type {target_type} "
-                                 f"is not compatible with {relationship} {source_processor.__class__.__name__} "
-                                 f"output type {source_type}")
-        except TypeError:
-            logger.warning(
-                "Cannot use issubclass() for type compatibility check between "
-                "%s (%s) and %s (%s). Skipping compatibility check.",
-                source_processor.__class__.__name__,
-                source_type,
-                target_processor.__class__.__name__,
-                target_type)
+        # Use Pydantic-based type compatibility checking
+        if not source_processor.is_output_compatible_with(target_processor.input_type):
+            raise ValueError(f"Processor {target_processor.__class__.__name__} input type {target_type} "
+                             f"is not compatible with {relationship} {source_processor.__class__.__name__} "
+                             f"output type {source_type}")
 
     async def _pre_start(self) -> None:
 
@@ -350,36 +336,21 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             last_processor = self._processors[-1]
 
             # validate that the first processor's input type is compatible with the exporter's input type
-            try:
-                if not issubclass(self.input_class, first_processor.input_class):
-                    raise ValueError(f"Processor {first_processor.__class__.__name__} input type "
-                                     f"{first_processor.input_type} is not compatible with the "
-                                     f"{self.input_type} input type")
-            except TypeError as e:
-                # Handle cases where classes are generic types that can't be used with issubclass
-                logger.warning(
-                    "Cannot validate type compatibility between %s (%s) "
-                    "and exporter (%s): %s. Skipping validation.",
-                    first_processor.__class__.__name__,
-                    first_processor.input_type,
-                    self.input_type,
-                    e)
-
+            if not first_processor.is_compatible_with_input(self.input_type):
+                logger.error("First processor %s input=%s incompatible with exporter input=%s",
+                             first_processor.__class__.__name__,
+                             first_processor.input_type,
+                             self.input_type)
+                raise ValueError("First processor incompatible with exporter input")
             # Validate that the last processor's output type is compatible with the exporter's output type
-            try:
-                if not DecomposedType.is_type_compatible(last_processor.output_type, self.output_type):
-                    raise ValueError(f"Processor {last_processor.__class__.__name__} output type "
-                                     f"{last_processor.output_type} is not compatible with the "
-                                     f"{self.output_type} output type")
-            except TypeError as e:
-                # Handle cases where classes are generic types that can't be used with issubclass
-                logger.warning(
-                    "Cannot validate type compatibility between %s (%s) "
-                    "and exporter (%s): %s. Skipping validation.",
-                    last_processor.__class__.__name__,
-                    last_processor.output_type,
-                    self.output_type,
-                    e)
+            # Use DecomposedType.is_type_compatible for the final export stage to allow batch compatibility
+            # This enables BatchingProcessor[T] -> Exporter[T] patterns where the exporter handles both T and list[T]
+            if not DecomposedType.is_type_compatible(last_processor.output_type, self.output_type):
+                logger.error("Last processor %s output=%s incompatible with exporter output=%s",
+                             last_processor.__class__.__name__,
+                             last_processor.output_type,
+                             self.output_type)
+                raise ValueError("Last processor incompatible with exporter output")
 
         # Lock the pipeline to prevent further modifications
         self._pipeline_locked = True
@@ -432,12 +403,15 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
                 await self.export_processed(processed_item)
             else:
                 logger.debug("Skipping export of empty batch")
-        elif isinstance(processed_item, self.output_class):
+        elif self.validate_output_type(processed_item):
             await self.export_processed(processed_item)
         else:
             if raise_on_invalid:
-                raise ValueError(f"Processed item {processed_item} is not a valid output type. "
-                                 f"Expected {self.output_class} or list[{self.output_class}]")
+                logger.error("Invalid processed item type for export: %s (expected %s or list[%s])",
+                             type(processed_item),
+                             self.output_type,
+                             self.output_type)
+                raise ValueError("Invalid processed item type for export")
             logger.warning("Processed item %s is not a valid output type for export", processed_item)
 
     async def _continue_pipeline_after(self, source_processor: Processor, item: Any) -> None:
@@ -512,7 +486,7 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             event (IntermediateStep): The event to be exported.
         """
         # Convert IntermediateStep to PipelineInputT and create export task
-        if isinstance(event, self.input_class):
+        if self.validate_input_type(event):
             input_item: PipelineInputT = event  # type: ignore
             coro = self._export_with_processing(input_item)
             self._create_export_task(coro)

nat/observability/mixin/redaction_config_mixin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -25,9 +25,10 @@ class RedactionConfigMixin(BaseModel):
     """
     redaction_enabled: bool = Field(default=False, description="Whether to enable redaction processing.")
     redaction_value: str = Field(default="[REDACTED]", description="Value to replace redacted attributes with.")
-    redaction_attributes: list[str] = Field(default_factory=lambda: ["input.value", "output.value", "metadata"],
-                                            description="Span attributes to redact when redaction is triggered.")
+    redaction_attributes: list[str] = Field(default_factory=lambda: ["input.value", "output.value", "nat.metadata"],
+                                            description="Attributes to redact when redaction is triggered.")
     force_redaction: bool = Field(default=False, description="Always redact regardless of other conditions.")
+    redaction_tag: str | None = Field(default=None, description="Tag to add to spans when redaction is triggered.")
 
 
 class HeaderRedactionConfigMixin(RedactionConfigMixin):
@@ -38,4 +39,4 @@ class HeaderRedactionConfigMixin(RedactionConfigMixin):
 
     Note: The callback function must be provided directly to the processor at runtime.
     """
-    redaction_header: str = Field(default="x-redaction-key", description="Header to check for redaction decisions.")
+    redaction_headers: list[str] = Field(default_factory=list, description="Headers to check for redaction decisions.")

nat/observability/mixin/tagging_config_mixin.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+from collections.abc import Mapping
 from enum import Enum
 from typing import Generic
 from typing import TypeVar
@@ -20,7 +22,17 @@ from typing import TypeVar
 from pydantic import BaseModel
 from pydantic import Field
 
-TagValueT = TypeVar("TagValueT")
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+TagMappingT = TypeVar("TagMappingT", bound=Mapping)
+
+
+class BaseTaggingConfigMixin(BaseModel, Generic[TagMappingT]):
+    """Base mixin for tagging spans."""
+    tags: TagMappingT | None = Field(default=None, description="Tags to add to the span.")
 
 
 class PrivacyLevel(str, Enum):
@@ -31,20 +43,20 @@ class PrivacyLevel(str, Enum):
     HIGH = "high"
 
 
-class TaggingConfigMixin(BaseModel, Generic[TagValueT]):
-    """Generic mixin for tagging spans with typed values.
+PrivacyTagSchema = TypedDict(
+    "PrivacyTagSchema",
+    {
+        "privacy.level": PrivacyLevel,
+    },
+    total=True,
+)
 
-    This mixin provides a flexible tagging system where both the tag key
-    and value type can be customized for different use cases.
-    """
-    tag_key: str | None = Field(default=None, description="Key to use when tagging traces.")
-    tag_value: TagValueT | None = Field(default=None, description="Value to tag the traces with.")
 
+class PrivacyTaggingConfigMixin(BaseTaggingConfigMixin[PrivacyTagSchema]):
+    """Mixin for privacy level tagging on spans."""
+    pass
 
-class PrivacyTaggingConfigMixin(TaggingConfigMixin[PrivacyLevel]):
-    """Mixin for privacy level tagging on spans.
 
-    Specializes TaggingConfigMixin to work with PrivacyLevel enum values,
-    providing a typed interface for privacy-related span tagging.
-    """
+class CustomTaggingConfigMixin(BaseTaggingConfigMixin[dict[str, str]]):
+    """Mixin for string key-value tagging on spans."""
     pass