nvidia-nat 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

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,
+    )

nat/data_models/optimizer.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-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 pathlib import Path
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class OptimizerMetric(BaseModel):
+    """
+    Parameters used by the workflow optimizer to define a metric to optimize.
+    """
+    evaluator_name: str = Field(description="Name of the metric to optimize.")
+    direction: str = Field(description="Direction of the optimization. Can be 'maximize' or 'minimize'.")
+    weight: float = Field(description="Weight of the metric in the optimization process.", default=1.0)
+
+
+class NumericOptimizationConfig(BaseModel):
+    """
+    Configuration for numeric/enum optimization (Optuna).
+    """
+    enabled: bool = Field(default=True, description="Enable numeric optimization")
+    n_trials: int = Field(description="Number of trials for numeric optimization.", default=20)
+
+
+class PromptGAOptimizationConfig(BaseModel):
+    """
+    Configuration for prompt optimization using a Genetic Algorithm.
+    """
+    enabled: bool = Field(default=False, description="Enable GA-based prompt optimization")
+
+    # Prompt optimization function hooks
+    prompt_population_init_function: str | None = Field(
+        default=None,
+        description="Optional function name to initialize/mutate candidate prompts.",
+    )
+    prompt_recombination_function: str | None = Field(
+        default=None,
+        description="Optional function name to recombine two parent prompts into a child.",
+    )
+
+    # Genetic algorithm configuration
+    ga_population_size: int = Field(
+        description="Population size for genetic algorithm prompt optimization.",
+        default=24,
+    )
+    ga_generations: int = Field(
+        description="Number of generations to evolve in GA prompt optimization.",
+        default=15,
+    )
+    ga_offspring_size: int | None = Field(
+        description="Number of offspring to produce per generation. Defaults to population_size - elitism.",
+        default=None,
+    )
+    ga_crossover_rate: float = Field(
+        description="Probability of applying crossover during reproduction.",
+        default=0.8,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_mutation_rate: float = Field(
+        description="Probability of mutating a child after crossover.",
+        default=0.3,
+        ge=0.0,
+        le=1.0,
+    )
+    ga_elitism: int = Field(
+        description="Number of top individuals carried over unchanged each generation.",
+        default=2,
+    )
+    ga_selection_method: str = Field(
+        description="Parent selection strategy: 'tournament' or 'roulette'.",
+        default="tournament",
+    )
+    ga_tournament_size: int = Field(
+        description="Tournament size when using tournament selection.",
+        default=3,
+    )
+    ga_parallel_evaluations: int = Field(
+        description="Max number of individuals to evaluate concurrently per generation.",
+        default=8,
+    )
+    ga_diversity_lambda: float = Field(
+        description="Strength of diversity penalty (0 disables). Penalizes identical/near-identical prompts.",
+        default=0.0,
+        ge=0.0,
+    )
+
+
+class OptimizerConfig(BaseModel):
+    """
+    Parameters used by the workflow optimizer.
+    """
+    output_path: Path | None = Field(
+        default=None,
+        description="Path to the output directory where the results will be saved.",
+    )
+
+    eval_metrics: dict[str, OptimizerMetric] | None = Field(
+        description="List of evaluation metrics to optimize.",
+        default=None,
+    )
+
+    reps_per_param_set: int = Field(
+        default=3,
+        description="Number of repetitions per parameter set for the optimization.",
+    )
+
+    target: float | None = Field(
+        description=(
+            "Target value for the optimization. If set, the optimization will stop when this value is reached."),
+        default=None,
+    )
+
+    multi_objective_combination_mode: str = Field(
+        description="Method to combine multiple objectives into a single score.",
+        default="harmonic",
+    )
+
+    # Nested configs
+    numeric: NumericOptimizationConfig = NumericOptimizationConfig()
+    prompt: PromptGAOptimizationConfig = PromptGAOptimizationConfig()
+
+
+class OptimizerRunConfig(BaseModel):
+    """
+    Parameters used for an Optimizer R=run
+    """
+    # Eval parameters
+
+    config_file: Path | BaseModel  # allow for instantiated configs to be passed in
+    dataset: str | Path | None  # dataset file path can be specified in the config file
+    result_json_path: str = "$"
+    endpoint: str | None = None  # only used when running the workflow remotely
+    endpoint_timeout: int = 300
+    override: tuple[tuple[str, str], ...] = ()

nat/data_models/span.py

@@ -128,10 +128,48 @@ class SpanStatus(BaseModel):
     message: str | None = Field(default=None, description="The status message of the span.")
 
 
+def _generate_nonzero_trace_id() -> int:
+    """Generate a non-zero 128-bit trace ID."""
+    return uuid.uuid4().int
+
+
+def _generate_nonzero_span_id() -> int:
+    """Generate a non-zero 64-bit span ID."""
+    return uuid.uuid4().int >> 64
+
+
 class SpanContext(BaseModel):
-    trace_id: int = Field(default_factory=lambda: uuid.uuid4().int, description="The 128-bit trace ID of the span.")
-    span_id: int = Field(default_factory=lambda: uuid.uuid4().int & ((1 << 64) - 1),
-                         description="The 64-bit span ID of the span.")
+    trace_id: int = Field(default_factory=_generate_nonzero_trace_id,
+                          description="The OTel-syle 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=_generate_nonzero_span_id,
+                         description="The OTel-syle 64-bit span ID of the span.")
+
+    @field_validator("trace_id", mode="before")
+    @classmethod
+    def _validate_trace_id(cls, v: int | str | None) -> int:
+        """Regenerate if trace_id is None; raise an exception if trace_id is invalid;"""
+        if isinstance(v, str):
+            v = uuid.UUID(v).int
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_trace_id()
+        if v <= 0 or v >> 128:
+            raise ValueError(f"Invalid trace_id: must be a non-zero 128-bit integer, got {v}")
+        return v
+
+    @field_validator("span_id", mode="before")
+    @classmethod
+    def _validate_span_id(cls, v: int | str | None) -> int:
+        """Regenerate if span_id is None; raise an exception if span_id is invalid;"""
+        if isinstance(v, str):
+            try:
+                v = int(v, 16)
+            except ValueError:
+                raise ValueError(f"span_id unable to be parsed: {v}")
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_span_id()
+        if v <= 0 or v >> 64:
+            raise ValueError(f"Invalid span_id: must be a non-zero 64-bit integer, got {v}")
+        return v
 
 
 class Span(BaseModel):

nat/data_models/swe_bench_model.py

@@ -39,7 +39,7 @@ class SWEBenchInput(BaseModel):
 
     # Handle improperly formatted JSON strings for list fields
     @field_validator("FAIL_TO_PASS", "PASS_TO_PASS", mode="before")
-    def parse_list_fields(cls, value):  # pylint: disable=no-self-argument
+    def parse_list_fields(cls, value):
         if isinstance(value, str):
             # Attempt to parse the string as a list
             return json.loads(value)

nat/data_models/temperature_mixin.py

@@ -0,0 +1,44 @@
+# 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 re
+
+from pydantic import BaseModel
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
+
+
+class TemperatureMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="temperature",
+        default_if_supported=0.0,
+        keys=("model_name", "model", "azure_deployment"),
+        unsupported=(re.compile(r"gpt-?5", re.IGNORECASE), ),
+):
+    """
+    Mixin class for temperature configuration. Unsupported on models like gpt-5.
+
+    Attributes:
+        temperature: Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.
+    """
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Sampling temperature in [0, 1]. Defaults to 0.0 when supported on the model.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))

