refactor: Use ImageType enum values everywhere

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>
This commit is contained in:
Ihar Hrachyshka 2025-03-19 13:11:57 -04:00
parent 40bd49c6a4
commit b0d96f172e
3 changed files with 15 additions and 13 deletions

View file

@ -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: 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) conda)
--image-name IMAGE_NAME --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) found. (default: None)
--print-deps-only Print the dependencies for the stack only, without building the stack (default: False) --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) --run Run the stack after building using the same image type, name, and other applicable arguments (default: False)

View file

@ -21,6 +21,7 @@ from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.validation import Validator from prompt_toolkit.validation import Validator
from termcolor import cprint from termcolor import cprint
from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.table import print_table from llama_stack.cli.table import print_table
from llama_stack.distribution.build import ( from llama_stack.distribution.build import (
SERVER_DEPENDENCIES, SERVER_DEPENDENCIES,
@ -62,10 +63,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
if args.list_templates: if args.list_templates:
return _run_template_list_cmd() return _run_template_list_cmd()
if args.image_type == "venv": if args.image_type == ImageType.VENV.value:
current_venv = os.environ.get("VIRTUAL_ENV") current_venv = os.environ.get("VIRTUAL_ENV")
image_name = args.image_name or current_venv 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") current_conda_env = os.environ.get("CONDA_DEFAULT_ENV")
image_name = args.image_name or current_conda_env image_name = args.image_name or current_conda_env
else: else:
@ -84,7 +85,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
build_config.image_type = args.image_type build_config.image_type = args.image_type
else: else:
cprint( 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", color="red",
) )
sys.exit(1) sys.exit(1)
@ -98,15 +99,15 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
) )
image_type = prompt( 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( validator=Validator.from_callable(
lambda x: x in ["container", "conda", "venv"], lambda x: x in [e.value for e in ImageType],
error_message="Invalid image type, please enter conda or container or venv", 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: if not image_name:
cprint( cprint(
f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`", f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`",

View file

@ -6,6 +6,7 @@
import argparse import argparse
import textwrap import textwrap
from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.subcommand import Subcommand from llama_stack.cli.subcommand import Subcommand
@ -46,16 +47,16 @@ class StackBuild(Subcommand):
self.parser.add_argument( self.parser.add_argument(
"--image-type", "--image-type",
type=str, 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.", help="Image Type to use for the build. If not specified, will use the image type from the template config.",
choices=["conda", "container", "venv"], choices=[e.value for e in ImageType],
default="conda", default=ImageType.CONDA.value,
) )
self.parser.add_argument( self.parser.add_argument(
"--image-name", "--image-name",
type=str, type=str,
help=textwrap.dedent( 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. the build. If not specified, currently active Conda environment will be used if found.
""" """
), ),