aiqtoolkit 1.2.0a20250803__py3-none-any.whl → 1.2.0a20250806__py3-none-any.whl

aiq/data_models/discovery_metadata.py

@@ -21,6 +21,7 @@ import typing
 from enum import Enum
 from functools import lru_cache
 from pathlib import Path
+from types import ModuleType
 from typing import TYPE_CHECKING
 
 from pydantic import BaseModel
@@ -115,6 +116,55 @@ class DiscoveryMetadata(BaseModel):
                     return data.get(root_package, None)
         return None
 
+    @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 AIQ Toolkit component.
+        """
+        from aiq.runtime.loader import get_all_aiq_entrypoints_distro_mapping
+
+        if module is None:
+            return "aiqtoolkit"
+
+        # Get the mapping of module names to distro names
+        mapping = get_all_aiq_entrypoints_distro_mapping()
+        module_package = module.__package__
+
+        if module_package is None:
+            return "aiqtoolkit"
+
+        # 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 AIQ Toolkit 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 "aiqtoolkit"
+
+    @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 AIQ Toolkit component.
+        """
+        module = inspect.getmodule(config_type)
+        return DiscoveryMetadata.get_distribution_name_from_module(module)
+
     @staticmethod
     @lru_cache
     def get_distribution_name(root_package: str) -> str:
@@ -123,6 +173,7 @@ class DiscoveryMetadata(BaseModel):
         root package name 'aiq'. They provide mapping in a metadata file
         for optimized installation.
         """
+
         distro_name = DiscoveryMetadata.get_distribution_name_from_private_data(root_package)
         return distro_name if distro_name else root_package
 
@@ -142,8 +193,7 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(config_type)
-            root_package: str = module.__package__.split(".")[0]
-            distro_name = DiscoveryMetadata.get_distribution_name(root_package)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_config_type(config_type)
 
             if not distro_name:
                 # raise an exception
@@ -187,12 +237,13 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(fn)
-            root_package: str = module.__package__.split(".")[0]
-            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
+
             try:
-                version = importlib.metadata.version(root_package) if root_package != "" else ""
+                # 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", root_package)
+                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)
@@ -201,7 +252,7 @@ class DiscoveryMetadata(BaseModel):
         if isinstance(wrapper_type, LLMFrameworkEnum):
             wrapper_type = wrapper_type.value
 
-        return DiscoveryMetadata(package=root_package,
+        return DiscoveryMetadata(package=distro_name,
                                  version=version,
                                  component_type=component_type,
                                  component_name=wrapper_type,
@@ -220,7 +271,6 @@ class DiscoveryMetadata(BaseModel):
         """
 
         try:
-            package_name = DiscoveryMetadata.get_distribution_name(package_name)
             try:
                 metadata = importlib.metadata.metadata(package_name)
                 description = metadata.get("Summary", "")
@@ -263,12 +313,11 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(config_type)
-            root_package: str = module.__package__.split(".")[0]
-            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            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", root_package)
+                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)
@@ -279,7 +328,7 @@ class DiscoveryMetadata(BaseModel):
 
         description = generate_config_type_docs(config_type=config_type)
 
