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

nat/data_models/discovery_metadata.py

@@ -177,7 +177,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         description = generate_config_type_docs(config_type=config_type)
@@ -217,7 +217,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         if isinstance(wrapper_type, LLMFrameworkEnum):
@@ -252,7 +252,7 @@ class DiscoveryMetadata(BaseModel):
                 description = ""
                 package_version = package_version or ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         return DiscoveryMetadata(package=package_name,
@@ -290,7 +290,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type

nat/data_models/evaluate.py

@@ -57,6 +57,9 @@ class EvalOutputConfig(BaseModel):
     dir: Path = Path("./.tmp/nat/examples/default/")
     # S3 prefix for the workflow and evaluation results
     remote_dir: str | None = None
+    # Custom function to pre-evaluation process the eval input
+    # Format: "module.path.function_name"
+    custom_pre_eval_process_function: str | None = None
     # Custom scripts to run after the workflow and evaluation results are saved
     custom_scripts: dict[str, EvalCustomScriptConfig] = {}
     # S3 config for uploading the contents of the output directory
@@ -108,7 +111,7 @@ class EvalConfig(BaseModel):
     @classmethod
     def rebuild_annotations(cls):
 
-        from nat.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+        from nat.cli.type_registry import GlobalTypeRegistry
 
         type_registry = GlobalTypeRegistry.get()
 

nat/data_models/function.py

@@ -15,6 +15,10 @@
 
 import typing
 
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
 from .common import BaseModelRegistryTag
 from .common import TypedBaseModel
 
@@ -23,8 +27,38 @@ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
     pass
 
 
+class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for function groups.
+
+    Function groups enable sharing of configurations and resources across multiple functions.
+    """
+    include: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be added to the global Function registry",
+    )
+    exclude: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be excluded from default access to the group",
+    )
+
+    @field_validator("include", "exclude")
+    @classmethod
+    def _validate_fields_include_exclude(cls, value: list[str]) -> list[str]:
+        if len(set(value)) != len(value):
+            raise ValueError("Function names must be unique")
+        return sorted(value)
+
+    @model_validator(mode="after")
+    def _validate_include_exclude(self):
+        if self.include and self.exclude:
+            raise ValueError("include and exclude cannot be used together")
+        return self
+
+
 class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
     pass
 
 
 FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)
+
+FunctionGroupConfigT = typing.TypeVar("FunctionGroupConfigT", bound=FunctionGroupBaseConfig)

nat/data_models/function_dependencies.py

@@ -23,6 +23,7 @@ class FunctionDependencies(BaseModel):
     A class to represent the dependencies of a function.
     """
     functions: set[str] = Field(default_factory=set)
+    function_groups: set[str] = Field(default_factory=set)
     llms: set[str] = Field(default_factory=set)
     embedders: set[str] = Field(default_factory=set)
     memory_clients: set[str] = Field(default_factory=set)
@@ -33,6 +34,10 @@ class FunctionDependencies(BaseModel):
     def serialize_functions(self, v: set[str]) -> list[str]:
         return list(v)
 
+    @field_serializer("function_groups", when_used="json")
+    def serialize_function_groups(self, v: set[str]) -> list[str]:
+        return list(v)
+
     @field_serializer("llms", when_used="json")
     def serialize_llms(self, v: set[str]) -> list[str]:
         return list(v)
@@ -54,19 +59,22 @@ class FunctionDependencies(BaseModel):
         return list(v)
 
     def add_function(self, function: str):
-        self.functions.add(function)  # pylint: disable=no-member
+        self.functions.add(function)
+
+    def add_function_group(self, function_group: str):
+        self.function_groups.add(function_group)  # pylint: disable=no-member
 
     def add_llm(self, llm: str):
-        self.llms.add(llm)  # pylint: disable=no-member
+        self.llms.add(llm)
 
     def add_embedder(self, embedder: str):
-        self.embedders.add(embedder)  # pylint: disable=no-member
+        self.embedders.add(embedder)
 
     def add_memory_client(self, memory_client: str):
-        self.memory_clients.add(memory_client)  # pylint: disable=no-member
+        self.memory_clients.add(memory_client)
 
     def add_object_store(self, object_store: str):
-        self.object_stores.add(object_store)  # pylint: disable=no-member
+        self.object_stores.add(object_store)
 
     def add_retriever(self, retriever: str):
-        self.retrievers.add(retriever)  # pylint: disable=no-member
+        self.retrievers.add(retriever)

nat/data_models/gated_field_mixin.py

@@ -0,0 +1,242 @@
+# 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 Sequence
+from dataclasses import dataclass
+from re import Pattern
+
+from pydantic import model_validator
+
+
+@dataclass
+class GatedFieldMixinConfig:
+    """Configuration for a gated field mixin."""
+
+    field_name: str
+    default_if_supported: object | None
+    unsupported: Sequence[Pattern[str]] | None
+    supported: Sequence[Pattern[str]] | None
+    keys: Sequence[str]
+
+
+class GatedFieldMixin:
+    """
+    A mixin that gates a field based on specified keys.
+
+    This should be used to automatically validate a field based on a given key.
+
+    Parameters
+    ----------
+    field_name: `str`
+                The name of the field.
+    default_if_supported: `object | None`
+                          The default value of the field if it is supported for the key.
+    keys: `Sequence[str]`
+          A sequence of keys that are used to validate the field.
+    unsupported: `Sequence[Pattern[str]] | None`
+                 A sequence of regex patterns that match the key names NOT supported for the field.
+                 Defaults to None.
+    supported: `Sequence[Pattern[str]] | None`
+               A sequence of regex patterns that match the key names supported for the field.
+               Defaults to None.
+    """
+
+    def __init_subclass__(
+        cls,
+        field_name: str | None = None,
+        default_if_supported: object | None = None,
+        keys: Sequence[str] | None = None,
+        unsupported: Sequence[Pattern[str]] | None = None,
+        supported: Sequence[Pattern[str]] | None = None,
+    ) -> None:
+        """Store the class variables for the field and define the gated field validator."""
+        super().__init_subclass__()
+
+        # Check if this class directly inherits from GatedFieldMixin
+        has_gated_field_mixin = GatedFieldMixin in cls.__bases__
+
+        if has_gated_field_mixin:
+            if keys is None:
+                raise ValueError("keys must be provided when subclassing GatedFieldMixin")
+            if field_name is None:
+                raise ValueError("field_name must be provided when subclassing GatedFieldMixin")
+
+            cls._setup_direct_mixin(field_name, default_if_supported, unsupported, supported, keys)
+
+        # Always try to collect mixins and create validators for multiple inheritance
+        # This handles both direct inheritance and deep inheritance chains
+        all_mixins = cls._collect_all_mixin_configs()
+        if all_mixins:
+            cls._create_combined_validator(all_mixins)
+
+    @classmethod
+    def _setup_direct_mixin(
+        cls,
+        field_name: str,
+        default_if_supported: object | None,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> None:
+        """Set up a class that directly inherits from GatedFieldMixin."""
+        cls._validate_mixin_parameters(unsupported, supported, keys)
+
+        # Create and store validator
+        validator = cls._create_gated_field_validator(field_name, default_if_supported, unsupported, supported, keys)
+        validator_name = f"_gated_field_validator_{field_name}"
+        setattr(cls, validator_name, validator)
+
+        # Store mixin info for multiple inheritance
+        if not hasattr(cls, "_gated_field_mixins"):
+            cls._gated_field_mixins = []
+
+        cls._gated_field_mixins.append(
+            GatedFieldMixinConfig(
+                field_name,
+                default_if_supported,
+                unsupported,
+                supported,
+                keys,
+            ))
+
+    @classmethod
+    def _validate_mixin_parameters(
+        cls,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> None:
+        """Validate that all required parameters are provided."""
+        if unsupported is None and supported is None:
+            raise ValueError("Either unsupported or supported must be provided")
+        if unsupported is not None and supported is not None:
+            raise ValueError("Only one of unsupported or supported must be provided")
+        if len(keys) == 0:
+            raise ValueError("keys must be provided and non-empty when subclassing GatedFieldMixin")
+
+    @classmethod
+    def _create_gated_field_validator(
+        cls,
+        field_name: str,
+        default_if_supported: object | None,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ):
+        """Create the model validator function."""
+
+        @model_validator(mode="after")
+        def gated_field_validator(self):
+            """Validate the gated field."""
+            current_value = getattr(self, field_name, None)
+            is_supported = cls._check_field_support(self, unsupported, supported, keys)
+            if not is_supported:
+                if current_value is not None:
+                    blocking_key = cls._find_blocking_key(self, unsupported, supported, keys)
+                    value = getattr(self, blocking_key, "<unknown>")
+                    raise ValueError(f"{field_name} is not supported for {blocking_key}: {value}")
+            elif current_value is None:
+                setattr(self, field_name, default_if_supported)
+            return self
+
+        return gated_field_validator
+
+    @classmethod
+    def _check_field_support(
+        cls,
+        instance: object,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> bool:
+        """Check if a specific field is supported based on its configuration and keys."""
+        seen = False
+        for key in keys:
+            if not hasattr(instance, key):
+                continue
+            seen = True
+            value = str(getattr(instance, key))
+            if supported is not None:
+                if any(p.search(value) for p in supported):
+                    return True
+            elif unsupported is not None:
+                if any(p.search(value) for p in unsupported):
+                    return False
+        return True if not seen else (unsupported is not None)
+
+    @classmethod
+    def _find_blocking_key(
+        cls,
+        instance: object,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> str:
+        """Find which key is blocking the field."""
+        for key in keys:
+            if not hasattr(instance, key):
+                continue
+            value = str(getattr(instance, key))
+            if supported is not None:
+                if not any(p.search(value) for p in supported):
+                    return key
+            elif unsupported is not None:
+                if any(p.search(value) for p in unsupported):
+                    return key
+
+        return "<unknown>"
+
+    @classmethod
+    def _collect_all_mixin_configs(cls) -> list[GatedFieldMixinConfig]:
+        """Collect all mixin configurations from base classes."""
+        all_mixins = []
+        for base in cls.__bases__:
+            if hasattr(base, "_gated_field_mixins"):
+                all_mixins.extend(base._gated_field_mixins)
+        return all_mixins
+
+    @classmethod
+    def _create_combined_validator(cls, all_mixins: list[GatedFieldMixinConfig]) -> None:
+        """Create a combined validator that handles all fields."""
+
+        @model_validator(mode="after")
+        def combined_gated_field_validator(self):
+            """Validate all gated fields."""
+            for mixin_config in all_mixins:
+                field_name_local = mixin_config.field_name
+                current_value = getattr(self, field_name_local, None)
+                if not self._check_field_support_instance(mixin_config):
+                    if current_value is not None:
+                        blocking_key = self._find_blocking_key_instance(mixin_config)
+                        value = getattr(self, blocking_key, "<unknown>")
+                        raise ValueError(f"{field_name_local} is not supported for {blocking_key}: {value}")
+                elif current_value is None:
+                    setattr(self, field_name_local, mixin_config.default_if_supported)
+
+            return self
+
+        cls._combined_gated_field_validator = combined_gated_field_validator
+
+        # Add helper methods
+        def _check_field_support_instance(self, mixin_config: GatedFieldMixinConfig) -> bool:
+            """Check if a specific field is supported based on its configuration and keys."""
+            return cls._check_field_support(self, mixin_config.unsupported, mixin_config.supported, mixin_config.keys)
+
+        def _find_blocking_key_instance(self, mixin_config: GatedFieldMixinConfig) -> str:
+            """Find which key is blocking the field."""
+            return cls._find_blocking_key(self, mixin_config.unsupported, mixin_config.supported, mixin_config.keys)
+
+        cls._check_field_support_instance = _check_field_support_instance
+        cls._find_blocking_key_instance = _find_blocking_key_instance

nat/data_models/intermediate_step.py

@@ -142,7 +142,7 @@ class IntermediateStepPayload(BaseModel):
     UUID: str = Field(default_factory=lambda: str(uuid.uuid4()))
 
     @property
-    def event_category(self) -> IntermediateStepCategory:  # pylint: disable=too-many-return-statements
+    def event_category(self) -> IntermediateStepCategory:
         match self.event_type:
             case IntermediateStepType.LLM_START:
                 return IntermediateStepCategory.LLM
@@ -180,7 +180,7 @@ class IntermediateStepPayload(BaseModel):
                 raise ValueError(f"Unknown event type: {self.event_type}")
 
     @property
-    def event_state(self) -> IntermediateStepState:  # pylint: disable=too-many-return-statements
+    def event_state(self) -> IntermediateStepState:
         match self.event_type:
             case IntermediateStepType.LLM_START:
                 return IntermediateStepState.START
@@ -290,7 +290,7 @@ class IntermediateStep(BaseModel):
         return self.payload.usage_info
 
     @property
-    def UUID(self) -> str:  # pylint: disable=invalid-name
+    def UUID(self) -> str:
         return self.payload.UUID
 
     @property

nat/data_models/optimizable.py

@@ -0,0 +1,119 @@
+# 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 Sequence
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from optuna import Trial
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import model_validator
+
+T = TypeVar("T", int, float, bool, str)
+
+
+# --------------------------------------------------------------------- #
+# 1.  Hyper‑parameter metadata container                                #
+# --------------------------------------------------------------------- #
+class SearchSpace(BaseModel, Generic[T]):
+    values: Sequence[T] | None = None
+    low: T | None = None
+    high: T | None = None
+    log: bool = False  # log scale
+    step: float | None = None
+    is_prompt: bool = False
+    prompt: str | None = None  # prompt to optimize
+    prompt_purpose: str | None = None  # purpose of the prompt
+
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
+
+    @model_validator(mode="after")
+    def validate_search_space_parameters(self):
+        """Validate that either values is provided, or both high and low."""
+        if self.values is not None:
+            # If values is provided, we don't need high/low
+            if self.high is not None or self.low is not None:
+                raise ValueError("SearchSpace 'values' is mutually exclusive with 'high' and 'low'")
+            return self
+
+        return self
+
+    # Helper for Optuna Trials
+    def suggest(self, trial: Trial, name: str):
+        if self.is_prompt:
+            raise ValueError("Prompt optimization not currently supported using Optuna. "
+                             "Use the genetic algorithm implementation instead.")
+        if self.values is not None:
+            return trial.suggest_categorical(name, self.values)
+        if isinstance(self.low, int):
+            return trial.suggest_int(name, self.low, self.high, log=self.log, step=self.step)
+        return trial.suggest_float(name, self.low, self.high, log=self.log, step=self.step)
+
+
+def OptimizableField(
+    default: Any,
+    *,
+    space: SearchSpace | None = None,
+    merge_conflict: str = "overwrite",
+    **fld_kw,
+):
+    # 1. Pull out any user‑supplied extras (must be a dict)
+    user_extra = fld_kw.pop("json_schema_extra", None) or {}
+    if not isinstance(user_extra, dict):
+        raise TypeError("`json_schema_extra` must be a mapping.")
+
+    # 2. If the space is a prompt, ensure a concrete base prompt exists
+    if space is not None and getattr(space, "is_prompt", False):
+        if getattr(space, "prompt", None) is None:
+            if default is None:
+                raise ValueError("Prompt-optimized fields require a base prompt: provide a "
+                                 "non-None field default or set space.prompt.")
+            # Default prompt not provided in space; fall back to the field's default
+            space.prompt = default
+
+    # 3. Prepare our own metadata
+    ours = {"optimizable": True}
+    if space is not None:
+        ours["search_space"] = space
+
+    # 4. Merge with user extras according to merge_conflict policy
+    intersect = ours.keys() & user_extra.keys()
+    if intersect:
+        if merge_conflict == "error":
+            raise ValueError("`json_schema_extra` already contains reserved key(s): "
+                             f"{', '.join(intersect)}")
+        if merge_conflict == "keep":
+            # remove the ones the user already set so we don't overwrite them
+            ours = {k: v for k, v in ours.items() if k not in intersect}
+
+    merged_extra = {**user_extra, **ours}  # ours wins if 'overwrite'
+
+    # 5. Return a normal Pydantic Field with merged extras
+    return Field(default, json_schema_extra=merged_extra, **fld_kw)
+
+
+class OptimizableMixin(BaseModel):
+    optimizable_params: list[str] = Field(default_factory=list,
+                                          description="List of parameters that can be optimized.",
+                                          exclude=True)
+
+    search_space: dict[str, SearchSpace] = Field(
+        default_factory=dict,
+        description="Optional search space overrides for optimizable parameters.",
+        exclude=True,
+    )