feat: add a graceful quit for build process

Signed-off-by: reidliu <reid201711@gmail.com>
This commit is contained in:
reidliu 2025-02-19 22:24:55 +08:00
parent a66b4c4c81
commit a1d25035a8

View file

@ -55,6 +55,13 @@ def available_templates_specs() -> Dict[str, BuildConfig]:
return template_specs
def _check_quit_input(user_input: str) -> bool:
if user_input.lower() == "quit" or user_input.lower() == "q":
cprint("Quit the build process.", color="yellow")
return True
return False
def run_stack_build_command(args: argparse.Namespace) -> None:
if args.list_templates:
return _run_template_list_cmd()
@ -80,6 +87,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
)
return
elif not args.config and not args.template:
cprint("=" * 34, color="cyan")
cprint("Welcome to Llama Stack Build Setup", color="cyan", attrs=["bold"])
cprint("=" * 34, color="cyan")
cprint("Type 'quit' or 'q' to cancel at any time.", color="yellow")
name = prompt(
"> Enter a name for your Llama Stack (e.g. my-local-stack): ",
validator=Validator.from_callable(
@ -88,15 +99,21 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
),
)
if _check_quit_input(name):
return
image_type = prompt(
"> Enter the image type you want your Llama Stack to be built as (container or conda or venv): ",
validator=Validator.from_callable(
lambda x: x in ["container", "conda", "venv"],
lambda x: x in ["container", "conda", "venv", "q", "quit"],
error_message="Invalid image type, please enter conda or container or venv",
),
default="conda",
)
if _check_quit_input(image_type):
return
if image_type == "conda":
if not image_name:
cprint(
@ -127,6 +144,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
providers = dict()
for api, providers_for_api in get_provider_registry().items():
available_providers = [x for x in providers_for_api.keys() if x not in ("remote", "remote::sample")]
available_providers += ["q", "quit"]
api_provider = prompt(
"> Enter provider for API {}: ".format(api.value),
completer=WordCompleter(available_providers),
@ -137,6 +155,9 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
),
)
if _check_quit_input(api_provider):
return
providers[api.value] = api_provider
description = prompt(
@ -144,6 +165,9 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
default="",
)
if _check_quit_input(description):
return
distribution_spec = DistributionSpec(
providers=providers,
description=description,