nvidia-nat 1.2.0__py3-none-any.whl

nat/data_models/discovery_metadata.py

@@ -0,0 +1,305 @@
+# 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.metadata
+import inspect
+import logging
+import typing
+from enum import Enum
+from functools import lru_cache
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+from pydantic import field_validator
+
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.component import ComponentEnum
+from nat.utils.metadata_utils import generate_config_type_docs
+
+if TYPE_CHECKING:
+    from nat.cli.type_registry import ToolWrapperBuildCallableT
+    from nat.data_models.common import TypedBaseModelT
+
+logger = logging.getLogger(__name__)
+
+
+class DiscoveryStatusEnum(str, Enum):
+    SUCCESS = "success"
+    FAILURE = "failure"
+
+
+class DiscoveryContractFieldsEnum(str, Enum):
+    PACKAGE = "package"
+    VERSION = "version"
+    COMPONENT_TYPE = "component_type"
+    COMPONENT_NAME = "component_name"
+    DESCRIPTION = "description"
+    DEVELOPER_NOTES = "developer_notes"
+
+
+class DiscoveryMetadata(BaseModel):
+    """A data model representing metadata about each registered component to faciliate its discovery.
+
+    Args:
+        package (str): The name of the package containing the NAT component.
+        version (str): The version number of the package containing the NAT component.
+        component_type (ComponentEnum): The type of NAT component this metadata represents.
+        component_name (str): The registered name of the NAT component.
+        description (str): Description of the NAT component pulled from its config objects docstrings.
+        developer_notes (str): Other notes to a developers to aid in the use of the component.
+        status (DiscoveryStatusEnum): Provides the status of the metadata discovery process.
+    """
+
+    package: str = ""
+    version: str = ""
+    component_type: ComponentEnum = ComponentEnum.UNDEFINED
+    component_name: str = ""
+    description: str = ""
+    developer_notes: str = ""
+    status: DiscoveryStatusEnum = DiscoveryStatusEnum.SUCCESS
+
+    @field_validator("description", mode="before")
+    @classmethod
+    def ensure_description_string(cls, v: typing.Any):
+        if not isinstance(v, str):
+            return ""
+        return v
+
+    @staticmethod
+    def get_preferred_item(items: list, preferred: str) -> str:
+        return preferred if preferred in items else items[0]
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_metadata(root_package_name: str) -> str | None:
+        """
+        This is not performant and is only present to be used (not used
+        currently) as a fallback when the distro name doesn't match the
+        module name and private_data is not available to map it.
+        """
+        mapping = importlib.metadata.packages_distributions()
+        try:
+            distro_names = mapping.get(root_package_name, [None])
+            distro_name = DiscoveryMetadata.get_preferred_item(distro_names, "nvidia-nat")
+        except KeyError:
+            return root_package_name
+
+        return distro_name if distro_name else root_package_name
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_module(module: ModuleType | None) -> str:
+        """Get the distribution name from the config type using the mapping of module names to distro names.
+
+        Args:
+            module (ModuleType): A registered component's module.
+
+        Returns:
+            str: The distribution name of the NAT component.
+        """
+        from nat.runtime.loader import get_all_entrypoints_distro_mapping
+
+        if module is None:
+            return "nvidia-nat"
+
+        # Get the mapping of module names to distro names
+        mapping = get_all_entrypoints_distro_mapping()
+        module_package = module.__package__
+
+        if module_package is None:
+            return "nvidia-nat"
+
+        # Traverse the module package parts in reverse order to find the distro name
+        # This is because the module package is the root package for the NAT component
+        # and the distro name is the name of the package that contains the component
+        module_package_parts = module_package.split(".")
+        for part_idx in range(len(module_package_parts), 0, -1):
+            candidate_module_name = ".".join(module_package_parts[0:part_idx])
+            candidate_distro_name = mapping.get(candidate_module_name, None)
+            if candidate_distro_name is not None:
+                return candidate_distro_name
+
+        return "nvidia-nat"
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_config_type(config_type: type["TypedBaseModelT"]) -> str:
+        """Get the distribution name from the config type using the mapping of module names to distro names.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+
+        Returns:
+            str: The distribution name of the NAT component.
+        """
+        module = inspect.getmodule(config_type)
+        return DiscoveryMetadata.get_distribution_name_from_module(module)
+
+    @staticmethod
+    def from_config_type(config_type: type["TypedBaseModelT"],
+                         component_type: ComponentEnum = ComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
+        """Generates discovery metadata from a NAT config object.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+            component_type (ComponentEnum, optional): The type of the registered component. Defaults to
+            ComponentEnum.UNDEFINED.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(config_type)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_config_type(config_type)
+
+            if not distro_name:
+                # raise an exception
+                logger.error("Encountered issue getting distro_name for module %s", module.__name__)
+                return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+            try:
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                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)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        description = generate_config_type_docs(config_type=config_type)
+
+        return DiscoveryMetadata(package=distro_name,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=config_type.static_type(),
+                                 description=description)
+
+    @staticmethod
+    def from_fn_wrapper(fn: "ToolWrapperBuildCallableT",
+                        wrapper_type: LLMFrameworkEnum | str,
+                        component_type: ComponentEnum = ComponentEnum.TOOL_WRAPPER) -> "DiscoveryMetadata":
+        """Generates discovery metadata from function with specified wrapper type.
+
+        Args:
+            fn (ToolWrapperBuildCallableT): A tool wrapper callable to source component metadata.
+            wrapper_type (LLMFrameworkEnum): The wrapper to apply to the callable to faciliate inter-framwork
+            interoperability.
+
+            component_type (ComponentEnum, optional): The type of the registered component. Defaults to
+            ComponentEnum.TOOL_WRAPPER.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(fn)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
+
+            try:
+                # version = importlib.metadata.version(root_package) if root_package != "" else ""
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                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)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        if isinstance(wrapper_type, LLMFrameworkEnum):
+            wrapper_type = wrapper_type.value
+
+        return DiscoveryMetadata(package=distro_name,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=wrapper_type,
+                                 description=fn.__doc__ or "")
+
+    @staticmethod
+    def from_package_name(package_name: str, package_version: str | None) -> "DiscoveryMetadata":
+        """Generates discovery metadata from an installed package name.
+
+        Args:
+            package_name (str): The name of the NAT plugin package containing registered components.
+            package_version (str, optional): The version of the package, Defaults to None.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            try:
+                metadata = importlib.metadata.metadata(package_name)
+                description = metadata.get("Summary", "")
+                if (package_version is None):
+                    package_version = importlib.metadata.version(package_name)
+            except importlib.metadata.PackageNotFoundError:
+                logger.warning("Package metadata not found for %s", package_name)
+                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)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        return DiscoveryMetadata(package=package_name,
+                                 version=package_version,
+                                 component_type=ComponentEnum.PACKAGE,
+                                 component_name=package_name,
+                                 description=description)
+
+    @staticmethod
+    def from_provider_framework_map(config_type: type["TypedBaseModelT"],
+                                    wrapper_type: LLMFrameworkEnum | str | None,
+                                    provider_type: ComponentEnum,
+                                    component_type: ComponentEnum = ComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
+        """Generates discovery metadata from provider and framework mapping information.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+            wrapper_type (LLMFrameworkEnum | str): The wrapper to apply to the callable to faciliate inter-framwork
+            interoperability.
+
+            provider_type (ComponentEnum): The type of provider the registered component supports.
+            component_type (ComponentEnum, optional): The type of the registered component. Defaults to
+            ComponentEnum.UNDEFINED.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(config_type)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
+            try:
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                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)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type
+        component_name = f"{config_type.static_type()} ({provider_type.value}) - {wrapper_type}"
+
+        description = generate_config_type_docs(config_type=config_type)
+
+        return DiscoveryMetadata(package=distro_name,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=component_name,
+                                 description=description)

