[GTM-836]Rework Init workflow (#4377)

* Rework Init workflow

* minor format

* refactor

* add comments

* fix pyright alongside some improvements

* add demolink for blank template

* fix darglint

* Add more templates and keep template name in kebab case

* revert getting other templates since we'll use the submodules approach

* remove debug statement

* Improvements based on standup comments

* Add redirect logic

* changes based on new flow

---------

Co-authored-by: Masen Furer <[email protected]>

Author: ElijahAhianyo

reflex/constants/base.py

@@ -97,6 +97,18 @@ class Templates(SimpleNamespace):
     # The default template
     DEFAULT = "blank"
 
+    # The AI template
+    AI = "ai"
+
+    # The option for the user to choose a remote template.
+    CHOOSE_TEMPLATES = "choose-templates"
+
+    # The URL to find reflex templates.
+    REFLEX_TEMPLATES_URL = "https://reflex.dev/templates"
+
+    # Demo url for the default template.
+    DEFAULT_TEMPLATE_URL = "https://blank-template.reflex.run"
+
     # The reflex.build frontend host
     REFLEX_BUILD_FRONTEND = "https://flexgen.reflex.run"
 

reflex/reflex.py

@@ -17,7 +17,7 @@
 from reflex.config import environment, get_config
 from reflex.custom_components.custom_components import custom_components_cli
 from reflex.state import reset_disk_state_manager
-from reflex.utils import console, redir, telemetry
+from reflex.utils import console, telemetry
 
 # Disable typer+rich integration for help panels
 typer.core.rich = None  # type: ignore
@@ -89,38 +89,16 @@ def _init(
     # Set up the web project.
     prerequisites.initialize_frontend_dependencies()
 
-    # Integrate with reflex.build.
-    generation_hash = None
-    if ai:
-        if template is None:
-            # If AI is requested and no template specified, redirect the user to reflex.build.
-            generation_hash = redir.reflex_build_redirect()
-        elif prerequisites.is_generation_hash(template):
-            # Otherwise treat the template as a generation hash.
-            generation_hash = template
-        else:
-            console.error(
-                "Cannot use `--template` option with `--ai` option. Please remove `--template` option."
-            )
-            raise typer.Exit(2)
-        template = constants.Templates.DEFAULT
-
     # Initialize the app.
-    template = prerequisites.initialize_app(app_name, template)
-
-    # If a reflex.build generation hash is available, download the code and apply it to the main module.
-    if generation_hash:
-        prerequisites.initialize_main_module_index_from_generation(
-            app_name, generation_hash=generation_hash
-        )
+    template = prerequisites.initialize_app(app_name, template, ai)
 
     # Initialize the .gitignore.
     prerequisites.initialize_gitignore()
 
     # Initialize the requirements.txt.
     prerequisites.initialize_requirements_txt()
 
-    template_msg = "" if not template else f" using the {template} template"
+    template_msg = f" using the {template} template" if template else ""
     # Finish initializing the app.
     console.success(f"Initialized {app_name}{template_msg}")
 

reflex/utils/prerequisites.py

@@ -34,7 +34,7 @@
 from reflex import constants, model
 from reflex.compiler import templates
 from reflex.config import Config, environment, get_config
-from reflex.utils import console, net, path_ops, processes
+from reflex.utils import console, net, path_ops, processes, redir
 from reflex.utils.exceptions import (
     GeneratedCodeHasNoFunctionDefs,
     raise_system_package_missing_error,
@@ -1211,7 +1211,7 @@ def check_schema_up_to_date():
                 )
 
 
-def prompt_for_template(templates: list[Template]) -> str:
+def prompt_for_template_options(templates: list[Template]) -> str:
     """Prompt the user to specify a template.
 
     Args:
@@ -1223,9 +1223,14 @@ def prompt_for_template(templates: list[Template]) -> str:
     # Show the user the URLs of each template to preview.
     console.print("\nGet started with a template:")
 
+    def format_demo_url_str(url: str) -> str:
+        return f" ({url})" if url else ""
+
     # Prompt the user to select a template.
     id_to_name = {
-        str(idx): f"{template.name} ({template.demo_url}) - {template.description}"
+        str(
+            idx
+        ): f"{template.name.replace('_', ' ').replace('-', ' ')}{format_demo_url_str(template.demo_url)} - {template.description}"
         for idx, template in enumerate(templates)
     }
     for id in range(len(id_to_name)):
@@ -1380,15 +1385,119 @@ def create_config_init_app_from_remote_template(app_name: str, template_url: str
     shutil.rmtree(unzip_dir)
 
 
-def initialize_app(app_name: str, template: str | None = None) -> str | None:
-    """Initialize the app either from a remote template or a blank app. If the config file exists, it is considered as reinit.
+def initialize_default_app(app_name: str):
+    """Initialize the default app.
 
     Args:
         app_name: The name of the app.
-        template: The name of the template to use.
+    """
+    create_config(app_name)
+    initialize_app_directory(app_name)
+
+
+def validate_and_create_app_using_remote_template(app_name, template, templates):
+    """Validate and create an app using a remote template.
+
+    Args:
+        app_name: The name of the app.
+        template: The name of the template.
+        templates: The available templates.
+
+    Raises:
+        Exit: If the template is not found.
+    """
+    # If user selects a template, it needs to exist.
+    if template in templates:
+        template_url = templates[template].code_url
+    else:
+        # Check if the template is a github repo.
+        if template.startswith("https://github.com"):
+            template_url = f"{template.strip('/').replace('.git', '')}/archive/main.zip"
+        else:
+            console.error(f"Template `{template}` not found.")
+            raise typer.Exit(1)
+
+    if template_url is None:
+        return
+
+    create_config_init_app_from_remote_template(
+        app_name=app_name, template_url=template_url
+    )
+
+
+def generate_template_using_ai(template: str | None = None) -> str:
+    """Generate a template using AI(Flexgen).
+
+    Args:
+        template: The name of the template.
+
+    Returns:
+        The generation hash.
 
     Raises:
-        Exit: If template is directly provided in the command flag and is invalid.
+        Exit: If the template and ai flags are used.
+    """
+    if template is None:
+        # If AI is requested and no template specified, redirect the user to reflex.build.
+        return redir.reflex_build_redirect()
+    elif is_generation_hash(template):
+        # Otherwise treat the template as a generation hash.
+        return template
+    else:
+        console.error(
+            "Cannot use `--template` option with `--ai` option. Please remove `--template` option."
+        )
+        raise typer.Exit(2)
+
+
+def fetch_remote_templates(
+    template: Optional[str] = None,
+) -> tuple[str, dict[str, Template]]:
+    """Fetch the available remote templates.
+
+    Args:
+        template: The name of the template.
+
+    Returns:
+        The selected template and the available templates.
+
+    Raises:
+        Exit: If the template is not valid or if the template is not specified.
+    """
+    available_templates = {}
+
+    try:
+        # Get the available templates
+        available_templates = fetch_app_templates(constants.Reflex.VERSION)
+    except Exception as e:
+        console.warn("Failed to fetch templates. Falling back to default template.")
+        console.debug(f"Error while fetching templates: {e}")
+        template = constants.Templates.DEFAULT
+
+    if template == constants.Templates.DEFAULT:
+        return template, available_templates
+
+    if template in available_templates:
+        return template, available_templates
+
+    else:
+        if template is not None:
+            console.error(f"{template!r} is not a valid template name.")
+        console.print(
+            f"Go to the templates page ({constants.Templates.REFLEX_TEMPLATES_URL}) and copy the command to init with a template."
+        )
+        raise typer.Exit(0)
+
+
+def initialize_app(
+    app_name: str, template: str | None = None, ai: bool = False
+) -> str | None:
+    """Initialize the app either from a remote template or a blank app. If the config file exists, it is considered as reinit.
+
+    Args:
+        app_name: The name of the app.
+        template: The name of the template to use.
+        ai: Whether to use AI to generate the template.
 
     Returns:
         The name of the template.
@@ -1401,54 +1510,73 @@ def initialize_app(app_name: str, template: str | None = None) -> str | None:
         telemetry.send("reinit")
         return
 
+    generation_hash = None
+    if ai:
+        generation_hash = generate_template_using_ai(template)
+        template = constants.Templates.DEFAULT
+
     templates: dict[str, Template] = {}
 
     # Don't fetch app templates if the user directly asked for DEFAULT.
-    if template is None or (template != constants.Templates.DEFAULT):
-        try:
-            # Get the available templates
-            templates = fetch_app_templates(constants.Reflex.VERSION)
-            if template is None and len(templates) > 0:
-                template = prompt_for_template(list(templates.values()))
-        except Exception as e:
-            console.warn("Failed to fetch templates. Falling back to default template.")
-            console.debug(f"Error while fetching templates: {e}")
-        finally:
-            template = template or constants.Templates.DEFAULT
+    if template is not None and (template not in (constants.Templates.DEFAULT,)):
+        template, templates = fetch_remote_templates(template)
+
+    if template is None:
+        template = prompt_for_template_options(get_init_cli_prompt_options())
+        if template == constants.Templates.AI:
+            generation_hash = generate_template_using_ai()
+            # change to the default to allow creation of default app
+            template = constants.Templates.DEFAULT
+        elif template == constants.Templates.CHOOSE_TEMPLATES:
+            template, templates = fetch_remote_templates()
 
     # If the blank template is selected, create a blank app.
-    if template == constants.Templates.DEFAULT:
+    if template in (constants.Templates.DEFAULT,):
         # Default app creation behavior: a blank app.
-        create_config(app_name)
-        initialize_app_directory(app_name)
+        initialize_default_app(app_name)
     else:
-        # Fetch App templates from the backend server.
-        console.debug(f"Available templates: {templates}")
-
-        # If user selects a template, it needs to exist.
-        if template in templates:
-            template_url = templates[template].code_url
-        else:
-            # Check if the template is a github repo.
-            if template.startswith("https://github.com"):
-                template_url = (
-                    f"{template.strip('/').replace('.git', '')}/archive/main.zip"
-                )
-            else:
-                console.error(f"Template `{template}` not found.")
-                raise typer.Exit(1)
-
-        if template_url is None:
-            return
-
-        create_config_init_app_from_remote_template(
-            app_name=app_name, template_url=template_url
+        validate_and_create_app_using_remote_template(
+            app_name=app_name, template=template, templates=templates
         )
 
+    # If a reflex.build generation hash is available, download the code and apply it to the main module.
+    if generation_hash:
+        initialize_main_module_index_from_generation(
+            app_name, generation_hash=generation_hash
+        )
     telemetry.send("init", template=template)
+
     return template
 
 
+def get_init_cli_prompt_options() -> list[Template]:
+    """Get the CLI options for initializing a Reflex app.
+
+    Returns:
+        The CLI options.
+    """
+    return [
+        Template(
+            name=constants.Templates.DEFAULT,
+            description="A blank Reflex app.",
+            demo_url=constants.Templates.DEFAULT_TEMPLATE_URL,
+            code_url="",
+        ),
+        Template(
+            name=constants.Templates.AI,
+            description="Generate a template using AI [Experimental]",
+            demo_url="",
+            code_url="",
+        ),
+        Template(
+            name=constants.Templates.CHOOSE_TEMPLATES,
+            description="Choose an existing template.",
+            demo_url="",
+            code_url="",
+        ),
+    ]
+
+
 def initialize_main_module_index_from_generation(app_name: str, generation_hash: str):
     """Overwrite the `index` function in the main module with reflex.build generated code.
 

reflex/utils/redir.py

@@ -10,6 +10,18 @@
 from . import console
 
 
+def open_browser(target_url: str) -> None:
+    """Open a browser window to target_url.
+
+    Args:
+        target_url: The URL to open in the browser.
+    """
+    if not webbrowser.open(target_url):
+        console.warn(
+            f"Unable to automatically open the browser. Please navigate to {target_url} in your browser."
+        )
+
+
 def open_browser_and_wait(
     target_url: str, poll_url: str, interval: int = 2
 ) -> httpx.Response:
@@ -23,10 +35,7 @@ def open_browser_and_wait(
     Returns:
         The response from the poll_url.
     """
-    if not webbrowser.open(target_url):
-        console.warn(
-            f"Unable to automatically open the browser. Please navigate to {target_url} in your browser."
-        )
+    open_browser(target_url)
     console.info("[b]Complete the workflow in the browser to continue.[/b]")
     while True:
         try: