nvidia-nat 1.3a20250819__py3-none-any.whl → 1.3.0a20250823__py3-none-any.whl

nat/builder/function_info.py

@@ -231,7 +231,7 @@ class FunctionDescriptor:
         else:
             annotations = [param.annotation for param in sig.parameters.values()]
 
-            is_input_typed = all([a != sig.empty for a in annotations])  # pylint: disable=use-a-generator
+            is_input_typed = all([a != sig.empty for a in annotations])
 
             input_type = tuple[*annotations] if is_input_typed else None  # noqa: syntax-error
 
@@ -372,18 +372,16 @@ class FunctionInfo:
 
             if (stream_to_single_fn is not None):
                 raise ValueError("Cannot provide both single_fn and stream_to_single_fn")
-        else:
-            if (stream_to_single_fn is not None and stream_fn is None):
-                raise ValueError("stream_fn must be provided if stream_to_single_fn is provided")
+        elif (stream_to_single_fn is not None and stream_fn is None):
+            raise ValueError("stream_fn must be provided if stream_to_single_fn is provided")
 
         if (stream_fn is not None):
             final_stream_fn = stream_fn
 
             if (single_to_stream_fn is not None):
                 raise ValueError("Cannot provide both stream_fn and single_to_stream_fn")
-        else:
-            if (single_to_stream_fn is not None and single_fn is None):
-                raise ValueError("single_fn must be provided if single_to_stream_fn is provided")
+        elif (single_to_stream_fn is not None and single_fn is None):
+            raise ValueError("single_fn must be provided if single_to_stream_fn is provided")
 
         if (single_fn is None and stream_fn is None):
             raise ValueError("At least one of single_fn or stream_fn must be provided")

nat/builder/user_interaction_manager.py

@@ -61,13 +61,13 @@ class UserInteractionManager:
 
         uuid_req = str(uuid.uuid4())
         status = InteractionStatus.IN_PROGRESS
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
         sys_human_interaction = InteractionPrompt(id=uuid_req, status=status, timestamp=timestamp, content=content)
 
         resp = await self._context_state.user_input_callback.get()(sys_human_interaction)
 
         # Rebuild a InteractionResponse object with the response
-        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
         status = InteractionStatus.COMPLETED
         sys_human_interaction = InteractionResponse(id=uuid_req, status=status, timestamp=timestamp, content=resp)
 

nat/builder/workflow.py

@@ -83,6 +83,9 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
 
         return self._entry_fn.has_single_output
 
