nvidia-nat 1.2.0__py3-none-any.whl

nat/data_models/config.py

@@ -0,0 +1,410 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import sys
+import typing
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Discriminator
+from pydantic import ValidationError
+from pydantic import ValidationInfo
+from pydantic import ValidatorFunctionWrapHandler
+from pydantic import field_validator
+
+from nat.data_models.evaluate import EvalConfig
+from nat.data_models.front_end import FrontEndBaseConfig
+from nat.data_models.function import EmptyFunctionConfig
+from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.logging import LoggingBaseConfig
+from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
+
+from .authentication import AuthProviderBaseConfig
+from .common import HashableBaseModel
+from .common import TypedBaseModel
+from .embedder import EmbedderBaseConfig
+from .llm import LLMBaseConfig
+from .memory import MemoryBaseConfig
+from .object_store import ObjectStoreBaseConfig
+from .retriever import RetrieverBaseConfig
+
+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
+
+    new_errors = []
+    logged_once = False
+    needs_reraise = False
+    for e in err.errors():
+
+        error_type = e['type']
+        if error_type == 'union_tag_invalid' and "ctx" in e and not logged_once:
+            requested_type = e["ctx"]["tag"]
+
+            if (info.field_name in ('workflow', 'functions')):
+                registered_keys = GlobalTypeRegistry.get().get_registered_functions()
+            elif (info.field_name == "authentication"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_auth_providers()
+            elif (info.field_name == "llms"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_llm_providers()
+            elif (info.field_name == "embedders"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_embedder_providers()
+            elif (info.field_name == "memory"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_memorys()
+            elif (info.field_name == "object_stores"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_object_stores()
+            elif (info.field_name == "retrievers"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_retriever_providers()
+            elif (info.field_name == "tracing"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_telemetry_exporters()
+            elif (info.field_name == "logging"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_logging_method()
+            elif (info.field_name == "evaluators"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_evaluators()
+            elif (info.field_name == "front_ends"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_front_ends()
+            elif (info.field_name == "ttc_strategies"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_ttc_strategies()
+
+            else:
+                assert False, f"Unknown field name {info.field_name} in validator"
+
+            # Check and see if the there are multiple full types which match this short type
+            matching_keys = [k for k in registered_keys if k.local_name == requested_type]
+
+            assert len(matching_keys) != 1, "Exact match should have been found. Contact developers"
+
+            matching_key_names = [x.full_type for x in matching_keys]
+            registered_key_names = [x.full_type for x in registered_keys]
+
+            if (len(matching_keys) == 0):
+                # This is a case where the requested type is not found. Show a helpful message about what is
+                # available
+                logger.error(("Requested %s type `%s` not found. "
+                              "Have you ensured the necessary package has been installed with `uv pip install`?"
+                              "\nAvailable %s names:\n - %s\n"),
+                             info.field_name,
+                             requested_type,
+                             info.field_name,
+                             '\n - '.join(registered_key_names))
+            else:
+                # This is a case where the requested type is ambiguous.
+                logger.error(("Requested %s type `%s` is ambiguous. "
+                              "Matched multiple %s by their local name: %s. "
+                              "Please use the fully qualified %s name."
+                              "\nAvailable %s names:\n - %s\n"),
+                             info.field_name,
+                             requested_type,
+                             info.field_name,
+                             matching_key_names,
+                             info.field_name,
+                             info.field_name,
+                             '\n - '.join(registered_key_names))
+
+            # Only show one error
+            logged_once = True
+
+        elif error_type == 'missing':
+            location = e["loc"]
+            if len(location) > 1:  # remove the _type field from the location
+                e['loc'] = (location[0], ) + location[2:]
+                needs_reraise = True
+
+        new_errors.append(e)
+
+    if needs_reraise:
+        raise ValidationError.from_exception_data(title=err.title, line_errors=new_errors)
+
+
+class TelemetryConfig(BaseModel):
+
+    logging: dict[str, LoggingBaseConfig] = {}
+    tracing: dict[str, TelemetryExporterBaseConfig] = {}
+
+    @field_validator("logging", "tracing", mode="wrap")
+    @classmethod
+    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
+
+        try:
+            return handler(value)
+        except ValidationError as err:
+            _process_validation_error(err, handler, info)
+            raise
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        from nat.cli.type_registry import GlobalTypeRegistry
+
+        type_registry = GlobalTypeRegistry.get()
+
+        TracingAnnotation = dict[str,
+                                 typing.Annotated[type_registry.compute_annotation(TelemetryExporterBaseConfig),
+                                                  Discriminator(TypedBaseModel.discriminator)]]
+
+        LoggingAnnotation = dict[str,
+                                 typing.Annotated[type_registry.compute_annotation(LoggingBaseConfig),
+                                                  Discriminator(TypedBaseModel.discriminator)]]
+
+        should_rebuild = False
+
+        tracing_field = cls.model_fields.get("tracing")
+        if tracing_field is not None and tracing_field.annotation != TracingAnnotation:
+            tracing_field.annotation = TracingAnnotation
+            should_rebuild = True
+
+        logging_field = cls.model_fields.get("logging")
+        if logging_field is not None and logging_field.annotation != LoggingAnnotation:
+            logging_field.annotation = LoggingAnnotation
+            should_rebuild = True
+
+        if (should_rebuild):
+            return cls.model_rebuild(force=True)
+
+        return False
+
+
+class GeneralConfig(BaseModel):
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    use_uvloop: bool = True
+    """
+    Whether to use uvloop for the event loop. This can provide a significant speedup in some cases. Disable to provide
+    better error messages when debugging.
+    """
+
+    telemetry: TelemetryConfig = TelemetryConfig()
+
+    # FrontEnd Configuration
+    front_end: FrontEndBaseConfig = FastApiFrontEndConfig()
+
+    @field_validator("front_end", mode="wrap")
+    @classmethod
+    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
+
+        try:
+            return handler(value)
+        except ValidationError as err:
+            _process_validation_error(err, handler, info)
+            raise
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        from nat.cli.type_registry import GlobalTypeRegistry
+
+        type_registry = GlobalTypeRegistry.get()
+
+        FrontEndAnnotation = typing.Annotated[type_registry.compute_annotation(FrontEndBaseConfig),
+                                              Discriminator(TypedBaseModel.discriminator)]
+
+        should_rebuild = False
+
+        front_end_field = cls.model_fields.get("front_end")
+        if front_end_field is not None and front_end_field.annotation != FrontEndAnnotation:
+            front_end_field.annotation = FrontEndAnnotation
+            should_rebuild = True
+
+        if (TelemetryConfig.rebuild_annotations()):
+            should_rebuild = True
+
+        if (should_rebuild):
+            return cls.model_rebuild(force=True)
+
+        return False
+
+
+class Config(HashableBaseModel):
+
+    model_config = ConfigDict(extra="forbid")
+
+    # Global Options
+    general: GeneralConfig = GeneralConfig()
+
+    # Functions Configuration
+    functions: dict[str, FunctionBaseConfig] = {}
+
+    # LLMs Configuration
+    llms: dict[str, LLMBaseConfig] = {}
+
+    # Embedders Configuration
+    embedders: dict[str, EmbedderBaseConfig] = {}
+
+    # Memory Configuration
+    memory: dict[str, MemoryBaseConfig] = {}
+
+    # Object Stores Configuration
+    object_stores: dict[str, ObjectStoreBaseConfig] = {}
+
+    # Retriever Configuration
+    retrievers: dict[str, RetrieverBaseConfig] = {}
+
+    # TTC Strategies
+    ttc_strategies: dict[str, TTCStrategyBaseConfig] = {}
+
+    # Workflow Configuration
+    workflow: FunctionBaseConfig = EmptyFunctionConfig()
+
+    # Authentication Configuration
+    authentication: dict[str, AuthProviderBaseConfig] = {}
+
+    # Evaluation Options
+    eval: EvalConfig = EvalConfig()
+
+    def print_summary(self, stream: typing.TextIO = sys.stdout):
+        """Print a summary of the configuration"""
+
+        stream.write("\nConfiguration Summary:\n")
+        stream.write("-" * 20 + "\n")
+        if self.workflow:
+            stream.write(f"Workflow Type: {self.workflow.type}\n")
+
+        stream.write(f"Number of Functions: {len(self.functions)}\n")
+        stream.write(f"Number of LLMs: {len(self.llms)}\n")
+        stream.write(f"Number of Embedders: {len(self.embedders)}\n")
+        stream.write(f"Number of Memory: {len(self.memory)}\n")
+        stream.write(f"Number of Object Stores: {len(self.object_stores)}\n")
+        stream.write(f"Number of Retrievers: {len(self.retrievers)}\n")
+        stream.write(f"Number of TTC Strategies: {len(self.ttc_strategies)}\n")
+        stream.write(f"Number of Authentication Providers: {len(self.authentication)}\n")
+
+    @field_validator("functions",
+                     "llms",
+                     "embedders",
+                     "memory",
+                     "retrievers",
+                     "workflow",
+                     "ttc_strategies",
+                     "authentication",
+                     mode="wrap")
+    @classmethod
+    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
+
+        try:
+            return handler(value)
+        except ValidationError as err:
+            _process_validation_error(err, handler, info)
+            raise
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        from nat.cli.type_registry import GlobalTypeRegistry
+
+        type_registry = GlobalTypeRegistry.get()
+
+        LLMsAnnotation = dict[str,
+                              typing.Annotated[type_registry.compute_annotation(LLMBaseConfig),
+                                               Discriminator(TypedBaseModel.discriminator)]]
+
+        AuthenticationProviderAnnotation = dict[str,
+                                                typing.Annotated[
+                                                    type_registry.compute_annotation(AuthProviderBaseConfig),
+                                                    Discriminator(TypedBaseModel.discriminator)]]
+
+        EmbeddersAnnotation = dict[str,
+                                   typing.Annotated[type_registry.compute_annotation(EmbedderBaseConfig),
+                                                    Discriminator(TypedBaseModel.discriminator)]]
+
+        FunctionsAnnotation = dict[str,
+                                   typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
+                                                    Discriminator(TypedBaseModel.discriminator)]]
+
+        MemoryAnnotation = dict[str,
+                                typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
+                                                 Discriminator(TypedBaseModel.discriminator)]]
+
+        ObjectStoreAnnotation = dict[str,
+                                     typing.Annotated[type_registry.compute_annotation(ObjectStoreBaseConfig),
+                                                      Discriminator(TypedBaseModel.discriminator)]]
+
+        RetrieverAnnotation = dict[str,
+                                   typing.Annotated[type_registry.compute_annotation(RetrieverBaseConfig),
+                                                    Discriminator(TypedBaseModel.discriminator)]]
+
+        TTCStrategyAnnotation = dict[str,
+                                     typing.Annotated[type_registry.compute_annotation(TTCStrategyBaseConfig),
+                                                      Discriminator(TypedBaseModel.discriminator)]]
+
+        WorkflowAnnotation = typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
+                                              Discriminator(TypedBaseModel.discriminator)]
+
+        should_rebuild = False
+
+        auth_providers_field = cls.model_fields.get("authentication")
+        if auth_providers_field is not None and auth_providers_field.annotation != AuthenticationProviderAnnotation:
+            auth_providers_field.annotation = AuthenticationProviderAnnotation
+            should_rebuild = True
+
+        llms_field = cls.model_fields.get("llms")
+        if llms_field is not None and llms_field.annotation != LLMsAnnotation:
+            llms_field.annotation = LLMsAnnotation
+            should_rebuild = True
+
+        embedders_field = cls.model_fields.get("embedders")
+        if embedders_field is not None and embedders_field.annotation != EmbeddersAnnotation:
+            embedders_field.annotation = EmbeddersAnnotation
+            should_rebuild = True
+
+        functions_field = cls.model_fields.get("functions")
+        if functions_field is not None and functions_field.annotation != FunctionsAnnotation:
+            functions_field.annotation = FunctionsAnnotation
+            should_rebuild = True
+
+        memory_field = cls.model_fields.get("memory")
+        if memory_field is not None and memory_field.annotation != MemoryAnnotation:
+            memory_field.annotation = MemoryAnnotation
+            should_rebuild = True
+
+        object_stores_field = cls.model_fields.get("object_stores")
+        if object_stores_field is not None and object_stores_field.annotation != ObjectStoreAnnotation:
+            object_stores_field.annotation = ObjectStoreAnnotation
+            should_rebuild = True
+
+        retrievers_field = cls.model_fields.get("retrievers")
+        if retrievers_field is not None and retrievers_field.annotation != RetrieverAnnotation:
+            retrievers_field.annotation = RetrieverAnnotation
+            should_rebuild = True
+
+        ttc_strategies_field = cls.model_fields.get("ttc_strategies")
+        if ttc_strategies_field is not None and ttc_strategies_field.annotation != TTCStrategyAnnotation:
+            ttc_strategies_field.annotation = TTCStrategyAnnotation
+            should_rebuild = True
+
+        workflow_field = cls.model_fields.get("workflow")
+        if workflow_field is not None and workflow_field.annotation != WorkflowAnnotation:
+            workflow_field.annotation = WorkflowAnnotation
+            should_rebuild = True
+
+        if (GeneralConfig.rebuild_annotations()):
+            should_rebuild = True
+
+        if (EvalConfig.rebuild_annotations()):
+            should_rebuild = True
+
+        if (should_rebuild):
+            return cls.model_rebuild(force=True)
+
+        return False
+
+
+# Compatibility aliases with previous releases
+AIQConfig = Config

nat/data_models/dataset_handler.py

@@ -0,0 +1,169 @@
+# 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 json
+import typing
+from collections.abc import Callable
+from pathlib import Path
+
+import pandas as pd
+from pydantic import BaseModel
+from pydantic import Discriminator
+from pydantic import FilePath
+from pydantic import Tag
+
+from nat.data_models.common import BaseModelRegistryTag
+from nat.data_models.common import TypedBaseModel
+
+
+class EvalS3Config(BaseModel):
+
+    endpoint_url: str | None = None
+    region_name: str | None = None
+    bucket: str
+    access_key: str
+    secret_key: str
+
+
+class EvalFilterEntryConfig(BaseModel):
+    # values are lists of allowed/blocked values
+    field: dict[str, list[str | int | float]] = {}
+
+
+class EvalFilterConfig(BaseModel):
+    allowlist: EvalFilterEntryConfig | None = None
+    denylist: EvalFilterEntryConfig | None = None
+
+
+class EvalDatasetStructureConfig(BaseModel):
+    disable: bool = False
+    question_key: str = "question"
+    answer_key: str = "answer"
+    generated_answer_key: str = "generated_answer"
+    trajectory_key: str = "intermediate_steps"
+    expected_trajectory_key: str = "expected_intermediate_steps"
+
+
+# Base model
+class EvalDatasetBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+
+    id_key: str = "id"
+    structure: EvalDatasetStructureConfig = EvalDatasetStructureConfig()
+
+    # Filters
+    filter: EvalFilterConfig | None = EvalFilterConfig()
+
+    s3: EvalS3Config | None = None
+
+    remote_file_path: str | None = None  # only for s3
+    file_path: Path | str = Path(".tmp/nat/examples/default/default.json")
+
+
+class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_json, {}
+
+
+def read_jsonl(file_path: FilePath):
+    with open(file_path, 'r', encoding='utf-8') as f:
+        data = [json.loads(line) for line in f]
+    return pd.DataFrame(data)
+
+
+class EvalDatasetJsonlConfig(EvalDatasetBaseConfig, name="jsonl"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return read_jsonl, {}
+
+
+class EvalDatasetCsvConfig(EvalDatasetBaseConfig, name="csv"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_csv, {}
+
+
+class EvalDatasetParquetConfig(EvalDatasetBaseConfig, name="parquet"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_parquet, {}
+
+
+class EvalDatasetXlsConfig(EvalDatasetBaseConfig, name="xls"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_excel, {"engine": "openpyxl"}
+
+
+class EvalDatasetCustomConfig(EvalDatasetBaseConfig, name="custom"):
+    """
+    Configuration for custom dataset type that allows users to specify
+    a custom Python function to transform their dataset into EvalInput format.
+    """
+
+    function: str  # Direct import path to function, format: "module.path.function_name"
+    kwargs: dict[str, typing.Any] = {}  # Additional arguments to pass to the custom function
+
+    def parser(self) -> tuple[Callable, dict]:
+        """
+        Load and return the custom function for dataset transformation.
+
+        Returns:
+            Tuple of (custom_function, kwargs) where custom_function transforms
+            a dataset file into an EvalInput object.
+        """
+        custom_function = self._load_custom_function()
+        return custom_function, self.kwargs
+
+    def _load_custom_function(self) -> Callable:
+        """
+        Import and return the custom function using standard Python import path.
+        """
+        if not self.function:
+            raise ValueError("Function path cannot be empty")
+
+        # Split the function path to get module and function name
+        module_path, function_name = self.function.rsplit(".", 1)
+
+        # Import the module
+        module = importlib.import_module(module_path)
+
+        # Get the function from the module
+        if not hasattr(module, function_name):
+            raise AttributeError(f"Function '{function_name}' not found in module '{module_path}'")
+
+        custom_function = getattr(module, function_name)
+
+        if not callable(custom_function):
+            raise ValueError(f"'{self.function}' is not callable")
+
+        return custom_function
+
+
+# Union model with discriminator
+EvalDatasetConfig = typing.Annotated[
+    typing.Annotated[EvalDatasetJsonConfig, Tag(EvalDatasetJsonConfig.static_type())]
+    | typing.Annotated[EvalDatasetCsvConfig, Tag(EvalDatasetCsvConfig.static_type())]
+    | typing.Annotated[EvalDatasetXlsConfig, Tag(EvalDatasetXlsConfig.static_type())]
+    | typing.Annotated[EvalDatasetParquetConfig, Tag(EvalDatasetParquetConfig.static_type())]
+    | typing.Annotated[EvalDatasetJsonlConfig, Tag(EvalDatasetJsonlConfig.static_type())]
+    | typing.Annotated[EvalDatasetCustomConfig, Tag(EvalDatasetCustomConfig.static_type())],
+    Discriminator(TypedBaseModel.discriminator)]