agent-starter-pack 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

src/cli/commands/setup_cicd.py

@@ -22,6 +22,7 @@ import tempfile
 import time
 from pathlib import Path
 
+import backoff
 import click
 from rich.console import Console
 
@@ -164,11 +165,12 @@ def update_build_triggers(tf_dir: Path) -> None:
 
 def prompt_for_repository_details(
     repository_name: str | None = None, repository_owner: str | None = None
-) -> tuple[str, str]:
+) -> tuple[str, str, bool]:
     """Interactive prompt for repository details with option to use existing repo."""
     # Get current GitHub username as default owner
     result = run_command(["gh", "api", "user", "--jq", ".login"], capture_output=True)
     default_owner = result.stdout.strip()
+    repository_exists = False
 
     if not (repository_name and repository_owner):
         console.print("\n📦 Repository Configuration", style="bold blue")
@@ -192,6 +194,7 @@ def prompt_for_repository_details(
                 )
         else:
             # Existing repository
+            repository_exists = True
             while True:
                 repo_url = click.prompt(
                     "Enter existing repository URL (e.g., https://github.com/owner/repo)"
@@ -233,10 +236,12 @@ def prompt_for_repository_details(
 
     if repository_name is None or repository_owner is None:
         raise ValueError("Repository name and owner must be provided")
-    return repository_name, repository_owner
+    return repository_name, repository_owner, repository_exists
 
 
-def setup_terraform_backend(tf_dir: Path, project_id: str, region: str) -> None:
+def setup_terraform_backend(
+    tf_dir: Path, project_id: str, region: str, repository_name: str
+) -> None:
     """Setup terraform backend configuration with GCS bucket"""
     console.print("\n🔧 Setting up Terraform backend...")
 
@@ -273,7 +278,7 @@ def setup_terraform_backend(tf_dir: Path, project_id: str, region: str) -> None:
         if dir_path.exists():
             # Use different state prefixes for dev and prod
             is_dev_dir = str(dir_path).endswith("/dev")
-            state_prefix = "dev" if is_dev_dir else "prod"
+            state_prefix = f"{repository_name}/{(is_dev_dir and 'dev') or 'prod'}"
 
             backend_file = dir_path / "backend.tf"
             backend_content = f'''terraform {{
@@ -383,6 +388,18 @@ console = Console()
     is_flag=True,
     help="Skip confirmation prompts and proceed automatically",
 )
+@click.option(
+    "--repository-exists",
+    is_flag=True,
+    default=False,
+    help="Flag indicating if the repository already exists",
+)
+@backoff.on_exception(
+    backoff.expo,
+    (subprocess.CalledProcessError, click.ClickException),
+    max_tries=3,
+    jitter=backoff.full_jitter,
+)
 def setup_cicd(
     dev_project: str | None,
     staging_project: str,
@@ -398,6 +415,7 @@ def setup_cicd(
     local_state: bool,
     debug: bool,
     auto_approve: bool,
+    repository_exists: bool,
 ) -> None:
     """Set up CI/CD infrastructure using Terraform."""
 
@@ -482,12 +500,12 @@ def setup_cicd(
 
     # Only prompt for repository details if not provided via CLI
     if not (repository_name and repository_owner):
-        repository_name, repository_owner = prompt_for_repository_details(
-            repository_name, repository_owner
+        repository_name, repository_owner, repository_exists = (
+            prompt_for_repository_details(repository_name, repository_owner)
         )
     # Set default host connection name if not provided
     if not host_connection_name:
-        host_connection_name = "github-connection"
+        host_connection_name = f"git-{repository_name}"
     # Check and enable required APIs regardless of auth method
     required_apis = ["secretmanager.googleapis.com", "cloudbuild.googleapis.com"]
     ensure_apis_enabled(cicd_project, required_apis)
@@ -548,6 +566,7 @@ def setup_cicd(
         github_pat=github_pat,
         github_app_installation_id=github_app_installation_id,
         git_provider=git_provider,
+        repository_exists=repository_exists,
     )
 
     tf_dir = Path("deployment/terraform")
@@ -562,7 +581,12 @@ def setup_cicd(
     # Setup Terraform backend if not using local state
     if not local_state:
         console.print("\n🔧 Setting up remote Terraform backend...")
-        setup_terraform_backend(tf_dir, cicd_project, region)
+        setup_terraform_backend(
+            tf_dir=tf_dir,
+            project_id=cicd_project,
+            region=region,
+            repository_name=repository_name,
+        )
         console.print("✅ Remote Terraform backend configured")
     else:
         console.print("\n📝 Using local Terraform state (remote backend disabled)")
@@ -600,6 +624,7 @@ def setup_cicd(
     new_vars["connection_exists"] = (
         "true" if detected_mode == "interactive" else "false"
     )
+    new_vars["repository_exists"] = "true" if config.repository_exists else "false"
 
     # Update or append variables
     with open(env_vars_path, "w") as f:
@@ -649,16 +674,35 @@ def setup_cicd(
             else:
                 run_command(["terraform", "init"], cwd=dev_tf_dir)
 
-            run_command(
-                [
-                    "terraform",
-                    "apply",
-                    "-auto-approve",
-                    "--var-file",
-                    "vars/env.tfvars",
-                ],
-                cwd=dev_tf_dir,
-            )
+            try:
+                run_command(
+                    [
+                        "terraform",
+                        "apply",
+                        "-auto-approve",
+                        "--var-file",
+                        "vars/env.tfvars",
+                    ],
+                    cwd=dev_tf_dir,
+                )
+            except subprocess.CalledProcessError as e:
+                if "Error acquiring the state lock" in str(e):
+                    console.print(
+                        "[yellow]State lock error detected, retrying without lock...[/yellow]"
+                    )
+                    run_command(
+                        [
+                            "terraform",
+                            "apply",
+                            "-auto-approve",
+                            "--var-file",
+                            "vars/env.tfvars",
+                            "-lock=false",
+                        ],
+                        cwd=dev_tf_dir,
+                    )
+                else:
+                    raise
 
             console.print("✅ Dev environment Terraform configuration applied")
     elif dev_tf_dir.exists():
@@ -673,10 +717,35 @@ def setup_cicd(
         else:
             run_command(["terraform", "init"], cwd=tf_dir)
 
-        run_command(
-            ["terraform", "apply", "-auto-approve", "--var-file", "vars/env.tfvars"],
-            cwd=tf_dir,
-        )
+        try:
+            run_command(
+                [
+                    "terraform",
+                    "apply",
+                    "-auto-approve",
+                    "--var-file",
+                    "vars/env.tfvars",
+                ],
+                cwd=tf_dir,
+            )
+        except subprocess.CalledProcessError as e:
+            if "Error acquiring the state lock" in str(e):
+                console.print(
+                    "[yellow]State lock error detected, retrying without lock...[/yellow]"
+                )
+                run_command(
+                    [
+                        "terraform",
+                        "apply",
+                        "-auto-approve",
+                        "--var-file",
+                        "vars/env.tfvars",
+                        "-lock=false",
+                    ],
+                    cwd=tf_dir,
+                )
+            else:
+                raise
 
         console.print("✅ Prod/Staging Terraform configuration applied")
 

src/cli/main.py

@@ -19,6 +19,7 @@ from rich.console import Console
 
 from .commands.create import create
 from .commands.setup_cicd import setup_cicd
+from .utils import display_update_message
 
 console = Console()
 
@@ -45,7 +46,8 @@ def print_version(ctx: click.Context, param: click.Parameter, value: bool) -> No
     help="Show the version and exit.",
 )
 def cli() -> None:
-    pass
+    # Check for updates at startup
+    display_update_message()
 
 
 # Register commands

src/cli/utils/__init__.py

@@ -12,24 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .datastores import DATASTORE_TYPES, get_datastore_info
 from .gcp import verify_credentials
 from .logging import handle_cli_error
 from .template import (
     get_available_agents,
     get_deployment_targets,
     get_template_path,
+    load_template_config,
     process_template,
-    prompt_data_ingestion,
+    prompt_datastore_selection,
     prompt_deployment_target,
 )
+from .version import display_update_message
 
 __all__ = [
+    "DATASTORE_TYPES",
+    "display_update_message",
     "get_available_agents",
+    "get_datastore_info",
     "get_deployment_targets",
     "get_template_path",
     "handle_cli_error",
+    "load_template_config",
     "process_template",
-    "prompt_data_ingestion",
+    "prompt_datastore_selection",
     "prompt_deployment_target",
     "verify_credentials",
 ]

src/cli/utils/cicd.py

@@ -22,6 +22,7 @@ from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 
+import backoff
 import click
 from rich.console import Console
 from rich.prompt import IntPrompt, Prompt
@@ -309,6 +310,7 @@ class ProjectConfig:
     project_name: str | None = None
     repository_name: str | None = None
     repository_owner: str | None = None
+    repository_exists: bool | None = None
     host_connection_name: str | None = None
     github_pat: str | None = None
     github_app_installation_id: str | None = None
@@ -414,6 +416,7 @@ def ensure_apis_enabled(project_id: str, apis: list[str]) -> None:
                 project_id,
                 f"--member=serviceAccount:{cloudbuild_sa}",
                 "--role=roles/secretmanager.admin",
+                "--condition=None",
             ]
         )
         console.print("✅ Permissions granted to Cloud Build service account")
@@ -428,6 +431,15 @@ def ensure_apis_enabled(project_id: str, apis: list[str]) -> None:
     time.sleep(10)
 
 
+@backoff.on_exception(
+    backoff.expo,
+    subprocess.CalledProcessError,
+    max_tries=3,
+    max_time=60,
+    on_backoff=lambda details: console.print(
+        f"⚠️ Command failed, retrying in {details['wait']:.1f}s (attempt {details['tries']})"
+    ),
+)
 def run_command(
     cmd: list[str] | str,
     check: bool = True,

src/cli/utils/datastores.py

@@ -0,0 +1,32 @@
+"""Datastore types and descriptions for data ingestion."""
+
+# Dictionary mapping datastore types to their descriptions
+DATASTORES = {
+    "vertex_ai_search": {
+        "name": "Vertex AI Search",
+        "description": "Managed, serverless document store that enables Google-quality search and RAG for generative AI.",
+    },
+    "vertex_ai_vector_search": {
+        "name": "Vertex AI Vector Search",
+        "description": "Scalable vector search engine for building search, recommendation systems, and generative AI applications. Based on ScaNN algorithm.",
+    },
+}
+
+DATASTORE_TYPES = list(DATASTORES.keys())
+
+
+def get_datastore_info(datastore_type: str) -> dict:
+    """Get information about a datastore type.
+
+    Args:
+        datastore_type: The datastore type key
+
+    Returns:
+        Dictionary with datastore information
+
+    Raises:
+        ValueError: If the datastore type is not valid
+    """
+    if datastore_type not in DATASTORES:
+        raise ValueError(f"Invalid datastore type: {datastore_type}")
+    return DATASTORES[datastore_type]

src/cli/utils/gcp.py

@@ -14,7 +14,6 @@
 
 # ruff: noqa: E722
 import subprocess
-from importlib import metadata
 
 import google.auth
 from google.api_core.client_options import ClientOptions
@@ -27,14 +26,13 @@ from google.cloud.aiplatform_v1beta1.types.prediction_service import (
     CountTokensRequest,
 )
 
+from src.cli.utils.version import PACKAGE_NAME, get_current_version
+
 
 def get_user_agent() -> tuple[str, str]:
     """Returns custom user agent header tuple (version, agent string)."""
-    try:
-        version = metadata.version(distribution_name="ags")
-    except metadata.PackageNotFoundError:
-        version = "0.0.0"
-    return version, f"ags/{version}"
+    version = get_current_version()
+    return version, f"{PACKAGE_NAME}/{version}"
 
 
 def get_client_info() -> ClientInfo:

src/cli/utils/template.py

@@ -22,6 +22,12 @@ from typing import Any
 
 import yaml
 from cookiecutter.main import cookiecutter
+from rich.console import Console
+from rich.prompt import Prompt
+
+from src.cli.utils.version import get_current_version
+
+from .datastores import DATASTORES
 
 
 @dataclass
@@ -172,15 +178,13 @@ def prompt_deployment_target(agent_name: str) -> str:
     if not targets:
         return ""
 
-    from rich.console import Console
-
     console = Console()
     console.print("\n> Please select a deployment target:")
     for idx, target in enumerate(targets, 1):
         info = TARGET_INFO.get(target, {})
         display_name = info.get("display_name", target)
         description = info.get("description", "")
-        console.print(f"{idx}. {display_name} - {description}")
+        console.print(f"{idx}. [bold]{display_name}[/] - [dim]{description}[/]")
 
     from rich.prompt import IntPrompt
 
@@ -192,8 +196,38 @@ def prompt_deployment_target(agent_name: str) -> str:
     return targets[choice - 1]
 
 
-def prompt_data_ingestion(agent_name: str) -> bool:
-    """Ask user if they want to include data pipeline if the agent supports it."""
+def prompt_datastore_selection(
+    agent_name: str, from_cli_flag: bool = False
+) -> str | None:
+    """Ask user to select a datastore type if the agent supports data ingestion.
+
+    Args:
+        agent_name: Name of the agent
+        from_cli_flag: Whether this is being called due to explicit --include-data-ingestion flag
+    """
+    console = Console()
+
+    # If this is from CLI flag, skip the "would you like to include" prompt
+    if from_cli_flag:
+        console.print("\n> Please select a datastore type for your data:")
+
+        # Display options with descriptions
+        for i, (_key, info) in enumerate(DATASTORES.items(), 1):
+            console.print(
+                f"{i}. [bold]{info['name']}[/] - [dim]{info['description']}[/]"
+            )
+
+        choice = Prompt.ask(
+            "\nEnter the number of your choice",
+            choices=[str(i) for i in range(1, len(DATASTORES) + 1)],
+            default="1",
+        )
+
+        # Convert choice number to datastore type
+        datastore_type = list(DATASTORES.keys())[int(choice) - 1]
+        return datastore_type
+
+    # Otherwise, proceed with normal flow
     template_path = (
         pathlib.Path(__file__).parent.parent.parent.parent
         / "agents"
@@ -203,26 +237,76 @@ def prompt_data_ingestion(agent_name: str) -> bool:
     config = load_template_config(template_path)
 
     if config:
-        # If requires_data_ingestion is true, return True without prompting
+        # If requires_data_ingestion is true, prompt for datastore type without asking if they want it
         if config.get("settings", {}).get("requires_data_ingestion"):
-            return True
+            console.print("\n> This agent includes a data ingestion pipeline.")
+            console.print("> Please select a datastore type for your data:")
 
-        # Only prompt if the agent has optional data processing support
-        if "data_ingestion" in config.get("settings", {}):
-            from rich.prompt import Prompt
+            # Display options with descriptions
+            for i, (_key, info) in enumerate(DATASTORES.items(), 1):
+                console.print(
+                    f"{i}. [bold]{info['name']}[/] - [dim]{info['description']}[/]"
+                )
+            choice = Prompt.ask(
+                "\nEnter the number of your choice",
+                choices=[str(i) for i in range(1, len(DATASTORES) + 1)],
+                default="1",
+            )
 
-            return (
+            # Convert choice number to datastore type
+            datastore_type = list(DATASTORES.keys())[int(choice) - 1]
+            return datastore_type
+
+        # Only prompt if the agent has optional data ingestion support
+        if "requires_data_ingestion" in config.get("settings", {}):
+            include = (
                 Prompt.ask(
-                    "\n> This agent supports a data pipeline. Would you like to include it?",
+                    "\n> This agent supports data ingestion. Would you like to include a data pipeline?",
                     choices=["y", "n"],
                     default="n",
                 ).lower()
                 == "y"
             )
-    return False
 
+            if include:
+                console.print("\n> Please select a datastore type for your data:")
 
-def get_template_path(agent_name: str, debug: bool = False) -> str:
+                # Display options with descriptions
+                for i, (_key, info) in enumerate(DATASTORES.items(), 1):
+                    console.print(
+                        f"{i}. [bold]{info['name']}[/] - [dim]{info['description']}[/]"
+                    )
+
+                choice = Prompt.ask(
+                    "\nEnter the number of your choice",
+                    choices=[str(i) for i in range(1, len(DATASTORES) + 1)],
+                    default="1",
+                )
+
+                # Convert choice number to datastore type
+                datastore_type = list(DATASTORES.keys())[int(choice) - 1]
+                return datastore_type
+
+    # If we get here, we need to prompt for datastore selection for explicit --include-data-ingestion flag
+    console.print(
+        "\n> Please select a datastore type for your data ingestion pipeline:"
+    )
+    # Display options with descriptions
+    for i, (_key, info) in enumerate(DATASTORES.items(), 1):
+        console.print(f"{i}. [bold]{info['name']}[/] - [dim]{info['description']}[/]")
+
+    choice = Prompt.ask(
+        "\nEnter the number of your choice",
+        choices=[str(i) for i in range(1, len(DATASTORES) + 1)],
+        default="1",
+    )
+
+    # Convert choice number to datastore type
+    datastore_type = list(DATASTORES.keys())[int(choice) - 1]
+    return datastore_type
+
+
+def get_template_path(agent_name: str, debug: bool = False) -> pathlib.Path:
     """Get the absolute path to the agent template directory."""
     current_dir = pathlib.Path(__file__).parent.parent.parent.parent
     template_path = current_dir / "agents" / agent_name / "template"
@@ -235,14 +319,17 @@ def get_template_path(agent_name: str, debug: bool = False) -> str:
     if not template_path.exists():
         raise ValueError(f"Template directory not found at {template_path}")
 
-    return str(template_path)
+    return template_path
 
 
-def copy_data_ingestion_files(project_template: pathlib.Path) -> None:
-    """Copy data processing files to the project template.
+def copy_data_ingestion_files(
+    project_template: pathlib.Path, datastore_type: str
+) -> None:
+    """Copy data processing files to the project template for cookiecutter templating.
 
     Args:
         project_template: Path to the project template directory
+        datastore_type: Type of datastore to use for data ingestion
     """
     data_ingestion_src = pathlib.Path(__file__).parent.parent.parent / "data_ingestion"
     data_ingestion_dst = project_template / "data_ingestion"
@@ -251,7 +338,10 @@ def copy_data_ingestion_files(project_template: pathlib.Path) -> None:
         logging.debug(
             f"Copying data processing files from {data_ingestion_src} to {data_ingestion_dst}"
         )
+
         copy_files(data_ingestion_src, data_ingestion_dst, overwrite=True)
+
+        logging.debug(f"Data ingestion files prepared for datastore: {datastore_type}")
     else:
         logging.warning(
             f"Data processing source directory not found at {data_ingestion_src}"
@@ -260,10 +350,11 @@ def copy_data_ingestion_files(project_template: pathlib.Path) -> None:
 
 def process_template(
     agent_name: str,
-    template_dir: str,
+    template_dir: pathlib.Path,
     project_name: str,
     deployment_target: str | None = None,
     include_data_ingestion: bool = False,
+    datastore: str | None = None,
     output_dir: pathlib.Path | None = None,
 ) -> None:
     """Process the template directory and create a new project.
@@ -278,11 +369,11 @@ def process_template(
     """
     logging.debug(f"Processing template from {template_dir}")
     logging.debug(f"Project name: {project_name}")
-    logging.debug(f"Include pipeline: {include_data_ingestion}")
+    logging.debug(f"Include pipeline: {datastore}")
     logging.debug(f"Output directory: {output_dir}")
 
     # Get paths
-    agent_path = pathlib.Path(template_dir).parent  # Get parent of template dir
+    agent_path = template_dir.parent  # Get parent of template dir
     logging.debug(f"agent path: {agent_path}")
     logging.debug(f"agent path exists: {agent_path.exists()}")
     logging.debug(
@@ -340,19 +431,15 @@ def process_template(
                     )
 
             # 3. Copy data ingestion files if needed
-            template_config = load_template_config(pathlib.Path(template_dir))
-            requires_data_ingestion = template_config.get("settings", {}).get(
-                "requires_data_ingestion", False
-            )
-            should_include_data_ingestion = (
-                include_data_ingestion or requires_data_ingestion
-            )
-
-            if should_include_data_ingestion:
-                logging.debug("3. Including data processing files")
-                copy_data_ingestion_files(project_template)
+            if include_data_ingestion and datastore:
+                logging.debug(
+                    f"3. Including data processing files with datastore: {datastore}"
+                )
+                copy_data_ingestion_files(project_template, datastore)
 
             # 4. Process frontend files
+            # Load template config
+            template_config = load_template_config(pathlib.Path(template_dir))
             frontend_type = template_config.get("settings", {}).get(
                 "frontend_type", DEFAULT_FRONTEND
             )
@@ -399,18 +486,11 @@ def process_template(
             template_config = load_template_config(pathlib.Path(template_dir))
 
             # Check if data processing should be included
-            requires_data_ingestion = template_config.get("settings", {}).get(
-                "requires_data_ingestion", False
-            )
-            should_include_data_ingestion = (
-                include_data_ingestion or requires_data_ingestion
-            )
-
-            if should_include_data_ingestion:
+            if include_data_ingestion and datastore:
                 logging.debug(
-                    "Including data processing files based on template config or user request"
+                    f"Including data processing files with datastore: {datastore}"
                 )
-                copy_data_ingestion_files(project_template)
+                copy_data_ingestion_files(project_template, datastore)
 
             # Create cookiecutter.json in the template root
             # Process extra dependencies
@@ -427,12 +507,14 @@ def process_template(
             cookiecutter_config = {
                 "project_name": "my-project",
                 "agent_name": agent_name,
+                "package_version": get_current_version(),
                 "agent_description": template_config.get("description", ""),
                 "deployment_target": deployment_target or "",
                 "frontend_type": frontend_type,
                 "extra_dependencies": [extra_deps],
                 "otel_instrumentations": otel_instrumentations,
-                "data_ingestion": should_include_data_ingestion,
+                "data_ingestion": include_data_ingestion,  # Use explicit flag for cookiecutter
+                "datastore_type": datastore if datastore else "",
                 "_copy_without_render": [
                     "*.ipynb",  # Don't render notebooks
                     "*.json",  # Don't render JSON files
@@ -535,10 +617,10 @@ def process_template(
 
 def should_exclude_path(path: pathlib.Path, agent_name: str) -> bool:
     """Determine if a path should be excluded based on the agent type."""
-    if agent_name == "multimodal_live_api":
-        # Exclude the unit test utils folder and app/utils folder for multimodal_live_api
+    if agent_name == "live_api":
+        # Exclude the unit test utils folder and app/utils folder for live_api
         if "tests/unit/test_utils" in str(path) or "app/utils" in str(path):
-            logging.debug(f"Excluding path for multimodal_live_api: {path}")
+            logging.debug(f"Excluding path for live_api: {path}")
             return True
     return False