+    async def get_all_exporters(self) -> dict[str, BaseExporter]:
+        return await self._exporter_manager.get_all_exporters()
+
     @asynccontextmanager
     async def run(self, message: InputT):
         """

nat/builder/workflow_builder.py

@@ -127,7 +127,6 @@ class ConfiguredTTCStrategy:
     instance: StrategyBase
 
 
-# pylint: disable=too-many-public-methods
 class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
     def __init__(self, *, general_config: GeneralConfig | None = None, registry: TypeRegistry | None = None):

nat/cli/commands/evaluate.py

@@ -97,7 +97,7 @@ async def run_and_evaluate(config: EvaluationRunConfig):
 
 @eval_command.result_callback(replace=True)
 def process_nat_eval(
-    processors,  # pylint: disable=unused-argument
+    processors,
     *,
     config_file: Path,
     dataset: Path,

nat/cli/commands/info/list_components.py

@@ -26,14 +26,13 @@ from nat.registry_handlers.schemas.search import SearchFields
 logger = logging.getLogger(__name__)
 
 
-async def search_artifacts(  # pylint: disable=R0917
-        registry_handler_config: RegistryHandlerBaseConfig,
-        component_types: list[ComponentEnum],
-        visualize: bool,
-        query: str,
-        num_results: int,
-        query_fields: list[SearchFields],
-        save_path: str | None) -> None:
+async def search_artifacts(registry_handler_config: RegistryHandlerBaseConfig,
+                           component_types: list[ComponentEnum],
+                           visualize: bool,
+                           query: str,
+                           num_results: int,
+                           query_fields: list[SearchFields],
+                           save_path: str | None) -> None:
 
     from nat.cli.type_registry import GlobalTypeRegistry
     from nat.registry_handlers.schemas.search import SearchQuery

nat/cli/commands/info/list_mcp.py

@@ -297,8 +297,7 @@ def ping(url: str, timeout: int, json_output: bool) -> None:
 
     if json_output:
         click.echo(result.model_dump_json(indent=2))
+    elif result.status == "healthy":
+        click.echo(f"Server at {result.url} is healthy (response time: {result.response_time_ms}ms)")
     else:
-        if result.status == "healthy":
-            click.echo(f"Server at {result.url} is healthy (response time: {result.response_time_ms}ms)")
-        else:
-            click.echo(f"Server at {result.url} {result.status}: {result.error}")
+        click.echo(f"Server at {result.url} {result.status}: {result.error}")

nat/cli/commands/registry/search.py

@@ -29,14 +29,13 @@ from nat.utils.data_models.schema_validator import validate_yaml
 logger = logging.getLogger(__name__)
 
 
-async def search_artifacts(  # pylint: disable=R0917
-        registry_handler_config: RegistryHandlerBaseConfig,
-        query: str,
-        search_fields: list[SearchFields],
-        visualize: bool,
-        component_types: list[ComponentEnum],
-        save_path: str | None = None,
-        n_results: int = 10) -> None:
+async def search_artifacts(registry_handler_config: RegistryHandlerBaseConfig,
+                           query: str,
+                           search_fields: list[SearchFields],
+                           visualize: bool,
+                           component_types: list[ComponentEnum],
+                           save_path: str | None = None,
+                           n_results: int = 10) -> None:
 
     from nat.cli.type_registry import GlobalTypeRegistry
     from nat.registry_handlers.schemas.search import SearchQuery
@@ -116,14 +115,13 @@ async def search_artifacts(  # pylint: disable=R0917
     required=False,
     help=("The component types to include in search."),
 )
-def search(  # pylint: disable=R0917
-        config_file: str,
-        channel: str,
-        fields: list[str],
-        query: str,
-        component_types: list[ComponentEnum],
-        n_results: int,
-        output_path: str) -> None:
+def search(config_file: str,
+           channel: str,
+           fields: list[str],
+           query: str,
+           component_types: list[ComponentEnum],
+           n_results: int,
+           output_path: str) -> None:
     """
     Search for NAT artifacts with the specified configuration.
     """

nat/cli/commands/start.py

@@ -35,7 +35,6 @@ logger = logging.getLogger(__name__)
 
 class StartCommandGroup(click.Group):
 
-    # pylint: disable=too-many-positional-arguments
     def __init__(
         self,
         name: str | None = None,

nat/cli/commands/workflow/templates/pyproject.toml.j2

@@ -3,6 +3,9 @@ build-backend = "setuptools.build_meta"
 {% if editable %}requires = ["setuptools >= 64", "setuptools-scm>=8"]
 
 [tool.setuptools_scm]
+# NAT uses the --first-parent flag to avoid tags from previous releases which have been merged into the develop branch
+# from causing an unexpected version change. This can be safely removed if developing outside of the NAT repository.
+git_describe_command = "git describe --long --first-parent"
 root = "{{ rel_path_to_repo_root}}"{% else %}requires = ["setuptools >= 64"]{% endif %}
 
 [project]

nat/cli/commands/workflow/templates/register.py.j2

@@ -1,4 +1,3 @@
-# pylint: disable=unused-import
 # flake8: noqa
 
 # Import any tools which need to be automatically registered here

nat/cli/commands/workflow/workflow_commands.py

@@ -161,7 +161,6 @@ def get_workflow_path_from_name(workflow_name: str):
     default="NAT function template. Please update the description.",
     help="""A description of the component being created. Will be used to populate the docstring and will describe the
          component when inspecting installed components using 'nat info component'""")
-# pylint: disable=missing-param-doc
 def create_command(workflow_name: str, install: bool, workflow_dir: str, description: str):
     """
     Create a new NAT workflow using templates.

nat/cli/type_registry.py

@@ -298,7 +298,7 @@ class RegisteredPackage(BaseModel):
     discovery_metadata: DiscoveryMetadata
 
 
-class TypeRegistry:  # pylint: disable=too-many-public-methods
+class TypeRegistry:
 
     def __init__(self) -> None:
         # Telemetry Exporters
@@ -588,8 +588,8 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         except KeyError as err:
             raise KeyError(
                 f"An invalid Embedder config and wrapper combination was supplied. Config: `{config_type}`, "
-                "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Embedder client but "
-                "there is no registered conversion from that Embedder provider to LLM framework: {wrapper_type}. "
+                f"Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Embedder client but "
+                f"there is no registered conversion from that Embedder provider to LLM framework: {wrapper_type}. "
                 "Please provide an Embedder configuration from one of the following providers: "
                 f"{set(self._embedder_client_provider_to_framework.keys())}") from err
 
