nvidia-nat 1.2.0rc5__py3-none-any.whl → 1.2.0rc6__py3-none-any.whl

{aiq → nat}/registry_handlers/schemas/publish.py

@@ -17,38 +17,38 @@ import logging
 
 from pydantic import BaseModel
 
-from aiq.data_models.component import AIQComponentEnum
-from aiq.data_models.discovery_metadata import DiscoveryMetadata
-from aiq.registry_handlers.schemas.status import StatusMessage
+from nat.data_models.component import ComponentEnum
+from nat.data_models.discovery_metadata import DiscoveryMetadata
+from nat.registry_handlers.schemas.status import StatusMessage
 
 logger = logging.getLogger(__name__)
 
 
-class BuiltAIQArtifact(BaseModel):
-    """An AIQ Toolkit artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
+class BuiltArtifact(BaseModel):
+    """A NAT artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
 
     Args:
-        whl (str): A base64 encoded string of an AIQ Toolkit package wheel (.whl).
+        whl (str): A base64 encoded string of a NAT package wheel (.whl).
 
-        metadata (dict[AIQComponentEnum, list[DiscoveryMetadata]]): Provides rich discover metadata for developers to
+        metadata (dict[ComponentEnum, list[DiscoveryMetadata]]): Provides rich discover metadata for developers to
         quickly find useful components.
     """
 
     whl: str
-    metadata: dict[AIQComponentEnum, list[DiscoveryMetadata]]
+    metadata: dict[ComponentEnum, list[DiscoveryMetadata]]
 
 
-class AIQArtifact(BaseModel):
-    """An AIQ Toolkit artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
+class Artifact(BaseModel):
+    """A NAT artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
 
     Args:
-        artifact (BuildAIQArtifact): An AIQ Toolkit artifact including base64 encoded string of wheel package and
+        artifact (BuiltArtifact): A NAT artifact including base64 encoded string of wheel package and
         corrosponding discovery metadata.
 
         whl_path (str): A local path to the built wheel package.
     """
 
-    artifact: BuiltAIQArtifact | None = None
+    artifact: BuiltArtifact | None = None
     whl_path: str
 
 
@@ -61,3 +61,8 @@ class PublishResponse(BaseModel):
     """
 
     status: StatusMessage
+
+
+# Compatibility aliases with previous releases
+BuiltAIQArtifact = BuiltArtifact
+AIQArtifact = Artifact

{aiq → nat}/registry_handlers/schemas/pull.py

@@ -17,8 +17,8 @@ import logging
 
 from pydantic import BaseModel
 
-from aiq.registry_handlers.schemas.package import PackageNameVersion
-from aiq.registry_handlers.schemas.status import StatusMessage
+from nat.registry_handlers.schemas.package import PackageNameVersion
+from nat.registry_handlers.schemas.status import StatusMessage
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +27,7 @@ class PulledPackage(BaseModel):
     """Represents a data model of a pulled package containing the package wheel and its name.
 
     Args:
-        whl (str): Base64 encoded string of the AIQ Toolkit python package wheel (.whl).
+        whl (str): Base64 encoded string of the NAT python package wheel (.whl).
         whl_name (str): A string representing the wheel filename.
     """
 
@@ -37,7 +37,7 @@ class PulledPackage(BaseModel):
 
 class PullResponse(BaseModel):
     """
-    Represents a data model of the expected respones from a AIQ Toolkit pull request, including detailed status
+    Represents a data model of the expected respones from a NAT pull request, including detailed status
     information.
 
     Args:
@@ -72,11 +72,11 @@ class PullRequestPackage(BaseModel):
 
 
 class PullRequestPackages(BaseModel):
-    """Represents a list of all packages th download and install in the local AIQ Toolkit environment.
+    """Represents a list of all packages th download and install in the local NAT environment.
 
     Args:
         packages (list[typing.Union[PackageNameVersion, PullPackageWhl]]): A list of packages that can be