nat/data_models/thinking_mixin.py

@@ -0,0 +1,86 @@
+# 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 re
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+
+# Currently the control logic for thinking is only implemented for Nemotron models
+_NEMOTRON_REGEX = re.compile(r"^nvidia/(llama|nvidia).*nemotron", re.IGNORECASE)
+# The keys are the fields that are used to determine if the model supports thinking
+_MODEL_KEYS = ("model_name", "model", "azure_deployment")
+
+
+class ThinkingMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="thinking",
+        default_if_supported=None,
+        keys=_MODEL_KEYS,
+        supported=(_NEMOTRON_REGEX, ),
+):
+    """
+    Mixin class for thinking configuration. Only supported on Nemotron models.
+
+    Attributes:
+        thinking: Whether to enable thinking. Defaults to None when supported on the model.
+    """
+    thinking: bool | None = Field(
+        default=None,
+        description="Whether to enable thinking. Defaults to None when supported on the model.",
+    )
+
+    @property
+    def thinking_system_prompt(self) -> str | None:
+        """
+        Returns the system prompt to use for thinking.
+        For NVIDIA Nemotron, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron v1.5, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron v1.0 or v1.1, returns "detailed thinking on" if enabled, else "detailed thinking off".
+        If thinking is not supported on the model, returns None.
+
+        Returns:
+            str | None: The system prompt to use for thinking.
+        """
+        if self.thinking is None:
+            return None
+
+        for key in _MODEL_KEYS:
+            model = getattr(self, key, None)
+            if not isinstance(model, str) or model is None:
+                continue
+
+            # Normalize name to reduce checks
+            model = model.lower().translate(str.maketrans("_.", "--"))
+
+            if model.startswith("nvidia/nvidia"):
+                return "/think" if self.thinking else "/no_think"
+
+            if model.startswith("nvidia/llama"):
+                if "v1-0" in model or "v1-1" in model or model.endswith("v1"):
+                    return f"detailed thinking {'on' if self.thinking else 'off'}"
+
+                if "v1-5" in model:
+                    # v1.5 models are updated to use the /think and /no_think system prompts
+                    return "/think" if self.thinking else "/no_think"
+
+                # Assume any other model is a newer model that uses the /think and /no_think system prompts
+                return "/think" if self.thinking else "/no_think"
+
+        # Unknown model
+        return None

nat/data_models/top_p_mixin.py

@@ -0,0 +1,44 @@
+# 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 re
+
+from pydantic import BaseModel
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
+
+
+class TopPMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="top_p",
+        default_if_supported=1.0,
+        keys=("model_name", "model", "azure_deployment"),
+        unsupported=(re.compile(r"gpt-?5", re.IGNORECASE), ),
+):
+    """
+    Mixin class for top-p configuration. Unsupported on models like gpt-5.
+
+    Attributes:
+        top_p: Top-p for distribution sampling. Defaults to 1.0 when supported on the model.
+    """
+    top_p: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Top-p for distribution sampling. Defaults to 1.0 when supported on the model.",
+        space=SearchSpace(high=1.0, low=0.5, step=0.1))

nat/embedder/azure_openai_embedder.py

@@ -0,0 +1,46 @@
+# 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 pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.embedder import EmbedderProviderInfo
+from nat.cli.register_workflow import register_embedder_provider
+from nat.data_models.embedder import EmbedderBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+
+
+class AzureOpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="azure_openai"):
+    """An Azure OpenAI embedder provider to be used with an embedder client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="Azure OpenAI API key to interact with hosted model.")
+    api_version: str = Field(default="2025-04-01-preview", description="Azure OpenAI API version.")
+    azure_endpoint: str | None = Field(validation_alias=AliasChoices("azure_endpoint", "base_url"),
+                                       serialization_alias="azure_endpoint",
+                                       default=None,
+                                       description="Base URL for the hosted model.")
+    azure_deployment: str = Field(validation_alias=AliasChoices("azure_deployment", "model_name", "model"),
+                                  serialization_alias="azure_deployment",
+                                  description="The Azure OpenAI hosted model/deployment name.")
+
+
+@register_embedder_provider(config_type=AzureOpenAIEmbedderModelConfig)
+async def azure_openai_embedder_model(config: AzureOpenAIEmbedderModelConfig, _builder: Builder):
+
+    yield EmbedderProviderInfo(config=config, description="An Azure OpenAI model for use with an Embedder client.")

nat/embedder/nim_embedder.py

@@ -50,7 +50,7 @@ class NIMEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="nim"):
                                        description=("The truncation strategy if the input on the "
                                                     "server side if it's too large."))
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
 
 @register_embedder_provider(config_type=NIMEmbedderModelConfig)

nat/embedder/openai_embedder.py

@@ -27,17 +27,16 @@ from nat.data_models.retry_mixin import RetryMixin
 class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
     """An OpenAI 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="OpenAI API key to interact with hosted model.")
     base_url: str | None = Field(default=None, description="Base url to the hosted model.")
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The OpenAI hosted model name.")
-    max_retries: int = Field(default=2, description="The max number of retries for the request.")
 
 
 @register_embedder_provider(config_type=OpenAIEmbedderModelConfig)
-async def openai_llm(config: OpenAIEmbedderModelConfig, builder: Builder):
+async def openai_embedder_model(config: OpenAIEmbedderModelConfig, _builder: Builder):
 
     yield EmbedderProviderInfo(config=config, description="An OpenAI model for use with an Embedder client.")

nat/embedder/register.py

@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 
 # Import any providers which need to be automatically registered here
+from . import azure_openai_embedder
 from . import nim_embedder
 from . import openai_embedder

nat/eval/config.py

@@ -27,7 +27,7 @@ class EvaluationRunConfig(BaseModel):
     """
     Parameters used for a single evaluation run.
     """
-    config_file: Path
+    config_file: Path | BaseModel
     dataset: str | None = None  # dataset file path can be specified in the config file
     result_json_path: str = "$"
     skip_workflow: bool = False
@@ -44,6 +44,8 @@ class EvaluationRunConfig(BaseModel):
     # number of passes at each concurrency, if 0 the dataset is adjusted to a multiple of the
     # concurrency. The is only used if adjust_dataset_size is true
     num_passes: int = 0
+    # timeout for waiting for trace export tasks to complete
+    export_timeout: float = 60.0
 
 
 class EvaluationRunOutput(BaseModel):