@@ -703,8 +703,8 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
         except KeyError as err:
             raise KeyError(
                 f"An invalid Retriever config and wrapper combination was supplied. Config: `{config_type}`, "
-                "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Retriever client but "
-                "there is no registered conversion from that Retriever provider to LLM framework: {wrapper_type}. "
+                f"Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Retriever client but "
+                f"there is no registered conversion from that Retriever provider to LLM framework: {wrapper_type}. "
                 "Please provide a Retriever configuration from one of the following providers: "
                 f"{set(self._retriever_client_provider_to_framework.keys())}") from err
 
@@ -779,7 +779,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         self._registration_changed()
 
-    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:  # pylint: disable=R0911
+    def get_infos_by_type(self, component_type: ComponentEnum) -> dict:
 
         if component_type == ComponentEnum.FRONT_END:
             return self._registered_front_end_infos
@@ -849,8 +849,7 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
 
         raise ValueError(f"Supplied an unsupported component type {component_type}")
 
-    def get_registered_types_by_component_type(  # pylint: disable=R0911
-            self, component_type: ComponentEnum) -> list[str]:
+    def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]:
 
         if component_type == ComponentEnum.FUNCTION:
             return [i.static_type() for i in self._registered_functions]
@@ -925,7 +924,6 @@ class TypeRegistry:  # pylint: disable=too-many-public-methods
             if (short_names[key.local_name] == 1):
                 type_list.append((key.local_name, key.config_type))
 
-        # pylint: disable=consider-alternative-union-syntax
         return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
     def compute_annotation(self, cls: type[TypedBaseModelT]):

nat/data_models/config.py

@@ -47,7 +47,7 @@ logger = logging.getLogger(__name__)
 
 
 def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-    from nat.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+    from nat.cli.type_registry import GlobalTypeRegistry
 
     new_errors = []
     logged_once = False

nat/data_models/evaluate.py

@@ -108,7 +108,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_dependencies.py

