Merge pull request #995 from guardrails-ai/feat/cli-create-install-local-models

Create CLI command to respect install local models

Author: CalebCourier

guardrails/cli/create.py

@@ -1,21 +1,14 @@
 import os
 import sys
 import time
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, cast
 
 import typer
 import json
 from rich.console import Console
 from rich.syntax import Syntax
 
 from guardrails.cli.guardrails import guardrails as gr_cli
-from guardrails.cli.hub.install import (  # JC: I don't like this import. Move fns?
-    install_hub_module,
-    add_to_hub_inits,
-    run_post_install,
-)
-from guardrails.cli.hub.utils import get_site_packages_location
-from guardrails.cli.server.hub_client import get_validator_manifest
 from guardrails.cli.hub.template import get_template
 
 console = Console()
@@ -30,6 +23,11 @@ def create_command(
     name: Optional[str] = typer.Option(
         default=None, help="The name of the guard to define in the file."
     ),
+    local_models: Optional[bool] = typer.Option(
+        None,
+        "--install-local-models/--no-install-local-models",
+        help="Install local models",
+    ),
     filepath: str = typer.Option(
         default="config.py",
         help="The path to which the configuration file should be saved.",
@@ -47,6 +45,8 @@ def create_command(
         help="Print out the validators to be installed without making any changes.",
     ),
 ):
+    # fix pyright typing issue
+    validators = cast(str, validators)
     filepath = check_filename(filepath)
 
     if not validators and template is not None:
@@ -56,7 +56,11 @@ def create_command(
             for validator in guard["validators"]:
                 validators_map[f"hub://{validator['id']}"] = True
         validators = ",".join(validators_map.keys())
-        installed_validators = split_and_install_validators(validators, dry_run)  # type: ignore
+        installed_validators = split_and_install_validators(
+            validators,
+            local_models,
+            dry_run,
+        )
         new_config_file = generate_template_config(
             template_dict, installed_validators, template_file_name
         )
@@ -67,7 +71,11 @@ def create_command(
         )
         sys.exit(1)
     else:
-        installed_validators = split_and_install_validators(validators, dry_run)  # type: ignore
+        installed_validators = split_and_install_validators(
+            validators,
+            local_models,
+            dry_run,
+        )
         if name is None and validators:
             name = "Guard"
             if len(installed_validators) > 0:
@@ -137,53 +145,52 @@ def check_filename(filename: Union[str, os.PathLike]) -> str:
     return filename  # type: ignore
 
 
-def split_and_install_validators(validators: str, dry_run: bool = False):
+def split_and_install_validators(
+    validators: str, local_models: Union[bool, None], dry_run: bool = False
+):
     """Given a comma-separated list of validators, check the hub to make sure
     all of them exist, install them, and return a list of 'imports'.
 
     If validators is empty, returns an empty list.
     """
+    from guardrails.hub.install import install
+
+    def install_local_models_confirm():
+        return typer.confirm(
+            "This validator has a Guardrails AI inference endpoint available. "
+            "Would you still like to install the"
+            " local models for local inference?",
+        )
+
     if not validators:
         return []
 
-    stripped_validators = list()
-    manifests = list()
-    site_packages = get_site_packages_location()
+    manifest_exports = list()
 
     # Split by comma, strip start and end spaces, then make sure there's a hub prefix.
     # If all that passes, download the manifest file so we know where to install.
     # hub://blah -> blah, then download the manifest.
-    console.print("Checking validators...")
-    with console.status("Checking validator manifests") as status:
-        for v in validators.split(","):
-            v = v.strip()
-            status.update(f"Prefetching {v}")
-            if not v.startswith("hub://"):
-                console.print(
-                    f"WARNING: Validator {v} does not appear to be a valid URI."
-                )
-                sys.exit(-1)
-            stripped_validator = v.lstrip("hub://")
-            stripped_validators.append(stripped_validator)
-            manifests.append(get_validator_manifest(stripped_validator))
-    console.print("Success!")
-
-    # We should make sure they exist.
     console.print("Installing...")
     with console.status("Installing validators") as status:
-        for manifest, validator in zip(manifests, stripped_validators):
-            status.update(f"Installing {validator}")
+        for v in validators.split(","):
+            validator_hub_uri = v.strip()
+            status.update(f"Installing {v}")
             if not dry_run:
-                install_hub_module(manifest, site_packages, quiet=True)
-                run_post_install(manifest, site_packages)
-                add_to_hub_inits(manifest, site_packages)
+                module = install(
+                    package_uri=validator_hub_uri,
+                    install_local_models=local_models,
+                    quiet=True,
+                    install_local_models_confirm=install_local_models_confirm,
+                )
+                exports = module.__validator_exports__
+                manifest_exports.append(exports[0])
             else:
-                console.print(f"Fake installing {validator}")
+                console.print(f"Fake installing {validator_hub_uri}")
                 time.sleep(1)
     console.print("Success!")
 
     # Pull the hub information from each of the installed validators and return it.
-    return [manifest.exports[0] for manifest in manifests]
+    return manifest_exports
 
 
 def generate_config_file(validators: List[str], name: Optional[str] = None) -> str:

guardrails/cli/hub/install.py

@@ -1,186 +1,10 @@
-import os
-import subprocess
 import sys
-from typing import List, Literal, Optional
+from typing import Optional
 
 import typer
 
-from guardrails.classes.generic import Stack
 from guardrails.cli.hub.hub import hub_command
-
-from guardrails.cli.hub.utils import (
-    get_hub_directory,
-    get_org_and_package_dirs,
-    pip_process,
-)
 from guardrails.cli.logger import logger
-from guardrails.cli.server.module_manifest import ModuleManifest
-
-
-def removesuffix(string: str, suffix: str) -> str:
-    if sys.version_info.minor >= 9:
-        return string.removesuffix(suffix)  # type: ignore
-    else:
-        if string.endswith(suffix):
-            return string[: -len(suffix)]
-        return string
-
-
-string_format: Literal["string"] = "string"
-json_format: Literal["json"] = "json"
-
-
-# NOTE: I don't like this but don't see another way without
-#  shimming the init file with all hub validators
-def add_to_hub_inits(manifest: ModuleManifest, site_packages: str):
-    org_package = get_org_and_package_dirs(manifest)
-    exports: List[str] = manifest.exports or []
-    sorted_exports = sorted(exports, reverse=True)
-    module_name = manifest.module_name
-    relative_path = ".".join([*org_package, module_name])
-    import_line = (
-        f"from guardrails.hub.{relative_path} import {', '.join(sorted_exports)}"
-    )
-
-    hub_init_location = os.path.join(site_packages, "guardrails", "hub", "__init__.py")
-    with open(hub_init_location, "a+") as hub_init:
-        hub_init.seek(0, 0)
-        content = hub_init.read()
-        if import_line in content:
-            hub_init.close()
-        else:
-            hub_init.seek(0, 2)
-            if len(content) > 0:
-                hub_init.write("\n")
-            hub_init.write(import_line)
-            hub_init.close()
-
-    namespace = org_package[0]
-    namespace_init_location = os.path.join(
-        site_packages, "guardrails", "hub", namespace, "__init__.py"
-    )
-    if os.path.isfile(namespace_init_location):
-        with open(namespace_init_location, "a+") as namespace_init:
-            namespace_init.seek(0, 0)
-            content = namespace_init.read()
-            if import_line in content:
-                namespace_init.close()
-            else:
-                namespace_init.seek(0, 2)
-                if len(content) > 0:
-                    namespace_init.write("\n")
-                namespace_init.write(import_line)
-                namespace_init.close()
-    else:
-        with open(namespace_init_location, "w") as namespace_init:
-            namespace_init.write(import_line)
-            namespace_init.close()
-
-
-def run_post_install(manifest: ModuleManifest, site_packages: str):
-    org_package = get_org_and_package_dirs(manifest)
-    post_install_script = manifest.post_install
-    if not post_install_script:
-        return
-
-    module_name = manifest.module_name
-    relative_path = os.path.join(
-        site_packages,
-        "guardrails",
-        "hub",
-        *org_package,
-        module_name,
-        post_install_script,
-    )
-
-    if os.path.isfile(relative_path):
-        try:
-            logger.debug("running post install script...")
-            command = [sys.executable, relative_path]
-            subprocess.check_output(command)
-        except subprocess.CalledProcessError as exc:
-            logger.error(
-                (
-                    f"Failed to run post install script for {manifest.id}\n"
-                    f"Exit code: {exc.returncode}\n"
-                    f"stdout: {exc.output}"
-                )
-            )
-            sys.exit(1)
-        except Exception as e:
-            logger.error(
-                f"An unexpected exception occurred while running the post install script for {manifest.id}!",  # noqa
-                e,
-            )
-            sys.exit(1)
-
-
-def get_install_url(manifest: ModuleManifest) -> str:
-    repo = manifest.repository
-    repo_url = repo.url
-    branch = repo.branch
-
-    git_url = repo_url
-    if not repo_url.startswith("git+"):
-        git_url = f"git+{repo_url}"
-
-    if branch is not None:
-        git_url = f"{git_url}@{branch}"
-
-    return git_url
-
-
-def install_hub_module(
-    module_manifest: ModuleManifest, site_packages: str, quiet: bool = False
-):
-    install_url = get_install_url(module_manifest)
-    install_directory = get_hub_directory(module_manifest, site_packages)
-
-    pip_flags = [f"--target={install_directory}", "--no-deps"]
-    if quiet:
-        pip_flags.append("-q")
-
-    # Install validator module in namespaced directory under guardrails.hub
-    download_output = pip_process("install", install_url, pip_flags, quiet=quiet)
-    if not quiet:
-        logger.info(download_output)
-
-    # Install validator module's dependencies in normal site-packages directory
-    inspect_output = pip_process(
-        "inspect",
-        flags=[f"--path={install_directory}"],
-        format=json_format,
-        quiet=quiet,
-        no_color=True,
-    )
-
-    # throw if inspect_output is a string. Mostly for pyright
-    if isinstance(inspect_output, str):
-        logger.error("Failed to inspect the installed package!")
-        sys.exit(1)
-
-    dependencies = (
-        Stack(*inspect_output.get("installed", []))
-        .at(0, {})
-        .get("metadata", {})  # type: ignore
-        .get("requires_dist", [])  # type: ignore
-    )
-    requirements = list(filter(lambda dep: "extra" not in dep, dependencies))
-    for req in requirements:
-        if "git+" in req:
-            install_spec = req.replace(" ", "")
-            dep_install_output = pip_process("install", install_spec, quiet=quiet)
-            if not quiet:
-                logger.info(dep_install_output)
-        else:
-            req_info = Stack(*req.split(" "))
-            name = req_info.at(0, "").strip()  # type: ignore
-            versions = req_info.at(1, "").strip("()")  # type: ignore
-            if name:
-                install_spec = name if not versions else f"{name}{versions}"
-                dep_install_output = pip_process("install", install_spec, quiet=quiet)
-                if not quiet:
-                    logger.info(dep_install_output)
 
 
 @hub_command.command()

guardrails/hub/install.py

@@ -1,8 +1,11 @@
 from contextlib import contextmanager
 from string import Template
-from typing import Callable
+from typing import Callable, cast
 
-from guardrails.hub.validator_package_service import ValidatorPackageService
+from guardrails.hub.validator_package_service import (
+    ValidatorPackageService,
+    ValidatorModuleType,
+)
 from guardrails.classes.credentials import Credentials
 
 from guardrails.cli.hub.console import console
@@ -33,7 +36,7 @@ def install(
     install_local_models=None,
     quiet: bool = True,
     install_local_models_confirm: Callable = default_local_models_confirm,
-):
+) -> ValidatorModuleType:
     """
     Install a validator package from a hub URI.
 
@@ -49,6 +52,9 @@ def install(
 
     Examples:
         >>> RegexMatch = install("hub://guardrails/regex_match").RegexMatch
+
+        >>> install("hub://guardrails/regex_match);
+        >>> import guardrails.hub.regex_match as regex_match
     """
 
     verbose_printer = console.print
@@ -122,7 +128,10 @@ def install(
     ValidatorPackageService.add_to_hub_inits(module_manifest, site_packages)
 
     # 5. Get Validator Class for the installed module
-    validators = ValidatorPackageService.get_validator_from_manifest(module_manifest)
+    installed_module = ValidatorPackageService.get_validator_from_manifest(
+        module_manifest
+    )
+    installed_module = cast(ValidatorModuleType, installed_module)
 
     # Print success messages
     cli_logger.info("Installation complete")
@@ -152,4 +161,7 @@ def install(
     quiet_printer(success_message_cli)  # type: ignore
     cli_logger.log(level=LEVELS.get("SPAM"), msg=success_message_logger)  # type: ignore
 
-    return validators
+    # Not a fan of this but allows the installation to be used in create command as is
+    installed_module.__validator_exports__ = module_manifest.exports
+
+    return installed_module