From a1d25035a8fd85c46d8d5dd87b653340f80abc7e Mon Sep 17 00:00:00 2001
From: reidliu <reid201711@gmail.com>
Date: Wed, 19 Feb 2025 22:24:55 +0800
Subject: [PATCH] feat: add a graceful quit for build process

Signed-off-by: reidliu <reid201711@gmail.com>
---
 llama_stack/cli/stack/_build.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py
index 76f03aa5c..cefa5f96d 100644
--- a/llama_stack/cli/stack/_build.py
+++ b/llama_stack/cli/stack/_build.py
@@ -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,