-            downloaded and installed in the local AIQ Toolkit environment.
+            downloaded and installed in the local NAT environment.
     """
 
     packages: list[PackageNameVersion | PullPackageWhl]

{aiq → nat}/registry_handlers/schemas/remove.py

@@ -17,8 +17,8 @@ import logging
 
 from pydantic import BaseModel
 
-from aiq.registry_handlers.schemas.package import PackageNameVersion
-from aiq.registry_handlers.schemas.status import StatusMessage
+from nat.registry_handlers.schemas.package import PackageNameVersion
+from nat.registry_handlers.schemas.status import StatusMessage
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/registry_handlers/schemas/search.py

@@ -18,8 +18,8 @@ from enum import Enum
 
 from pydantic import BaseModel
 
-from aiq.data_models.component import AIQComponentEnum
-from aiq.registry_handlers.schemas.status import StatusMessage
+from nat.data_models.component import ComponentEnum
+from nat.registry_handlers.schemas.status import StatusMessage
 
 logger = logging.getLogger(__name__)
 
@@ -42,18 +42,18 @@ class VisualizeFields(str, Enum):
 
 
 class SearchQuery(BaseModel):
-    """Represents the search criteria that will be used to discover useful AIQ Toolkit components.
+    """Represents the search criteria that will be used to discover useful NAT components.
 
     Args:
-        query (str): A query string used to find useful AIQ Toolkit components.
+        query (str): A query string used to find useful NAT components.
         fields (list[SearchFields]): The list of fields used when applying the query string.
-        component_types (list[AIQComponentEnum]): AIQ Toolkit components types to filter search results.
+        component_types (list[ComponentEnum]): NAT components types to filter search results.
         top_k (int): Specifies the number of search results to provide.
     """
 
     query: str = "*"
     fields: list[SearchFields] = [SearchFields.ALL]
-    component_types: list[AIQComponentEnum]
+    component_types: list[ComponentEnum]
     top_k: int = 10
 
 
@@ -61,16 +61,16 @@ class SearchResponseItem(BaseModel):
     """Represents an individual item in the search response, including elements of it's discovery metadata.
 
     Args:
-        package (str): The name of the AIQ Toolkit package that includes the component.
-        version (str): The version of the AIQ Toolkit package that includes the component.
-        component_type (AIQComponentEnum): Type of AIQ Toolkit component this item represents.
-        description (str): A description of this AIQ Toolkit component.
+        package (str): The name of the NAT package that includes the component.
+        version (str): The version of the NAT package that includes the component.
+        component_type (ComponentEnum): Type of NAT component this item represents.
+        description (str): A description of this NAT component.
         developer_notes (str): Additional details that would help a developer use this component.
     """
 
     package: str
     version: str
-    component_type: AIQComponentEnum
+    component_type: ComponentEnum
     component_name: str
     description: str
     developer_notes: str

{aiq → nat}/retriever/interface.py

@@ -16,10 +16,10 @@
 from abc import ABC
 from abc import abstractmethod
 
-from aiq.retriever.models import RetrieverOutput
+from nat.retriever.models import RetrieverOutput
 
 
-class AIQRetriever(ABC):
+class Retriever(ABC):
     """
     Abstract interface for interacting with data stores.
 
@@ -35,3 +35,7 @@ class AIQRetriever(ABC):
 
         """
         raise NotImplementedError
+
+
+# Compatibility aliases with previous releases
+AIQRetriever = Retriever

{aiq → nat}/retriever/milvus/register.py

@@ -16,12 +16,12 @@
 from pydantic import Field
 from pydantic import HttpUrl
 
-from aiq.builder.builder import Builder
-from aiq.builder.builder import LLMFrameworkEnum
-from aiq.builder.retriever import RetrieverProviderInfo
-from aiq.cli.register_workflow import register_retriever_client
-from aiq.cli.register_workflow import register_retriever_provider
-from aiq.data_models.retriever import RetrieverBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.builder import LLMFrameworkEnum
+from nat.builder.retriever import RetrieverProviderInfo
+from nat.cli.register_workflow import register_retriever_client
+from nat.cli.register_workflow import register_retriever_provider
+from nat.data_models.retriever import RetrieverBaseConfig
 
 
 class MilvusRetrieverConfig(RetrieverBaseConfig, name="milvus_retriever"):
