diff --git a/guardrails/cli/hub/utils.py b/guardrails/cli/hub/utils.py index c3f1b924a..84c085c1e 100644 --- a/guardrails/cli/hub/utils.py +++ b/guardrails/cli/hub/utils.py @@ -20,6 +20,87 @@ string_format = "string" +class PipProcessError(Exception): + action: str + package: str + stderr: str = "" + stdout: str = "" + returncode: int = 1 + + def __init__( + self, + action: str, + package: str, + stderr: str = "", + stdout: str = "", + returncode: int = 1, + ): + self.action = action + self.package = package + self.stderr = stderr + self.stdout = stdout + self.returncode = returncode + message = ( + f"PipProcessError: {action} on '{package}' failed with" + "return code {returncode}.\n" + f"Stdout:\n{stdout}\n" + f"Stderr:\n{stderr}" + ) + super().__init__(message) + + +def pip_process_with_custom_exception( + action: str, + package: str = "", + flags: List[str] = [], + format: Union[Literal["string"], Literal["json"]] = string_format, + quiet: bool = False, + no_color: bool = False, +) -> Union[str, dict]: + try: + if not quiet: + logger.debug(f"running pip {action} {' '.join(flags)} {package}") + command = [sys.executable, "-m", "pip", action] + command.extend(flags) + if package: + command.append(package) + + env = dict(os.environ) + if no_color: + env["NO_COLOR"] = "true" + + result = subprocess.run( + command, + env=env, + capture_output=True, # Capture both stdout and stderr + text=True, # Automatically decode to strings + check=True, # Automatically raise error on non-zero exit code + ) + + if format == json_format: + try: + remove_color_codes = re.compile(r"\x1b\[[0-9;]*m") + parsed_as_string = re.sub(remove_color_codes, "", result.stdout.strip()) + return json.loads(parsed_as_string) + except Exception: + logger.debug( + f"JSON parse exception in decoding output from pip {action}" + f" {package}. Falling back to accumulating the byte stream", + ) + accumulator = {} + parsed = BytesHeaderParser().parsebytes(result.stdout.encode()) + for key, value in parsed.items(): + accumulator[key] = value + return accumulator + + return result.stdout + + except subprocess.CalledProcessError as exc: + raise PipProcessError(action, package, exc.stderr, exc.stdout, exc.returncode) + except Exception as e: + raise PipProcessError(action, package, stderr=str(e), stdout="", returncode=1) + + def pip_process( action: str, package: str = "", diff --git a/guardrails/hub/validator_package_service.py b/guardrails/hub/validator_package_service.py index f54ccdbc9..c3478b3a3 100644 --- a/guardrails/hub/validator_package_service.py +++ b/guardrails/hub/validator_package_service.py @@ -12,7 +12,7 @@ from guardrails.logger import logger as guardrails_logger -from guardrails.cli.hub.utils import pip_process +from guardrails.cli.hub.utils import PipProcessError, pip_process_with_custom_exception from guardrails_hub_types import Manifest from guardrails.cli.server.hub_client import get_validator_manifest from guardrails.settings import settings @@ -251,7 +251,6 @@ def install_hub_module( validator_id ) validator_version = validator_version if validator_version else "" - full_package_name = f"{pep_503_package_name}{validator_version}" guardrails_token = settings.rc.token @@ -267,8 +266,41 @@ def install_hub_module( pip_flags.append("-q") # Install from guardrails hub pypi server with public pypi index as fallback - download_output = pip_process( - "install", full_package_name, pip_flags, quiet=quiet - ) - if not quiet: - logger.info(download_output) + + try: + full_package_name = f"{pep_503_package_name}[validators]{validator_version}" + download_output = pip_process_with_custom_exception( + "install", full_package_name, pip_flags, quiet=quiet + ) + if not quiet: + logger.info(download_output) + except PipProcessError: + try: + full_package_name = f"{pep_503_package_name}{validator_version}" + download_output = pip_process_with_custom_exception( + "install", full_package_name, pip_flags, quiet=quiet + ) + if not quiet: + logger.info(download_output) + except PipProcessError as e: + action = e.action + package = e.package + stderr = e.stderr + stdout = e.stdout + returncode = e.returncode + logger.error( + ( + f"Failed to {action} {package}\n" + f"Exit code: {returncode}\n" + f"stderr: {(stderr or '').strip()}\n" + f"stdout: {(stdout or '').strip()}" + ) + ) + raise + except Exception as e: + logger.error( + "An unexpected exception occurred while " + f"installing {validator_id}: ", + e, + ) + raise