@@ -54,19 +54,19 @@ 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_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/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/model_gated_field_mixin.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.
+
+from collections.abc import Sequence
+from re import Pattern
+from typing import Generic
+from typing import TypeVar
+
+from pydantic import BaseModel
+from pydantic import model_validator
+
+T = TypeVar("T")
+
+
+class ModelGatedFieldMixin(Generic[T]):
+    """
+    A mixin that gates a field based on model support.
+
+    This should be used to automatically validate a field based on a given model.
+
+    Parameters
+    ----------
+    field_name: `str`
+                The name of the field.
+    default_if_supported: `T`
+                          The default value of the field if it is supported for the model.
+    unsupported_models: `Sequence[Pattern[str]] | None`
+                        A sequence of regex patterns that match the model names NOT supported for the field.
+                        Defaults to None.
+    supported_models: `Sequence[Pattern[str]] | None`
+                      A sequence of regex patterns that match the model names supported for the field.
+                      Defaults to None.
+    model_keys: `Sequence[str]`
+                A sequence of keys that are used to validate the field.
+                Defaults to ("model_name", "model", "azure_deployment",)
+    """
+
+    def __init_subclass__(
+            cls,
+            field_name: str | None = None,
+            default_if_supported: T | None = None,
+            unsupported_models: Sequence[Pattern[str]] | None = None,
+            supported_models: Sequence[Pattern[str]] | None = None,
+            model_keys: Sequence[str] = ("model_name", "model", "azure_deployment"),
+    ) -> None:
+        """
+        Store the class variables for the field and define the model validator.
+        """
+        super().__init_subclass__()
+        if ModelGatedFieldMixin in cls.__bases__:
+            if field_name is None:
+                raise ValueError("field_name must be provided when subclassing ModelGatedFieldMixin")
+            if default_if_supported is None:
+                raise ValueError("default_if_supported must be provided when subclassing ModelGatedFieldMixin")
+            if unsupported_models is None and supported_models is None:
+                raise ValueError("Either unsupported_models or supported_models must be provided")
+            if unsupported_models is not None and supported_models is not None:
+                raise ValueError("Only one of unsupported_models or supported_models must be provided")
+            if model_keys is not None and len(model_keys) == 0:
+                raise ValueError("model_keys must be provided and non-empty when subclassing ModelGatedFieldMixin")
+            cls.field_name = field_name
+            cls.default_if_supported = default_if_supported
+            cls.unsupported_models = unsupported_models
+            cls.supported_models = supported_models
+            if model_keys is not None:
+                cls.model_keys = model_keys
+
+            @classmethod
+            def check_model(cls, model_name: str) -> bool:
+                """
+                Check if a model is supported for a given field.
+
+                Args:
+                    model_name: The name of the model to check.
+                """
+                unsupported = getattr(cls, "unsupported_models", None)
+                supported = getattr(cls, "supported_models", None)
+                if unsupported is not None:
+                    return not any(p.search(model_name) for p in unsupported)
+                if supported is not None:
+                    return any(p.search(model_name) for p in supported)
+                return False
+
+            cls._model_gated_field_check_model = check_model
+
+            @classmethod
+            def detect_support(cls, instance: BaseModel) -> str | None:
+                for key in getattr(cls, "model_keys"):
+                    if hasattr(instance, key):
+                        model_name_value = getattr(instance, key)
+                        is_supported = getattr(cls, "_model_gated_field_check_model")(str(model_name_value))
+                        return key if not is_supported else None
+                return None
+
+            cls._model_gated_field_detect_support = detect_support
+
+            @model_validator(mode="after")
+            def model_validate(self):
+                klass = self.__class__
+
+                field_name_local = getattr(klass, "field_name")
+                current_value = getattr(self, field_name_local, None)
+
+                found_key = klass._model_gated_field_detect_support(self)
+                if found_key is not None:
+                    if current_value is not None:
+                        raise ValueError(
+                            f"{field_name_local} is not supported for {found_key}: {getattr(self, found_key)}")
+                elif current_value is None:
+                    setattr(self, field_name_local, getattr(klass, "default_if_supported", None))
+                return self
+
+            cls._model_gated_field_model_validator = model_validate

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,36 @@
+# 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.model_gated_field_mixin import ModelGatedFieldMixin
+
+_UNSUPPORTED_TEMPERATURE_MODELS = (re.compile(r"gpt-?5", re.IGNORECASE), )
+
+
+class TemperatureMixin(
+        BaseModel,
+        ModelGatedFieldMixin[float],
+        field_name="temperature",
+        default_if_supported=0.0,
+        unsupported_models=_UNSUPPORTED_TEMPERATURE_MODELS,
+):
+    """
+    Mixin class for temperature configuration.
+    """
+    temperature: float | None = Field(default=None, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")

nat/data_models/top_p_mixin.py

@@ -0,0 +1,36 @@
+# 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.model_gated_field_mixin import ModelGatedFieldMixin
+
+_UNSUPPORTED_TOP_P_MODELS = (re.compile(r"gpt-?5", re.IGNORECASE), )
+
+
+class TopPMixin(
+        BaseModel,
+        ModelGatedFieldMixin[float],
+        field_name="top_p",
+        default_if_supported=1.0,
+        unsupported_models=_UNSUPPORTED_TOP_P_MODELS,
+):
+    """
+    Mixin class for top-p configuration.
+    """
+    top_p: float | None = Field(default=None, ge=0.0, le=1.0, description="Top-p for distribution sampling.")

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/openai_embedder.py

@@ -34,10 +34,9 @@ class OpenAIEmbedderModelConfig(EmbedderBaseConfig, RetryMixin, name="openai"):
     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

@@ -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):

nat/eval/dataset_handler/dataset_handler.py

@@ -146,13 +146,12 @@ class DatasetHandler:
             # When num_passes is specified, always use concurrency * num_passes
             # This respects the user's intent for exact number of passes
             target_size = self.concurrency * self.num_passes
+        # When num_passes = 0, use the largest multiple of concurrency <= original_size
+        # If original_size < concurrency, we need at least concurrency rows
+        elif original_size >= self.concurrency:
+            target_size = (original_size // self.concurrency) * self.concurrency
         else:
-            # When num_passes = 0, use the largest multiple of concurrency <= original_size
-            # If original_size < concurrency, we need at least concurrency rows
-            if original_size >= self.concurrency:
-                target_size = (original_size // self.concurrency) * self.concurrency
-            else:
-                target_size = self.concurrency
+            target_size = self.concurrency
 
         if target_size == 0:
             raise ValueError("Input dataset too small for even one batch at given concurrency.")