@@ -60,7 +60,7 @@ async def milvus_retriever(retriever_config: MilvusRetrieverConfig, builder: Bui
 async def milvus_retriever_client(config: MilvusRetrieverConfig, builder: Builder):
     from pymilvus import MilvusClient
 
-    from aiq.retriever.milvus.retriever import MilvusRetriever
+    from nat.retriever.milvus.retriever import MilvusRetriever
 
     embedder = await builder.get_embedder(embedder_name=config.embedding_model, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 

{aiq → nat}/retriever/milvus/retriever.py

@@ -20,10 +20,10 @@ from langchain_core.embeddings import Embeddings
 from pymilvus import MilvusClient
 from pymilvus.client.abstract import Hit
 
-from aiq.retriever.interface import AIQRetriever
-from aiq.retriever.models import AIQDocument
-from aiq.retriever.models import RetrieverError
-from aiq.retriever.models import RetrieverOutput
+from nat.retriever.interface import Retriever
+from nat.retriever.models import Document
+from nat.retriever.models import RetrieverError
+from nat.retriever.models import RetrieverOutput
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +32,7 @@ class CollectionNotFoundError(RetrieverError):
     pass
 
 
-class MilvusRetriever(AIQRetriever):
+class MilvusRetriever(Retriever):
     """
     Client for retrieving document chunks from a Milvus vectorstore
     """
@@ -213,16 +213,16 @@ def _wrap_milvus_results(res: list[Hit], content_field: str):
     return RetrieverOutput(results=[_wrap_milvus_single_results(r, content_field=content_field) for r in res])
 
 
-def _wrap_milvus_single_results(res: Hit | dict, content_field: str) -> AIQDocument:
+def _wrap_milvus_single_results(res: Hit | dict, content_field: str) -> Document:
     if not isinstance(res, (Hit, dict)):
         raise ValueError(f"Milvus search returned object of type {type(res)}. Expected 'Hit' or 'dict'.")
 
     if isinstance(res, Hit):
         metadata = {k: v for k, v in res.fields.items() if k != content_field}
         metadata.update({"distance": res.distance})
-        return AIQDocument(page_content=res.fields[content_field], metadata=metadata, document_id=res.id)
+        return Document(page_content=res.fields[content_field], metadata=metadata, document_id=res.id)
 
     fields = res["entity"]
     metadata = {k: v for k, v in fields.items() if k != content_field}
     metadata.update({"distance": res.get("distance")})
-    return AIQDocument(page_content=fields.get(content_field), metadata=metadata, document_id=res["id"])
+    return Document(page_content=fields.get(content_field), metadata=metadata, document_id=res["id"])

{aiq → nat}/retriever/models.py

@@ -21,22 +21,22 @@ from typing import Any
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.utils.type_converter import GlobalTypeConverter
+from nat.utils.type_converter import GlobalTypeConverter
 
 
-class AIQDocument(BaseModel):
+class Document(BaseModel):
     """
-    Object representing a retrieved document/chunk from a standard AIQ Toolkit Retriever.
+    Object representing a retrieved document/chunk from a standard NAT Retriever.
     """
     page_content: str = Field(description="Primary content of the document to insert or retrieve")
-    metadata: dict[str, Any] = Field(description="Metadata dictionary attached to the AIQDocument")
+    metadata: dict[str, Any] = Field(description="Metadata dictionary attached to the Document")
     document_id: str | None = Field(description="Unique ID for the document, if supported by the configured datastore",
                                     default=None)
 
     @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> AIQDocument:
+    def from_dict(cls, data: dict[str, Any]) -> Document:
         """
-        Deserialize an AIQDocument from a dictionary representation.
+        Deserialize an Document from a dictionary representation.
 
         Args:
             data (dict): A dictionary containing keys
@@ -49,7 +49,7 @@ class AIQDocument(BaseModel):
 
 
 class RetrieverOutput(BaseModel):
-    results: list[AIQDocument] = Field(description="A list of retrieved AIQDocuments")
+    results: list[Document] = Field(description="A list of retrieved Documents")
 
     def __len__(self):
         return len(self.results)
@@ -72,3 +72,6 @@ def retriever_output_to_str(obj: RetrieverOutput) -> str:
 
 GlobalTypeConverter.register_converter(retriever_output_to_dict)
 GlobalTypeConverter.register_converter(retriever_output_to_str)
+
+# Compatibility aliases with previous releases
+AIQDocument = Document

{aiq → nat}/retriever/nemo_retriever/register.py

@@ -16,11 +16,11 @@
 from pydantic import Field
 from pydantic import HttpUrl
 
-from aiq.builder.builder import Builder
-from aiq.builder.retriever import RetrieverProviderInfo
-from aiq.cli.register_workflow import register_retriever_client
-from aiq.cli.register_workflow import register_retriever_provider
-from aiq.data_models.retriever import RetrieverBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.retriever import RetrieverProviderInfo
+from nat.cli.register_workflow import register_retriever_client
+from nat.cli.register_workflow import register_retriever_provider
+from nat.data_models.retriever import RetrieverBaseConfig
 
 
 class NemoRetrieverConfig(RetrieverBaseConfig, name="nemo_retriever"):
@@ -48,7 +48,7 @@ async def nemo_retriever(retriever_config: NemoRetrieverConfig, builder: Builder
 
 @register_retriever_client(config_type=NemoRetrieverConfig, wrapper_type=None)
 async def nemo_retriever_client(config: NemoRetrieverConfig, builder: Builder):
-    from aiq.retriever.nemo_retriever.retriever import NemoRetriever
+    from nat.retriever.nemo_retriever.retriever import NemoRetriever
 
     retriever = NemoRetriever(**config.model_dump(exclude={"type", "top_k", "collection_name"}))
     optional_fields = ["collection_name", "top_k", "output_fields"]

{aiq → nat}/retriever/nemo_retriever/retriever.py

@@ -26,10 +26,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import HttpUrl
 
-from aiq.retriever.interface import AIQRetriever
-from aiq.retriever.models import AIQDocument
-from aiq.retriever.models import RetrieverError
-from aiq.retriever.models import RetrieverOutput
+from nat.retriever.interface import Retriever
+from nat.retriever.models import Document
+from nat.retriever.models import RetrieverError
+from nat.retriever.models import RetrieverOutput
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +51,7 @@ class CollectionUnavailableError(RetrieverError):
     pass
 
 
-class NemoRetriever(AIQRetriever):
+class NemoRetriever(Retriever):
     """
     Client for retrieving document chunks from a Nemo Retriever service.
     """
@@ -153,11 +153,11 @@ def _wrap_nemo_results(output: list[dict], content_field: str):
 
 
 def _wrap_nemo_single_results(output: dict, content_field: str):
-    return AIQDocument(page_content=output[content_field],
-                       metadata={
-                           k: v
-                           for k, v in output.items() if k != content_field
-                       })
+    return Document(page_content=output[content_field],
+                    metadata={
+                        k: v
+                        for k, v in output.items() if k != content_field
+                    })
 
 
 def _flatten(obj: dict, output_fields: list[str]) -> list[str]:

{aiq/embedder → nat/retriever}/register.py

@@ -18,7 +18,5 @@
 # isort:skip_file
 
 # Import any providers which need to be automatically registered here
-from . import nim_embedder
-from . import openai_embedder
-# Import any clients which need to be automatically registered here
-from . import langchain_client
+import nat.retriever.milvus.register
+import nat.retriever.nemo_retriever.register

{aiq → nat}/runtime/loader.py

@@ -24,14 +24,14 @@ from enum import auto
 from functools import lru_cache
 from functools import reduce
 
-from aiq.builder.workflow_builder import WorkflowBuilder
-from aiq.cli.type_registry import GlobalTypeRegistry
-from aiq.data_models.config import AIQConfig
-from aiq.runtime.session import AIQSessionManager
-from aiq.utils.data_models.schema_validator import validate_schema
-from aiq.utils.debugging_utils import is_debugger_attached
-from aiq.utils.io.yaml_tools import yaml_load
-from aiq.utils.type_utils import StrPath
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.cli.type_registry import GlobalTypeRegistry
+from nat.data_models.config import Config
+from nat.runtime.session import SessionManager
+from nat.utils.data_models.schema_validator import validate_schema
+from nat.utils.debugging_utils import is_debugger_attached
+from nat.utils.io.yaml_tools import yaml_load
+from nat.utils.type_utils import StrPath
 
 logger = logging.getLogger(__name__)
 
@@ -58,7 +58,7 @@ class PluginTypes(IntFlag):
     # Convenience flag for groups of plugin types
     CONFIG_OBJECT = COMPONENT | FRONT_END | EVALUATOR | AUTHENTICATION
     """
-    Any plugin that can be specified in the AIQ Toolkit configuration file.
+    Any plugin that can be specified in the NAT configuration file.
     """
     ALL = COMPONENT | FRONT_END | EVALUATOR | REGISTRY_HANDLER | AUTHENTICATION
     """
@@ -66,10 +66,10 @@ class PluginTypes(IntFlag):
     """
 
 
-def load_config(config_file: StrPath) -> AIQConfig:
+def load_config(config_file: StrPath) -> Config:
     """
-    This is the primary entry point for loading an AIQ Toolkit configuration file. It ensures that all plugins are
-    loaded and then validates the configuration file against the AIQConfig schema.
+    This is the primary entry point for loading a NAT configuration file. It ensures that all plugins are
+    loaded and then validates the configuration file against the Config schema.
 
     Parameters
     ----------
@@ -78,8 +78,8 @@ def load_config(config_file: StrPath) -> AIQConfig:
 
     Returns
     -------
-    AIQConfig
-        The validated AIQConfig object
+    Config
+        The validated Config object
     """
 
     # Ensure all of the plugins are loaded
@@ -87,17 +87,17 @@ def load_config(config_file: StrPath) -> AIQConfig:
 
     config_yaml = yaml_load(config_file)
 
-    # Validate configuration adheres to AIQ Toolkit schemas
-    validated_aiq_config = validate_schema(config_yaml, AIQConfig)
+    # Validate configuration adheres to NAT schemas
+    validated_nat_config = validate_schema(config_yaml, Config)
 
-    return validated_aiq_config
+    return validated_nat_config
 
 
 @asynccontextmanager
 async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
     """
-    Load the AIQ Toolkit configuration file and create an AIQRunner object. This is the primary entry point for running
-    AIQ Toolkit workflows.
+    Load the NAT configuration file and create an Runner object. This is the primary entry point for running
+    NAT workflows.
 
     Parameters
     ----------
@@ -114,7 +114,7 @@ async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
     # Must yield the workflow function otherwise it cleans up
     async with WorkflowBuilder.from_config(config=config) as workflow:
 
-        yield AIQSessionManager(workflow.build(), max_concurrency=max_concurrency)
+        yield SessionManager(workflow.build(), max_concurrency=max_concurrency)
 
 
 @lru_cache
@@ -128,32 +128,33 @@ def discover_entrypoints(plugin_type: PluginTypes):
     plugin_groups = []
 
     # Add the specified plugin type to the list of groups to load
+    # The aiq entrypoints are intentionally left in the list to maintain backwards compatibility.
     if (plugin_type & PluginTypes.COMPONENT):
-        plugin_groups.extend(["aiq.plugins", "aiq.components"])
+        plugin_groups.extend(["aiq.plugins", "aiq.components", "nat.plugins", "nat.components"])
     if (plugin_type & PluginTypes.FRONT_END):
-        plugin_groups.append("aiq.front_ends")
+        plugin_groups.extend(["aiq.front_ends", "nat.front_ends"])
     if (plugin_type & PluginTypes.REGISTRY_HANDLER):
-        plugin_groups.append("aiq.registry_handlers")
+        plugin_groups.extend(["aiq.registry_handlers", "nat.registry_handlers"])
     if (plugin_type & PluginTypes.EVALUATOR):
-        plugin_groups.append("aiq.evaluators")
+        plugin_groups.extend(["aiq.evaluators", "nat.evaluators"])
     if (plugin_type & PluginTypes.AUTHENTICATION):
-        plugin_groups.append("aiq.authentication_providers")
+        plugin_groups.extend(["aiq.authentication_providers", "nat.authentication_providers"])
 
     # Get the entry points for the specified groups
-    aiq_plugins = reduce(lambda x, y: list(x) + list(y), [entry_points.select(group=y) for y in plugin_groups])
+    nat_plugins = reduce(lambda x, y: list(x) + list(y), [entry_points.select(group=y) for y in plugin_groups])
 
-    return aiq_plugins
+    return nat_plugins
 
 
 @lru_cache
-def get_all_aiq_entrypoints_distro_mapping() -> dict[str, str]:
+def get_all_entrypoints_distro_mapping() -> dict[str, str]:
     """
-    Get the mapping of all AIQ entry points to their distribution names.
+    Get the mapping of all NAT entry points to their distribution names.
     """
 
     mapping = {}
-    aiq_entrypoints = discover_entrypoints(PluginTypes.ALL)
-    for ep in aiq_entrypoints:
+    nat_entrypoints = discover_entrypoints(PluginTypes.ALL)
+    for ep in nat_entrypoints:
         ep_module_parts = ep.module.split(".")
         current_parts = []
         for part in ep_module_parts:
@@ -171,14 +172,14 @@ def discover_and_register_plugins(plugin_type: PluginTypes):
     """
 
     # Get the entry points for the specified groups
-    aiq_plugins = discover_entrypoints(plugin_type)
+    nat_plugins = discover_entrypoints(plugin_type)
 
     count = 0
 
     # Pause registration hooks for performance. This is useful when loading a large number of plugins.
     with GlobalTypeRegistry.get().pause_registration_changed_hooks():
 
-        for entry_point in aiq_plugins:
+        for entry_point in nat_plugins:
             try:
                 logger.debug("Loading module '%s' from entry point '%s'...", entry_point.module, entry_point.name)
 
@@ -213,3 +214,7 @@ def discover_and_register_plugins(plugin_type: PluginTypes):
 
             finally:
                 count += 1
+
+
+# Compatibility alias
+get_all_aiq_entrypoints_distro_mapping = get_all_entrypoints_distro_mapping