-        return DiscoveryMetadata(package=root_package,
+        return DiscoveryMetadata(package=distro_name,
                                  version=version,
                                  component_type=component_type,
                                  component_name=component_name,

aiq/registry_handlers/package_utils.py

@@ -119,16 +119,15 @@ def build_wheel(package_root: str) -> WheelData:
 
     whl_version = Wheel(whl_path).version
 
-    return WheelData(
-        package_root=package_root,
-        package_name=module_name,  # should it be module name or distro name here
-        toml_project=toml_project,
-        toml_dependencies=toml_dependencies,
-        toml_aiq_packages=toml_packages,
-        union_dependencies=union_dependencies,
-        whl_path=whl_path,
-        whl_base64=whl_base64,
-        whl_version=whl_version)
+    return WheelData(package_root=package_root,
+                     package_name=toml_project_name,
+                     toml_project=toml_project,
+                     toml_dependencies=toml_dependencies,
+                     toml_aiq_packages=toml_packages,
+                     union_dependencies=union_dependencies,
+                     whl_path=whl_path,
+                     whl_base64=whl_base64,
+                     whl_version=whl_version)
 
 
 def build_package_metadata(wheel_data: WheelData | None) -> dict[AIQComponentEnum, list[dict | DiscoveryMetadata]]:
@@ -155,7 +154,7 @@ def build_package_metadata(wheel_data: WheelData | None) -> dict[AIQComponentEnu
     if (wheel_data is not None):
         registry.register_package(package_name=wheel_data.package_name, package_version=wheel_data.whl_version)
         for entry_point in aiq_plugins:
-            package_name = entry_point.module.split('.')[0]
+            package_name = entry_point.dist.name
             if (package_name == wheel_data.package_name):
                 continue
             if (package_name in wheel_data.union_dependencies):
@@ -163,8 +162,7 @@ def build_package_metadata(wheel_data: WheelData | None) -> dict[AIQComponentEnu
 
     else:
         for entry_point in aiq_plugins:
-            package_name = entry_point.module.split('.')[0]
-            registry.register_package(package_name=package_name)
+            registry.register_package(package_name=entry_point.dist.name)
 
     discovery_metadata = {}
     for component_type in AIQComponentEnum:

aiq/runtime/loader.py

@@ -21,6 +21,7 @@ import time
 from contextlib import asynccontextmanager
 from enum import IntFlag
 from enum import auto
+from functools import lru_cache
 from functools import reduce
 
 from aiq.builder.workflow_builder import WorkflowBuilder
@@ -116,6 +117,7 @@ async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
         yield AIQSessionManager(workflow.build(), max_concurrency=max_concurrency)
 
 
+@lru_cache
 def discover_entrypoints(plugin_type: PluginTypes):
     """
     Discover all the requested plugin types which were registered via an entry point group and return them.
@@ -143,6 +145,25 @@ def discover_entrypoints(plugin_type: PluginTypes):
     return aiq_plugins
 
 
+@lru_cache
+def get_all_aiq_entrypoints_distro_mapping() -> dict[str, str]:
+    """
+    Get the mapping of all AIQ entry points to their distribution names.
+    """
+
+    mapping = {}
+    aiq_entrypoints = discover_entrypoints(PluginTypes.ALL)
+    for ep in aiq_entrypoints:
+        ep_module_parts = ep.module.split(".")
+        current_parts = []
+        for part in ep_module_parts:
+            current_parts.append(part)
+            module_prefix = ".".join(current_parts)
+            mapping[module_prefix] = ep.dist.name
+
+    return mapping
+
+
 def discover_and_register_plugins(plugin_type: PluginTypes):
     """
     Discover all the requested plugin types which were registered via an entry point group and register them into the

aiq/utils/dump_distro_mapping.py

@@ -0,0 +1,32 @@
+# 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 argparse
+import json
+
+from aiq.runtime.loader import get_all_aiq_entrypoints_distro_mapping
+
+
+def dump_distro_mapping(path: str):
+    mapping = get_all_aiq_entrypoints_distro_mapping()
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(mapping, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--path", type=str, required=True)
+    args = parser.parse_args()
+    dump_distro_mapping(args.path)

{aiqtoolkit-1.2.0a20250803.dist-info → aiqtoolkit-1.2.0a20250806.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqtoolkit
-Version: 1.2.0a20250803
+Version: 1.2.0a20250806
 Summary: NVIDIA Agent Intelligence toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation

{aiqtoolkit-1.2.0a20250803.dist-info → aiqtoolkit-1.2.0a20250806.dist-info}/RECORD

@@ -103,7 +103,7 @@ aiq/data_models/component.py,sha256=_HeHlDmz2fgJlVfmz0tG_yCyes86hQNaQwSWz1FDQvk,
 aiq/data_models/component_ref.py,sha256=_6DeNWJSHN3GSoAFe0RCQkCJ31_WA9GR3nP6vsxMj7Q,4594
 aiq/data_models/config.py,sha256=I7-9kJqpr6IvqmTU4nnALjDmq8YRlrKhYzq6g7DrJeo,17124
 aiq/data_models/dataset_handler.py,sha256=liMB3xRohkr4VTMmNWPvWi9qhbhlJQfQK36g5Rknweo,4027
-aiq/data_models/discovery_metadata.py,sha256=OcITQc5VeML4bTHurrsMNiK_oB3z7wudMxcyN7LI8pY,12785
+aiq/data_models/discovery_metadata.py,sha256=ycrLQ2HTT4lOFZJGwvvGfXPAFRswygmNUWW0nKFNtHg,14662
 aiq/data_models/embedder.py,sha256=nPhthEQDtzAMGd8gFRB1ZfJpN5M9DJvv0h28ohHnTmI,1002
 aiq/data_models/evaluate.py,sha256=WBeABZsIa6W04MPj24SRu4s-ty2PkJ7_4SLojXmj5Pk,4704
 aiq/data_models/evaluator.py,sha256=bd2njsyQB2t6ClJ66gJiCjYHsQpWZwPD7rsU0J109TI,939
@@ -324,7 +324,7 @@ aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py,s
 aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py,sha256=nD46SCVi7zwxdPXRN5c61O58_OgMA2WCfaZdRJRqgP4,16600
 aiq/registry_handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/registry_handlers/metadata_factory.py,sha256=Urr_7mLLpoU0VPjl2Nknl9aZhzaAAwb66ydTGBBrIwc,2778
-aiq/registry_handlers/package_utils.py,sha256=WCENHOO55Vhwv9DbLYjR9SMu_ERzxMKX6oO9ywe9cpM,7468
+aiq/registry_handlers/package_utils.py,sha256=waCWoW351EACof5NS1l8VKued03pknArRub8kHW2nM0,7459
 aiq/registry_handlers/register.py,sha256=k2gvV0htVpfMdzsRvZGnTZp17JA2Na8sO9ocMaxDSmo,902
 aiq/registry_handlers/registry_handler_base.py,sha256=BYLH6R9y3pySBDH9HKUR16rM0atFRt594IDyyhgCQVQ,5842
 aiq/registry_handlers/local/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -355,7 +355,7 @@ aiq/retriever/nemo_retriever/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W
 aiq/retriever/nemo_retriever/register.py,sha256=ODV-TZfXzDs1VJHHLdj2kC05odirtlQZSeh9c1zw8AQ,2893
 aiq/retriever/nemo_retriever/retriever.py,sha256=IvScUr9XuDLiMR__I3QsboLaM52N5D5Qu94qtTOGQw8,6958
 aiq/runtime/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
-aiq/runtime/loader.py,sha256=drLE-OnrAPjFgWrQ3ZsnhHSbon8GkRsUk0p7Oe8Jn2I,7117
+aiq/runtime/loader.py,sha256=F_CHLj3KR8PBxzqu1sy70JLyNBDFG55bFm9kVNGoSdM,7707
 aiq/runtime/runner.py,sha256=CqmlVAYfrBh3ml3t2n3V693RaNyxtK9ScWT4S-Isbr8,6365
 aiq/runtime/session.py,sha256=i1pIqopZCBgGJqVUskKLiBnZYH-lTdMhvFu56dXAU5A,6206
 aiq/runtime/user_metadata.py,sha256=9EiBc-EEJzOdpf3Q1obHqAdY_kRlJ1T0TVvY0Jonk6o,3692
@@ -400,6 +400,7 @@ aiq/tool/memory_tools/delete_memory_tool.py,sha256=wdB_I8y-1D1OpNtBi6ZOg36vvNkba
 aiq/tool/memory_tools/get_memory_tool.py,sha256=-i0Bt5xYeapbbd2wtAgPc0pOv0Dx4jK1-yyHG7YCeQ0,2749
 aiq/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/utils/debugging_utils.py,sha256=6M4JhbHDNDnfmSRGmHvT5IgEeWSHBore3VngdE_PMqc,1332
+aiq/utils/dump_distro_mapping.py,sha256=5I3IGsmhoW12VZrInxx9CyyVSEtRmhRyo1bSdJJUj-o,1155
 aiq/utils/log_utils.py,sha256=dZLHt7qFqLlpPqMMFO9UVtSkOpMjFwz9tkmbAfOiNlg,1355
 aiq/utils/metadata_utils.py,sha256=lGYvc8Gk0az4qZDGeRbVz4L7B_b-Gnjss8JT4goqL5I,2897
 aiq/utils/optional_imports.py,sha256=jQSVBc2fBSRw-2d6r8cEwvh5-di2EUUPakuuo9QbbwA,4039
@@ -428,10 +429,10 @@ aiq/utils/reactive/base/observer_base.py,sha256=UAlyAY_ky4q2t0P81RVFo2Bs_R7z5Nde
 aiq/utils/reactive/base/subject_base.py,sha256=Ed-AC6P7cT3qkW1EXjzbd5M9WpVoeN_9KCe3OM3FLU4,2521
 aiq/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/utils/settings/global_settings.py,sha256=U9TCLdoZsKq5qOVGjREipGVv9e-FlStzqy5zv82_VYk,7454
-aiqtoolkit-1.2.0a20250803.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
-aiqtoolkit-1.2.0a20250803.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-aiqtoolkit-1.2.0a20250803.dist-info/METADATA,sha256=5pB_ls9-qYsouughtqE9OZMND7Q92kPgGZQEuYZXGzE,21564
-aiqtoolkit-1.2.0a20250803.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aiqtoolkit-1.2.0a20250803.dist-info/entry_points.txt,sha256=iZR3yrf1liXfbcLqn5_pUkLhZyr1bUw_Qh1d2i7gsv4,625
-aiqtoolkit-1.2.0a20250803.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
-aiqtoolkit-1.2.0a20250803.dist-info/RECORD,,
+aiqtoolkit-1.2.0a20250806.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
+aiqtoolkit-1.2.0a20250806.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+aiqtoolkit-1.2.0a20250806.dist-info/METADATA,sha256=aJrMPRk-wxuTuU-u_mo3nrAoPy6OnYFAgmxk8BJTLb4,21564
+aiqtoolkit-1.2.0a20250806.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqtoolkit-1.2.0a20250806.dist-info/entry_points.txt,sha256=iZR3yrf1liXfbcLqn5_pUkLhZyr1bUw_Qh1d2i7gsv4,625
+aiqtoolkit-1.2.0a20250806.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+aiqtoolkit-1.2.0a20250806.dist-info/RECORD,,