nat/data_models/embedder.py

@@ -0,0 +1,27 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class EmbedderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """ Base configuration for embedding model providers. """
+    pass
+
+
+EmbedderBaseConfigT = typing.TypeVar("EmbedderBaseConfigT", bound=EmbedderBaseConfig)

nat/data_models/evaluate.py

@@ -0,0 +1,127 @@
+# 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 typing
+from enum import Enum
+from pathlib import Path
+
+from pydantic import BaseModel
+from pydantic import Discriminator
+from pydantic import model_validator
+
+from nat.data_models.common import TypedBaseModel
+from nat.data_models.dataset_handler import EvalDatasetConfig
+from nat.data_models.dataset_handler import EvalS3Config
+from nat.data_models.evaluator import EvaluatorBaseConfig
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.profiler import ProfilerConfig
+
+
+class JobEvictionPolicy(str, Enum):
+    """Policy for evicting old jobs when max_jobs is exceeded."""
+    TIME_CREATED = "time_created"
+    TIME_MODIFIED = "time_modified"
+
+
+class EvalCustomScriptConfig(BaseModel):
+    # Path to the script to run
+    script: Path
+    # Keyword arguments to pass to the script
+    kwargs: dict[str, str] = {}
+
+
+class JobManagementConfig(BaseModel):
+    # Whether to append a unique job ID to the output directory for each run
+    append_job_id_to_output_dir: bool = False
+    # Maximum number of jobs to keep in the output directory. Oldest jobs will be evicted.
+    # A value of 0 means no limit.
+    max_jobs: int = 0
+    # Policy for evicting old jobs. Defaults to using time_created.
+    eviction_policy: JobEvictionPolicy = JobEvictionPolicy.TIME_CREATED
+
+
+class EvalOutputConfig(BaseModel):
+    # Output directory for the workflow and evaluation results
+    dir: Path = Path("./.tmp/nat/examples/default/")
+    # S3 prefix for the workflow and evaluation results
+    remote_dir: 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
+    s3: EvalS3Config | None = None
+    # Whether to cleanup the output directory before running the workflow
+    cleanup: bool = True
+    # Job management configuration (job id, eviction, etc.)
+    job_management: JobManagementConfig = JobManagementConfig()
+    # Filter for the workflow output steps
+    workflow_output_step_filter: list[IntermediateStepType] | None = None
+
+
+class EvalGeneralConfig(BaseModel):
+    max_concurrency: int = 8
+
+    # Workflow alias for displaying in evaluation UI, if not provided,
+    # the workflow type will be used
+    workflow_alias: str | None = None
+
+    # Output directory for the workflow and evaluation results
+    output_dir: Path = Path("./.tmp/nat/examples/default/")
+
+    # If present overrides output_dir
+    output: EvalOutputConfig | None = None
+
+    # Dataset for running the workflow and evaluating
+    dataset: EvalDatasetConfig | None = None
+
+    # Inference profiler
+    profiler: ProfilerConfig | None = None
+
+    # overwrite the output_dir with the output config if present
+    @model_validator(mode="before")
+    @classmethod
+    def override_output_dir(cls, values):
+        if values.get("output") and values["output"].get("dir"):
+            values["output_dir"] = values["output"]["dir"]
+        return values
+
+
+class EvalConfig(BaseModel):
+
+    # General Evaluation Options
+    general: EvalGeneralConfig = EvalGeneralConfig()
+
+    # Evaluators
+    evaluators: dict[str, EvaluatorBaseConfig] = {}
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        from nat.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+
+        type_registry = GlobalTypeRegistry.get()
+
+        EvaluatorsAnnotation = dict[str,
+                                    typing.Annotated[type_registry.compute_annotation(EvaluatorBaseConfig),
+                                                     Discriminator(TypedBaseModel.discriminator)]]
+
+        should_rebuild = False
+
+        evaluators_field = cls.model_fields.get("evaluators")
+        if evaluators_field is not None and evaluators_field.annotation != EvaluatorsAnnotation:
+            evaluators_field.annotation = EvaluatorsAnnotation
+            should_rebuild = True
+
+        if (should_rebuild):
+            cls.model_rebuild(force=True)

