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

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
@@ -175,8 +196,8 @@ def discover_and_register_plugins(plugin_type: PluginTypes):
                 # Log a warning if the plugin took a long time to load. This can be useful for debugging slow imports.
                 # The threshold is 300 ms if no plugins have been loaded yet, and 100 ms otherwise. Triple the threshold
                 # if a debugger is attached.
-                if (elapsed_time > (300.0 if count == 0 else 100.0) * (3 if is_debugger_attached() else 1)):
-                    logger.warning(
+                if (elapsed_time > (300.0 if count == 0 else 150.0) * (3 if is_debugger_attached() else 1)):
+                    logger.debug(
                         "Loading module '%s' from entry point '%s' took a long time (%f ms). "
                         "Ensure all imports are inside your registered functions.",
                         entry_point.module,

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/tool/server_tools.py

@@ -63,4 +63,4 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
                 f"Conversation Id: {conversation_id}")
 
     yield FunctionInfo.from_fn(_get_request_attributes,
-                               description="Returns the acquired user defined request attriubutes.")
+                               description="Returns the acquired user defined request attributes.")

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)