`_check_library` fails for packages with mismatched install/import names (e.g., google-genai)

[kpeez]

Hello! I am attempting to use the GeminiEmbeddingFunction from pymilvus[model] but I am running into an issue where the dependency check is failing, causing the package to be reinstalled each time I run the script. I believe this is due to the _check_library function in pymilvus/model/utils/__init__.py not being able to find the google-genai package due to a mismatch between the install/import names. I've included a proposed fix using importlib.metadata below. Let me know if you think this approach, makes sense; I'm happy to submit a PR if you think it's a good solution!

Description

The utility function pymilvus.model.utils._check_library attempts to determine if a required Python package is installed before potentially prompting the user to install it via prompt_install. However, its current implementation has issues with packages where the name used for installation (e.g., google-genai) differs significantly from the name used for importing (e.g., google.generativeai).

The core problems are:

1. importlib.util.find_spec Usage: _check_library currently relies on importlib.util.find_spec(libname).
   • If libname is the installable name (like "google-genai" as used in import_google), find_spec correctly returns None because you cannot import google-genai.
   • If libname were changed to the importable name (like "google.generativeai"), find_spec can still return None even if the package is installed, especially for modules nested within namespace packages like google, before they have been explicitly imported elsewhere. find_spec is not always reliable for checking installation status this way.
2. prompt_install Behavior: The pymilvus.model.utils.dependency_control.prompt_install function currently does not check if the package is already installed. It directly proceeds to run pip install.

Combined Effect: When import_google calls _check_library("google-genai", package="google-genai>=1.7.0"), the find_spec("google-genai") check fails (as expected), leading directly to prompt_install("google-genai>=1.7.0"), which then attempts to run pip install google-genai... every single time, even if the package is already present. This is inefficient and unnecessary.

Current Code Snippets

pymilvus/model/utils/__init__.py (Relevant parts):

import importlib.util
from typing import Optional
from pymilvus.model.utils.dependency_control import prompt_install

# Example problematic caller
def import_google():
    _check_library("google-genai", package="google-genai>=1.7.0")

# The problematic check function (original version)
def _check_library(libname: str, prompt: bool = True, package: Optional[str] = None):
    is_avail = False
    # This check fails for "google-genai" even if installed
    if importlib.util.find_spec(libname):
        is_avail = True
    # This is always triggered for google-genai
    if not is_avail and prompt:
        prompt_install(package if package else libname)
    return is_avail

pymilvus/model/utils/dependency_control.py:

import subprocess

def prompt_install(package: str, warn: bool = False):  # pragma: no cover
    cmd = f"pip install -q {package}"
    try:
        # No check here to see if package is already installed
        print(f"start to install package: {package}")
        subprocess.check_call(cmd, shell=True)
        print(f"successfully installed package: {package}")
    except subprocess.CalledProcessError as e:
        raise ValueError(f"install error {e}")

Proposed Solution: Use importlib.metadata

The standard and most reliable way to check if a distribution package is installed is using importlib.metadata. The _check_library function should be updated to use importlib.metadata.version() within a try...except PackageNotFoundError block.

The name to check should be derived from the package argument (stripping version specifiers) if provided, otherwise fallback to libname. This ensures we check for the distribution package name.

Recommended _check_library implementation:

import importlib.metadata
from typing import Optional
from pymilvus.model.utils.dependency_control import prompt_install

# ... other import functions ...

def _check_library(libname: str, prompt: bool = True, package: Optional[str] = None):
    # Determine the distribution package name to check with importlib.metadata
    # Use the 'package' name if provided (stripping version specifiers), otherwise 'libname'
    check_name = libname
    if package:
        # Basic parsing to remove version specifiers like >=, ==, <, etc.
        specifiers = ['>=', '==', '<=', '<', '!=', '~='] # Added ~=
        for spec in specifiers:
            if spec in package:
                check_name = package.split(spec)[0].strip() 
                break
        else:
            check_name = package.strip()
    else:
      check_name = libname.strip()


    is_avail = False
    try:
        importlib.metadata.version(check_name)
        is_avail = True
    except importlib.metadata.PackageNotFoundError:
        is_avail = False

    if not is_avail and prompt:
        install_target = package if package else libname
        prompt_install(install_target)

    return is_avail

Benefits of the Fix

• Correctness: Accurately detects if the required distribution package is installed, regardless of import name differences or namespace package complexities.
• Efficiency: Prevents redundant and time-consuming pip install calls when the package is already present.