nat/data_models/evaluator.py

@@ -0,0 +1,26 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class EvaluatorBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+EvaluatorBaseConfigT = typing.TypeVar("EvaluatorBaseConfigT", bound=EvaluatorBaseConfig)

nat/data_models/front_end.py

@@ -0,0 +1,26 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class FrontEndBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+FrontEndConfigT = typing.TypeVar("FrontEndConfigT", bound=FrontEndBaseConfig)

nat/data_models/function.py

@@ -0,0 +1,30 @@
+# 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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
+    pass
+
+
+FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)

nat/data_models/function_dependencies.py

@@ -0,0 +1,72 @@
+# 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 BaseModel
+from pydantic import Field
+from pydantic import field_serializer
+
+
+class FunctionDependencies(BaseModel):
+    """
+    A class to represent the dependencies of a function.
+    """
+    functions: 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)
+    object_stores: set[str] = Field(default_factory=set)
+    retrievers: set[str] = Field(default_factory=set)
+
+    @field_serializer("functions", when_used="json")
+    def serialize_functions(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)
+
+    @field_serializer("embedders", when_used="json")
+    def serialize_embedders(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("memory_clients", when_used="json")
+    def serialize_memory_clients(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("object_stores", when_used="json")
+    def serialize_object_stores(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("retrievers", when_used="json")
+    def serialize_retrievers(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    def add_function(self, function: str):
+        self.functions.add(function)  # pylint: disable=no-member
+
+    def add_llm(self, llm: str):
+        self.llms.add(llm)  # pylint: disable=no-member
+
+    def add_embedder(self, embedder: str):
+        self.embedders.add(embedder)  # pylint: disable=no-member
+
+    def add_memory_client(self, memory_client: str):
+        self.memory_clients.add(memory_client)  # pylint: disable=no-member
+
+    def add_object_store(self, object_store: str):
+        self.object_stores.add(object_store)  # pylint: disable=no-member
+
+    def add_retriever(self, retriever: str):
+        self.retrievers.add(retriever)  # pylint: disable=no-member