aiqtoolkit 1.2.0rc3__py3-none-any.whl → 1.2.0rc4__py3-none-any.whl

aiq/authentication/api_key/api_key_auth_provider.py

@@ -28,9 +28,13 @@ logger = logging.getLogger(__name__)
 
 class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
-    def __init__(self, config: APIKeyAuthProviderConfig, config_name: str | None = None) -> None:
-        assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyConfig")
+    # fmt: off
+    def __init__(self,
+                 config: APIKeyAuthProviderConfig,
+                 config_name: str | None = None) -> None:  # pylint: disable=unused-argument
+        assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyAuthProviderConfig")
         super().__init__(config)
+    # fmt: on
 
     async def _construct_authentication_header(self) -> BearerTokenCred:
         """

aiq/data_models/component_ref.py

@@ -130,7 +130,7 @@ class ObjectStoreRef(ComponentRef):
     """
 
     @property
-    @typing.override
+    @override
     def component_group(self):
         return ComponentGroup.OBJECT_STORES
 

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/object_store/in_memory_object_store.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
+
 from aiq.builder.builder import Builder
 from aiq.cli.register_workflow import register_object_store
 from aiq.data_models.object_store import KeyAlreadyExistsError
@@ -37,37 +39,37 @@ class InMemoryObjectStore(ObjectStore):
     """
 
     def __init__(self) -> None:
+        self._lock = asyncio.Lock()
         self._store: dict[str, ObjectStoreItem] = {}
 
     @override
     async def put_object(self, key: str, item: ObjectStoreItem) -> None:
-        if key in self._store:
-            raise KeyAlreadyExistsError(key)
-
-        self._store[key] = item
-        return
+        async with self._lock:
+            if key in self._store:
+                raise KeyAlreadyExistsError(key)
+            self._store[key] = item
 
     @override
     async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
-        self._store[key] = item
-        return
+        async with self._lock:
+            self._store[key] = item
 
     @override
     async def get_object(self, key: str) -> ObjectStoreItem:
-
-        if key not in self._store:
-            raise NoSuchKeyError(key)
-
-        return self._store[key]
+        async with self._lock:
+            value = self._store.get(key)
+            if value is None:
+                raise NoSuchKeyError(key)
+            return value
 
     @override
     async def delete_object(self, key: str) -> None:
-        if key not in self._store:
+        try:
+            async with self._lock:
+                self._store.pop(key)
+        except KeyError:
             raise NoSuchKeyError(key)
 
-        self._store.pop(key)
-        return
-
 
 @register_object_store(config_type=InMemoryObjectStoreConfig)
 async def in_memory_object_store(config: InMemoryObjectStoreConfig, builder: Builder):

aiq/registry_handlers/package_utils.py

@@ -20,6 +20,8 @@ import os
 import subprocess
 from functools import lru_cache
 
+from packaging.requirements import Requirement
+
 from aiq.data_models.component import AIQComponentEnum
 from aiq.data_models.discovery_metadata import DiscoveryMetadata
 from aiq.registry_handlers.schemas.package import WheelData
@@ -57,6 +59,372 @@ def get_module_name_from_distribution(distro_name: str) -> str | None:
     return None
 
 
+def parse_requirement(requirement: str) -> str:
+    """Extract the base package name from a requirement string.
+
+    This function extracts only the package name, ignoring extras, version specifiers,
+    and environment markers.
+
+    Args:
+        requirement (str): A requirement string like 'numpy>=1.20.0' or 'requests[security]~=2.28.0'
+
+    Returns:
+        str: The base package name (e.g., 'numpy' from 'numpy>=1.20.0',
+        'requests' from 'requests[security]~=2.28.0')
+    """
+    # Handle inline comments by splitting on '#' and taking the first part
+    clean_requirement = requirement.split('#')[0].strip()
+    if not clean_requirement:
+        return ""
+
+    try:
+        parsed = Requirement(clean_requirement)
+        return parsed.name.lower()
+    except Exception as e:
+        logger.warning("Failed to parse requirement '%s': %s. Skipping this dependency.", requirement, e)
+        return ""
+
+
+def resolve_extras_to_packages(package_name: str, extras: list[str]) -> set[str]:
+    """Resolve package extras to their actual package dependencies.
+
+    Args:
+        package_name (str): The base package name (e.g., 'aiqtoolkit')
+        extras (list[str]): List of extra names (e.g., ['langchain', 'telemetry'])
+
+    Returns:
+        set[str]: Set of additional package names that the extras resolve to
+        (e.g., {'aiqtoolkit-langchain', 'aiqtoolkit-opentelemetry', 'aiqtoolkit-phoenix',
+        'aiqtoolkit-weave', 'aiqtoolkit-ragaai'})
+    """
+    resolved_packages = set()
+
+    try:
+        # Get the distribution metadata for the package
+        dist = importlib.metadata.distribution(package_name)
+
+        # Parse all requirements to find optional dependencies
+        requires = dist.requires or []
+
+        for requirement_str in requires:
+            try:
+                req = Requirement(requirement_str)
+
+                # Check if this requirement has a marker that matches our extras
+                if req.marker:
+                    for extra in extras:
+                        # Try marker evaluation first
+                        try:
+                            if req.marker.evaluate({'extra': extra}):
+                                resolved_packages.add(req.name.lower())
+                                break
+                        except Exception:
+                            # Fallback to simple string check
+                            marker_str = str(req.marker)
+                            if f'extra == "{extra}"' in marker_str or f"extra == '{extra}'" in marker_str:
+                                resolved_packages.add(req.name.lower())
+                                break
+
+            except Exception as e:
+                logger.warning("Failed to parse requirement '%s' for extras resolution: %s", requirement_str, e)
+
+    except importlib.metadata.PackageNotFoundError:
+        logger.warning("Package '%s' not found for extras resolution", package_name)
+    except Exception as e:
+        logger.warning("Failed to resolve extras for package '%s': %s", package_name, e)
+
+    return resolved_packages
+
+
+def extract_dependencies_with_extras_resolved(pyproject_path: str) -> set[str]:
+    """Extract dependency names from pyproject.toml with extras properly resolved.
+
+    This function not only extracts the base package names but also resolves
+    any extras (e.g., package[extra1,extra2]) to their actual package dependencies.
+
+    Args:
+        pyproject_path (str): Path to the pyproject.toml file
+
+    Returns:
+        set[str]: Set of all dependency names including those resolved from extras
+
+    Example:
+        For a dependency like "aiqtoolkit[langchain,telemetry]~=1.2", this will return:
+        {'aiqtoolkit', 'aiqtoolkit-langchain', 'aiqtoolkit-opentelemetry', 'aiqtoolkit-phoenix', ...}
+
+    Raises:
+        FileNotFoundError: If the pyproject.toml file doesn't exist
+        ValueError: If the file cannot be parsed
+    """
+    import tomllib
+
+    if not os.path.exists(pyproject_path):
+        raise FileNotFoundError(f"pyproject.toml not found at {pyproject_path}")
+
+    try:
+        with open(pyproject_path, "rb") as f:
+            data = tomllib.load(f)
+    except Exception as e:
+        raise ValueError(f"Failed to parse pyproject.toml: {e}") from e
+
+    project_data = data.get("project", {})
+    all_dependencies = set()
+
+    def _process_dependency(dep_spec: str):
+        """Process a single dependency specification and resolve extras."""
+        # Handle inline comments
+        clean_req = dep_spec.split('#')[0].strip()
+        if not clean_req:
+            return
+
+        try:
+            parsed = Requirement(clean_req)
+            base_name = parsed.name.lower()
+            all_dependencies.add(base_name)
+
+            # If there are extras, try to resolve them
+            if parsed.extras:
+                resolved_extras = resolve_extras_to_packages(base_name, list(parsed.extras))
+                all_dependencies.update(resolved_extras)
+
+        except Exception as e:
+            logger.warning("Failed to process dependency '%s': %s", dep_spec, e)
+
+    # Process main dependencies
+    for dep_spec in project_data.get("dependencies", []):
+        _process_dependency(dep_spec)
+
+    # Process optional dependencies
+    optional_deps = project_data.get("optional-dependencies", {})
+    for _group_name, group_deps in optional_deps.items():
+        for dep_spec in group_deps:
+            _process_dependency(dep_spec)
+
+    return all_dependencies
+
+
+@lru_cache
+def get_distributions() -> list[importlib.metadata.Distribution]:
+    """Get all installed distributions. This is an expensive operation and should be cached."""
+    return list(importlib.metadata.distributions())
+
+
+def find_distribution_name(name: str) -> str | None:
+    """Try to find the correct distribution name for a given package name.
+
+    Uses dynamic discovery through importlib.metadata to find distributions
+    that provide the requested module/package name.
+
+    Args:
+        name (str): Package name to search for.
+
+    Returns:
+        str | None: The correct distribution name if found, None otherwise.
+    """
+    # First try the name as-is
+    try:
+        importlib.metadata.distribution(name)
+        return name
+    except importlib.metadata.PackageNotFoundError:
+        pass
+
+    # Try common case variations
+    variations = [
+        name.lower(),
+        name.upper(),
+        name.replace('-', '_'),
+        name.replace('_', '-'),
+    ]
+
+    # Try each variation
+    for variation in variations:
+        if variation != name:  # Skip the original name we already tried
+            try:
+                importlib.metadata.distribution(variation)
+                return variation
+            except importlib.metadata.PackageNotFoundError:
+                continue
+
+    # Search through all installed distributions to find one that provides this module
+    try:
+        for dist in get_distributions():
+            dist_name = dist.metadata['Name']
+
+            # Check top-level packages provided by this distribution
+            try:
+                # Try to get top-level packages from metadata
+                top_level_txt = dist.read_text('top_level.txt')
+                if top_level_txt:
+                    top_level_packages = set(top_level_txt.strip().split('\n'))
+                    if name in top_level_packages:
+                        return dist_name
+            except (FileNotFoundError, AttributeError):
+                # top_level.txt doesn't exist, try alternative method
+                pass
+
+            # Fallback: check file paths for top-level modules
+            try:
+                if hasattr(dist, 'files') and dist.files:
+                    top_level_from_files = {
+                        f.parts[0]
+                        for f in dist.files if len(f.parts) > 0 and not f.parts[0].endswith('.dist-info')
+                    }
+                    if name in top_level_from_files:
+                        return dist_name
+            except Exception:
+                # Some distributions might not have files info or it might be inaccessible
+                continue
+
+    except Exception as e:
+        logger.debug("Error searching distributions for %s: %s", name, e)
+
+    return None
+
+
+def get_transitive_dependencies(distribution_names: list[str]) -> dict[str, set[str]]:
+    """Get transitive dependencies from a list of Python distribution names.
+
+    This function recursively resolves all dependencies for the given distribution names,
+    returning a mapping of each package to its complete set of transitive dependencies.
+    This is useful when publishing plugins to remote registries that contain with nested dependencies,
+    ensuring that all dependencies are included in the AIQArtifact's metadata.
+
+    Args:
+        distribution_names (list[str]): List of Python distribution names (package names) to analyze.
+
+    Returns:
+        dict[str, set[str]]: Dictionary mapping each distribution name to its set of transitive dependencies.
+        The dependencies include both direct and indirect dependencies.
+    """
+    result: dict[str, set[str]] = {}
+    processing: set[str] = set()  # Track packages currently being processed (cycle detection)
+    completed: set[str] = set()  # Track packages that have been fully processed
+
+    def _get_dependencies_recursive(dist_name: str, path: set[str]) -> set[str]:
+        """Recursively get all dependencies for a distribution.
+
+        Args:
+            dist_name: The distribution name to process
+            path: Set of packages in the current dependency path (for cycle detection)
+        """
+        # If we've already computed this package's dependencies, return them
+        if dist_name in completed:
+            return result.get(dist_name, set())
+
+        # If we encounter this package in the current path, we have a cycle
+        if dist_name in path:
+            logger.debug("Cycle detected in dependency chain: %s", " -> ".join(list(path) + [dist_name]))
+            return set()
+
+        # If we're currently processing this package in another branch, return empty
+        # to avoid duplicate work (we'll get the full result when that branch completes)
+        if dist_name in processing:
+            return set()
+
+        processing.add(dist_name)
+        new_path = path | {dist_name}
+        dependencies = set()
+
+        try:
+            dist = importlib.metadata.distribution(dist_name)
+            requires = dist.requires or []
+
+            for requirement in requires:
+                # Skip requirements with extra markers (optional dependencies)
+                # These should only be included if the extra is explicitly requested
+                if 'extra ==' in requirement:
+                    continue
+
+                # Parse the requirement to get the package name
+                dep_name = parse_requirement(requirement)
+
+                # Skip self-references and empty names
+                if not dep_name or dep_name == dist_name.lower():
+                    continue
+
+                dependencies.add(dep_name)
+
+                # Recursively get dependencies of this dependency
+                try:
+                    transitive_deps = _get_dependencies_recursive(dep_name, new_path)
+                    dependencies.update(transitive_deps)
+                except importlib.metadata.PackageNotFoundError:
+                    # Check if this is likely a conditional dependency (has markers)
+                    is_conditional = any(marker in requirement for marker in [
+                        'python_version', 'sys_platform', 'platform_system', 'platform_machine', 'implementation_name',
+                        'implementation_version'
+                    ])
+
+                    if is_conditional:
+                        # This is expected - conditional dependencies aren't always installed
+                        logger.debug("Conditional dependency %s of %s is not installed: %s",
+                                     dep_name,
+                                     dist_name,
+                                     requirement)
+                    else:
+                        # This might be a real issue - a non-conditional dependency is missing
+                        logger.warning("Dependency %s of %s is not installed", dep_name, dist_name)
+                    continue
+
+        except importlib.metadata.PackageNotFoundError:
+            # Transitive dependencies that aren't found are usually conditional (platform/version specific)
+            # and this is expected behavior
+            logger.debug("Distribution %s not found (likely conditional dependency)", dist_name)
+            # Don't raise - just return empty dependencies for missing distributions
+        finally:
+            processing.remove(dist_name)
+
+        result[dist_name] = dependencies
+        completed.add(dist_name)
+        return dependencies
+
+    # Process each distribution name
+    for dist_name in distribution_names:
+        if dist_name not in completed:
+            try:
+                _get_dependencies_recursive(dist_name.lower(), set())
+            except importlib.metadata.PackageNotFoundError:
+                # Try to find the correct distribution name
+                correct_name = find_distribution_name(dist_name)
+                if correct_name:
+                    logger.debug("Found distribution '%s' for requested name '%s'", correct_name, dist_name)
+                    try:
+                        _get_dependencies_recursive(correct_name.lower(), set())
+                        # Map the original name to the results of the correct name
+                        if correct_name.lower() in result:
+                            result[dist_name] = result[correct_name.lower()]
+                        continue
+                    except importlib.metadata.PackageNotFoundError:
+                        pass
+
+                logger.error("Distribution %s not found (tried common variations)", dist_name)
+                result[dist_name] = set()
+
+    return result
+
+
+def get_all_transitive_dependencies(distribution_names: list[str]) -> set[str]:
+    """Get all unique transitive dependencies from a list of Python distribution names.
+
+    Returns a flattened set of all unique dependencies across all the provided distribution names.
+    This is useful when publishing plugins to remote registries that contain with nested dependencies,
+    ensuring that all dependencies are included in the AIQArtifact's metadata.
+
+    Args:
+        distribution_names: List of Python distribution names (package names) to analyze
+
+    Returns:
+        set[str]: Set of all unique transitive dependency names
+    """
+    deps_map = get_transitive_dependencies(distribution_names)
+    all_deps = set()
+
+    for deps in deps_map.values():
+        all_deps.update(deps)
+
+    return all_deps
+
+
 def build_wheel(package_root: str) -> WheelData:
     """Builds a Python .whl for the specified package and saves to disk, sets self._whl_path, and returned as bytes.
 
