From b0d96f172e583b156f621ace0ca3f29e51e8cb9c Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Wed, 19 Mar 2025 13:11:57 -0400 Subject: [PATCH] refactor: Use ImageType enum values everywhere Signed-off-by: Ihar Hrachyshka --- docs/source/distributions/building_distro.md | 2 +- llama_stack/cli/stack/_build.py | 17 +++++++++-------- llama_stack/cli/stack/build.py | 9 +++++---- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/docs/source/distributions/building_distro.md b/docs/source/distributions/building_distro.md index 9b8c1b9ad..e1e38d7ce 100644 --- a/docs/source/distributions/building_distro.md +++ b/docs/source/distributions/building_distro.md @@ -67,7 +67,7 @@ options: Image Type to use for the build. This can be either conda or container or venv. If not specified, will use the image type from the template config. (default: conda) --image-name IMAGE_NAME - [for image-type=conda|venv] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if + [for image-type=conda|container|venv] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if found. (default: None) --print-deps-only Print the dependencies for the stack only, without building the stack (default: False) --run Run the stack after building using the same image type, name, and other applicable arguments (default: False) diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index d87e3bd0b..e440799d1 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -21,6 +21,7 @@ from prompt_toolkit.completion import WordCompleter from prompt_toolkit.validation import Validator from termcolor import cprint +from llama_stack.cli.stack.utils import ImageType from llama_stack.cli.table import print_table from llama_stack.distribution.build import ( SERVER_DEPENDENCIES, @@ -62,10 +63,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None: if args.list_templates: return _run_template_list_cmd() - if args.image_type == "venv": + if args.image_type == ImageType.VENV.value: current_venv = os.environ.get("VIRTUAL_ENV") image_name = args.image_name or current_venv - elif args.image_type == "conda": + elif args.image_type == ImageType.CONDA.value: current_conda_env = os.environ.get("CONDA_DEFAULT_ENV") image_name = args.image_name or current_conda_env else: @@ -84,7 +85,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None: build_config.image_type = args.image_type else: cprint( - f"Please specify a image-type (container | conda | venv) for {args.template}", + f"Please specify a image-type ({' | '.join(e.value for e in ImageType)}) for {args.template}", color="red", ) sys.exit(1) @@ -98,15 +99,15 @@ def run_stack_build_command(args: argparse.Namespace) -> None: ) image_type = prompt( - "> Enter the image type you want your Llama Stack to be built as (container or conda or venv): ", + f"> Enter the image type you want your Llama Stack to be built as ({' or '.join(e.value for e in ImageType)}): ", validator=Validator.from_callable( - lambda x: x in ["container", "conda", "venv"], - error_message="Invalid image type, please enter conda or container or venv", + lambda x: x in [e.value for e in ImageType], + error_message=f"Invalid image type, please enter {' or '.join(e.value for e in ImageType)}", ), - default="conda", + default=ImageType.CONDA.value, ) - if image_type == "conda": + if image_type == ImageType.CONDA.value: if not image_name: cprint( f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`", diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py index 70d74c620..0ada7c615 100644 --- a/llama_stack/cli/stack/build.py +++ b/llama_stack/cli/stack/build.py @@ -6,6 +6,7 @@ import argparse import textwrap +from llama_stack.cli.stack.utils import ImageType from llama_stack.cli.subcommand import Subcommand @@ -46,16 +47,16 @@ class StackBuild(Subcommand): self.parser.add_argument( "--image-type", type=str, - help="Image Type to use for the build. This can be either conda or container or venv. If not specified, will use the image type from the template config.", - choices=["conda", "container", "venv"], - default="conda", + help="Image Type to use for the build. If not specified, will use the image type from the template config.", + choices=[e.value for e in ImageType], + default=ImageType.CONDA.value, ) self.parser.add_argument( "--image-name", type=str, help=textwrap.dedent( - """[for image-type=conda|venv] Name of the conda or virtual environment to use for + f"""[for image-type={"|".join(e.value for e in ImageType)}] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if found. """ ),