@@ -67,8 +435,6 @@ def build_wheel(package_root: str) -> WheelData:
         WheelData: Data model containing a built python wheel and its corresponding metadata.
     """
 
-    import importlib.util
-    import re
     import tomllib
 
     from pkginfo import Wheel
@@ -83,22 +449,27 @@ def build_wheel(package_root: str) -> WheelData:
 
     toml_project: dict = data.get("project", {})
     toml_project_name = toml_project.get("name", None)
-
-    assert toml_project_name is not None, f"Package name '{toml_project_name}' not found in pyproject.toml"
-    # replace "aiqtoolkit" substring with "aiq" to get the import name
-    module_name = get_module_name_from_distribution(toml_project_name)
-    assert module_name is not None, f"No modules found for package name '{toml_project_name}'"
-
-    assert importlib.util.find_spec(module_name) is not None, (f"Package {module_name} not "
-                                                               "installed, cannot discover components.")
-
     toml_packages = set(i for i in data.get("project", {}).get("entry-points", {}).get("aiq.plugins", {}))
-    toml_dependencies = set(
-        re.search(r"[a-zA-Z][a-zA-Z\d_-]*", package_name).group(0)
-        for package_name in toml_project.get("dependencies", []))
 
+    # Extract dependencies using the robust requirement parser with extras resolution
+    try:
+        toml_dependencies = extract_dependencies_with_extras_resolved(pyproject_toml_path)
+        logger.debug("Extracted dependencies with extras resolved: %s", toml_dependencies)
+    except Exception as e:
+        logger.warning("Failed to extract dependencies with extras resolution, falling back to basic extraction: %s", e)
+        # Fallback to basic extraction
+        toml_dependencies = set()
+        for dep_spec in toml_project.get("dependencies", []):
+            try:
+                dep_name = parse_requirement(dep_spec)
+                if dep_name:
+                    toml_dependencies.add(dep_name)
+            except Exception as e:
+                logger.warning("Failed to parse dependency '%s': %s", dep_spec, e)
+
+    toml_dependencies_transitive = get_all_transitive_dependencies(list(toml_dependencies))
     union_dependencies = toml_dependencies.union(toml_packages)
-    union_dependencies.add(toml_project_name)
+    union_dependencies.update(toml_dependencies_transitive)
 
     working_dir = os.getcwd()
     os.chdir(package_root)
@@ -119,16 +490,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 +525,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 +533,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/tool/code_execution/README.md

@@ -23,7 +23,6 @@ A secure, containerized Python code execution environment that allows safe execu
 
 The Code Execution Sandbox provides:
 - **Secure code execution** in isolated Docker containers
-- **Comprehensive error handling** with detailed stdout/stderr capture
 - **Multiple input formats** including raw code, dictionary format, and markdown
 - **Dependency management** with pre-installed libraries
 - **Flexible configuration** with customizable timeouts and output limits

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.0rc3.dist-info → aiqtoolkit-1.2.0rc4.dist-info}/METADATA

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

{aiqtoolkit-1.2.0rc3.dist-info → aiqtoolkit-1.2.0rc4.dist-info}/RECORD

@@ -20,7 +20,7 @@ aiq/authentication/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7
 aiq/authentication/interfaces.py,sha256=jfsbVx0MTrxZonyTEH0YFcSJoyFHVkfVHFhm9v3jMrY,3317
 aiq/authentication/register.py,sha256=dYChd2HRv-uv4m8Ebo2sLfbVnksT8D3jEWA-QT-MzX0,947
 aiq/authentication/api_key/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-aiq/authentication/api_key/api_key_auth_provider.py,sha256=08YhKWSS0t5ov1sagLC9DN8xM2Asep8eu5m5yeAzzj8,3947
+aiq/authentication/api_key/api_key_auth_provider.py,sha256=ZclKuexJShF8QVJ2vwhePk2q-VSDsJB4iZn2dkOCK4I,4057
 aiq/authentication/api_key/api_key_auth_provider_config.py,sha256=U9Rx93XVfEFXzT_fikkE0b3bX4eP8VRB2b52_YjAFjY,5472
 aiq/authentication/api_key/register.py,sha256=2W7GMWNNa_cZvB1uU4an5kNBrYVBY-kX0tgZ9daAF6Y,1147
 aiq/authentication/exceptions/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
@@ -100,10 +100,10 @@ aiq/data_models/api_server.py,sha256=dHlRq_jI8N6dFHiKVJ_y4mMOHKkKP9TPM_A6oc0GSCQ
 aiq/data_models/authentication.py,sha256=TShP47-z9vIrkV3S4sadm0QF7YUkEWtsToKSBiHS3NI,7349
 aiq/data_models/common.py,sha256=y_8AiWmTEaMjCMayVaFYddhv2AAou8Pr84isHgGxeUg,5874
 aiq/data_models/component.py,sha256=_HeHlDmz2fgJlVfmz0tG_yCyes86hQNaQwSWz1FDQvk,1737
-aiq/data_models/component_ref.py,sha256=_6DeNWJSHN3GSoAFe0RCQkCJ31_WA9GR3nP6vsxMj7Q,4594
+aiq/data_models/component_ref.py,sha256=LJbi5NzGQuq8QRhNGzAKNK1zfxvYTycELFiBy1DKNeM,4587
 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
@@ -250,7 +250,7 @@ aiq/memory/models.py,sha256=c5dA7nKHQ4AS1_ptQZcfC_oXO495-ehocnf_qXTE6c8,4319
 aiq/meta/module_to_distro.json,sha256=1XV7edobFrdDKvsSoynfodXg_hczUWpDrQzGkW9qqEs,28
 aiq/meta/pypi.md,sha256=N1fvWaio3KhnAw9yigeM-oWaLuT5i_C7U_2UVzyPbks,4386
 aiq/object_store/__init__.py,sha256=7JInpxFZUdqigaBaXDzUj0KTlv_2tiM-SuTOlSvadkg,847
-aiq/object_store/in_memory_object_store.py,sha256=_k_5V3sOZh-S3UtsVaKWPnQaeQc7eKY0-Ui6Tx2uXW4,2453
+aiq/object_store/in_memory_object_store.py,sha256=iMwy_JaA3uZkXL6K5ToaOtarC7UfmBuvRoqmLbuYieo,2644
 aiq/object_store/interfaces.py,sha256=5NbsE9TccihOf5ScG04hE1eNOaiajOZIUOeK_Kvukk8,2519
 aiq/object_store/models.py,sha256=yAEa7oZgax7OwObynv0MFCMo43xTGps5xz3SIY6-0VM,1425
 aiq/object_store/register.py,sha256=M93V17RxBXzGZaAmWboDw-S2XP7lrMEyzdlxXqC0f30,788
@@ -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=KZYtGNEIaG9pyG3gyHyumL6w8jsNF8iD-FWNWgP4oBY,22523
 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
@@ -370,7 +370,7 @@ aiq/tool/nvidia_rag.py,sha256=9mS3igONo1RywxXNj_ITh2-qD91x1R0f7uhOWMZQX3o,4178
 aiq/tool/register.py,sha256=Lwl6l_eEzS8LAELxClmniNhhLluRVZFYXhsk2ocQhNg,1491
 aiq/tool/retriever.py,sha256=DnuU4khpJkd4epDBGQsowDOqDBKFiLQrnyKXgU6IRW8,3724
 aiq/tool/server_tools.py,sha256=286hTIvX_8pEohmkqmgWTCLkTfMTjuR9v0zivcW17r4,3200
-aiq/tool/code_execution/README.md,sha256=GQy3pPOVspFHQdclQCCweIAY2r8rVZ6bXyCY2FcEc6M,4090
+aiq/tool/code_execution/README.md,sha256=QM8Fe8O7a9t2oDbG2d1PTIbddpGM4hhlIg2ZTYUKHY0,4019
 aiq/tool/code_execution/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/tool/code_execution/code_sandbox.py,sha256=ibe6BoYdWzmQWHdVNweVP-QW1WRGirfhvFr0iqm4fI8,10159
 aiq/tool/code_execution/register.py,sha256=i3MNxCb3lBeeSKMwQeg6X2oUpgUKdEoOw2bqUEIqy1E,3322
@@ -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.0rc3.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
-aiqtoolkit-1.2.0rc3.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-aiqtoolkit-1.2.0rc3.dist-info/METADATA,sha256=X2SwqPxRdjMtrQ0l2Q9Gg439pHoduJja350zK9gCLc0,21558
-aiqtoolkit-1.2.0rc3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aiqtoolkit-1.2.0rc3.dist-info/entry_points.txt,sha256=iZR3yrf1liXfbcLqn5_pUkLhZyr1bUw_Qh1d2i7gsv4,625
-aiqtoolkit-1.2.0rc3.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
-aiqtoolkit-1.2.0rc3.dist-info/RECORD,,
+aiqtoolkit-1.2.0rc4.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
+aiqtoolkit-1.2.0rc4.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+aiqtoolkit-1.2.0rc4.dist-info/METADATA,sha256=4mgm80zLDxo_rJhad7-nJNv--TYpcFJv5UO3eUDzUZM,21558
+aiqtoolkit-1.2.0rc4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqtoolkit-1.2.0rc4.dist-info/entry_points.txt,sha256=iZR3yrf1liXfbcLqn5_pUkLhZyr1bUw_Qh1d2i7gsv4,625
+aiqtoolkit-1.2.0rc4.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+aiqtoolkit-1.2.0rc4.dist-info/RECORD,,