agent-starter-pack 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

src/cli/utils/template.py

@@ -29,6 +29,9 @@ from rich.prompt import IntPrompt, Prompt
 from src.cli.utils.version import get_current_version
 
 from .datastores import DATASTORES
+from .remote_template import (
+    get_base_template_name,
+)
 
 ADK_FILES = ["app/__init__.py"]
 NON_ADK_FILES: list[str] = []
@@ -69,7 +72,7 @@ class TemplateConfig:
 
 
 OVERWRITE_FOLDERS = ["app", "frontend", "tests", "notebooks"]
-TEMPLATE_CONFIG_FILE = ".templateconfig.yaml"
+TEMPLATE_CONFIG_FILE = "templateconfig.yaml"
 DEPLOYMENT_FOLDERS = ["cloud_run", "agent_engine"]
 DEFAULT_FRONTEND = "streamlit"
 
@@ -94,7 +97,7 @@ def get_available_agents(deployment_target: str | None = None) -> dict:
 
     for agent_dir in agents_dir.iterdir():
         if agent_dir.is_dir() and not agent_dir.name.startswith("__"):
-            template_config_path = agent_dir / "template" / ".templateconfig.yaml"
+            template_config_path = agent_dir / ".template" / "templateconfig.yaml"
             if template_config_path.exists():
                 try:
                     with open(template_config_path) as f:
@@ -154,15 +157,20 @@ def load_template_config(template_dir: pathlib.Path) -> dict[str, Any]:
         return {}
 
 
-def get_deployment_targets(agent_name: str) -> list:
+def get_deployment_targets(
+    agent_name: str, remote_config: dict[str, Any] | None = None
+) -> list:
     """Get available deployment targets for the selected agent."""
-    template_path = (
-        pathlib.Path(__file__).parent.parent.parent.parent
-        / "agents"
-        / agent_name
-        / "template"
-    )
-    config = load_template_config(template_path)
+    if remote_config:
+        config = remote_config
+    else:
+        template_path = (
+            pathlib.Path(__file__).parent.parent.parent.parent
+            / "agents"
+            / agent_name
+            / ".template"
+        )
+        config = load_template_config(template_path)
 
     if not config:
         return []
@@ -171,9 +179,11 @@ def get_deployment_targets(agent_name: str) -> list:
     return targets if isinstance(targets, list) else [targets]
 
 
-def prompt_deployment_target(agent_name: str) -> str:
+def prompt_deployment_target(
+    agent_name: str, remote_config: dict[str, Any] | None = None
+) -> str:
     """Ask user to select a deployment target for the agent."""
-    targets = get_deployment_targets(agent_name)
+    targets = get_deployment_targets(agent_name, remote_config=remote_config)
 
     # Define deployment target friendly names and descriptions
     TARGET_INFO = {
@@ -206,6 +216,36 @@ def prompt_deployment_target(agent_name: str) -> str:
     return targets[choice - 1]
 
 
+def prompt_session_type_selection() -> str:
+    """Ask user to select a session type for Cloud Run deployment."""
+    console = Console()
+
+    session_types = {
+        "in_memory": {
+            "display_name": "In-memory session",
+            "description": "Session data stored in memory - ideal for stateless applications",
+        },
+        "alloydb": {
+            "display_name": "AlloyDB",
+            "description": "Use AlloyDB for session management. Comes with terraform resources for deployment.",
+        },
+    }
+
+    console.print("\n> Please select a session type:")
+    for idx, (_key, info) in enumerate(session_types.items(), 1):
+        console.print(
+            f"{idx}. [bold]{info['display_name']}[/] - [dim]{info['description']}[/]"
+        )
+
+    choice = IntPrompt.ask(
+        "\nEnter the number of your session type choice",
+        default=1,
+        show_default=True,
+    )
+
+    return list(session_types.keys())[choice - 1]
+
+
 def prompt_datastore_selection(
     agent_name: str, from_cli_flag: bool = False
 ) -> str | None:
@@ -242,7 +282,7 @@ def prompt_datastore_selection(
         pathlib.Path(__file__).parent.parent.parent.parent
         / "agents"
         / agent_name
-        / "template"
+        / ".template"
     )
     config = load_template_config(template_path)
 
@@ -319,7 +359,7 @@ def prompt_datastore_selection(
 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"
+    template_path = current_dir / "agents" / agent_name / ".template"
     if debug:
         logging.debug(f"Looking for template in: {template_path}")
         logging.debug(f"Template exists: {template_path.exists()}")
@@ -365,7 +405,10 @@ def process_template(
     deployment_target: str | None = None,
     include_data_ingestion: bool = False,
     datastore: str | None = None,
+    session_type: str | None = None,
     output_dir: pathlib.Path | None = None,
+    remote_template_path: pathlib.Path | None = None,
+    remote_config: dict[str, Any] | None = None,
 ) -> None:
     """Process the template directory and create a new project.
 
@@ -375,15 +418,33 @@ def process_template(
         project_name: Name of the project to create
         deployment_target: Optional deployment target (agent_engine or cloud_run)
         include_data_ingestion: Whether to include data pipeline components
+        datastore: Optional datastore type for data ingestion
+        session_type: Optional session type for cloud_run deployment
         output_dir: Optional output directory path, defaults to current directory
+        remote_template_path: Optional path to remote template for overlay
+        remote_config: Optional remote template configuration
     """
     logging.debug(f"Processing template from {template_dir}")
     logging.debug(f"Project name: {project_name}")
     logging.debug(f"Include pipeline: {datastore}")
     logging.debug(f"Output directory: {output_dir}")
 
-    # Get paths
-    agent_path = template_dir.parent  # Get parent of template dir
+    # Handle remote vs local templates
+    is_remote = remote_template_path is not None
+
+    if is_remote:
+        # For remote templates, determine the base template
+        base_template_name = get_base_template_name(remote_config or {})
+        agent_path = (
+            pathlib.Path(__file__).parent.parent.parent.parent
+            / "agents"
+            / base_template_name
+        )
+        logging.debug(f"Remote template using base: {base_template_name}")
+    else:
+        # For local templates, use the existing logic
+        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(
@@ -456,7 +517,7 @@ def process_template(
             copy_frontend_files(frontend_type, project_template)
             logging.debug(f"4. Processed frontend files for type: {frontend_type}")
 
-            # 5. Finally, copy agent-specific files to override everything else
+            # 5. Copy agent-specific files to override base template
             if agent_path.exists():
                 for folder in OVERWRITE_FOLDERS:
                     agent_folder = agent_path / folder
@@ -467,11 +528,27 @@ def process_template(
                             agent_folder, project_folder, agent_name, overwrite=True
                         )
 
+            # 6. Finally, overlay remote template files if present
+            if is_remote and remote_template_path:
+                logging.debug(
+                    f"6. Overlaying remote template files from {remote_template_path}"
+                )
+                copy_files(
+                    remote_template_path,
+                    project_template,
+                    agent_name=agent_name,
+                    overwrite=True,
+                )
+
             # Load and validate template config first
-            template_path = pathlib.Path(template_dir)
-            config = load_template_config(template_path)
+            if is_remote:
+                config = remote_config or {}
+            else:
+                template_path = pathlib.Path(template_dir)
+                config = load_template_config(template_path)
+
             if not config:
-                raise ValueError(f"Could not load template config from {template_path}")
+                raise ValueError("Could not load template config")
 
             # Validate deployment target
             available_targets = config.get("settings", {}).get("deployment_targets", [])
@@ -483,8 +560,8 @@ def process_template(
                     f"Invalid deployment target '{deployment_target}'. Available targets: {available_targets}"
                 )
 
-            # Load template config
-            template_config = load_template_config(pathlib.Path(template_dir))
+            # Use the already loaded config
+            template_config = config
 
             # Check if data processing should be included
             if include_data_ingestion and datastore:
@@ -527,6 +604,7 @@ def process_template(
                 "settings": settings,
                 "tags": tags,
                 "deployment_target": deployment_target or "",
+                "session_type": session_type or "",
                 "frontend_type": frontend_type,
                 "extra_dependencies": [extra_deps],
                 "data_ingestion": include_data_ingestion,
@@ -591,9 +669,36 @@ def process_template(
                         file_path.unlink()
                         logging.debug(f"Deleted {file_path}")
 
-                # After copying template files, handle the lock file
-                if deployment_target:
-                    # Get the source lock file path
+                # Handle pyproject.toml and uv.lock files
+                if is_remote and remote_template_path:
+                    # For remote templates, use their pyproject.toml and uv.lock if they exist
+                    remote_pyproject = remote_template_path / "pyproject.toml"
+                    remote_uv_lock = remote_template_path / "uv.lock"
+
+                    if remote_pyproject.exists():
+                        shutil.copy2(
+                            remote_pyproject, final_destination / "pyproject.toml"
+                        )
+                        logging.debug("Used pyproject.toml from remote template")
+
+                    if remote_uv_lock.exists():
+                        shutil.copy2(remote_uv_lock, final_destination / "uv.lock")
+                        logging.debug("Used uv.lock from remote template")
+                    elif deployment_target:
+                        # Fallback to base template lock file
+                        base_template_name = get_base_template_name(remote_config or {})
+                        lock_path = (
+                            pathlib.Path(__file__).parent.parent.parent.parent
+                            / "src"
+                            / "resources"
+                            / "locks"
+                            / f"uv-{base_template_name}-{deployment_target}.lock"
+                        )
+                        if lock_path.exists():
+                            shutil.copy2(lock_path, final_destination / "uv.lock")
+                            logging.debug(f"Used fallback lock file from {lock_path}")
+                elif deployment_target:
+                    # For local templates, use the existing logic
                     lock_path = (
                         pathlib.Path(__file__).parent.parent.parent.parent
                         / "src"
@@ -672,8 +777,12 @@ def copy_files(
             return True
         if "__pycache__" in str(path) or path.name == "__pycache__":
             return True
+        if ".git" in path.parts:
+            return True
         if agent_name is not None and should_exclude_path(path, agent_name):
             return True
+        if path.is_dir() and path.name == ".template":
+            return True
         return False
 
     if src.is_dir():

src/deployment_targets/agent_engine/app/agent_engine_app.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# mypy: disable-error-code="attr-defined"
+# mypy: disable-error-code="attr-defined,arg-type"
 {%- if "adk" in cookiecutter.tags %}
 import copy
 import datetime
 import json
 import logging
 import os
-from collections.abc import Mapping, Sequence
 from typing import Any
 
 import google.auth
@@ -57,7 +56,7 @@ class AgentEngineApp(AdkApp):
         feedback_obj = Feedback.model_validate(feedback)
         self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
 
-    def register_operations(self) -> Mapping[str, Sequence]:
+    def register_operations(self) -> dict[str, list[str]]:
         """Registers the operations of the Agent.
 
         Extends the base operations to include feedback registration functionality.
@@ -69,9 +68,10 @@ class AgentEngineApp(AdkApp):
     def clone(self) -> "AgentEngineApp":
         """Returns a clone of the ADK application."""
         template_attributes = self._tmpl_attrs
+
         return self.__class__(
-            agent=copy.deepcopy(template_attributes.get("agent")),
-            enable_tracing=template_attributes.get("enable_tracing"),
+            agent=copy.deepcopy(template_attributes["agent"]),
+            enable_tracing=bool(template_attributes.get("enable_tracing", False)),
             session_service_builder=template_attributes.get("session_service_builder"),
             artifact_service_builder=template_attributes.get(
                 "artifact_service_builder"
@@ -83,7 +83,7 @@ import datetime
 import json
 import logging
 import os
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping
 from typing import (
     Any,
 )
@@ -182,7 +182,7 @@ class AgentEngineApp:
         feedback_obj = Feedback.model_validate(feedback)
         self.logger.log_struct(feedback_obj.model_dump(), severity="INFO")
 
-    def register_operations(self) -> Mapping[str, Sequence]:
+    def register_operations(self) -> dict[str, list[str]]:
         """Registers the operations of the Agent.
 
         This mapping defines how different operation modes (e.g., "", "stream")

src/deployment_targets/agent_engine/tests/load_test/README.md

@@ -18,12 +18,7 @@ Follow these steps to execute load tests:
    It's recommended to use a separate terminal tab and create a virtual environment for Locust to avoid conflicts with your application's Python environment.
 
    ```bash
-   # Create and activate virtual environment
-   python3 -m venv .locust_env
-   source .locust_env/bin/activate
-   
-   # Install required packages
-   pip install locust==2.31.1 "google-cloud-aiplatform[langchain,reasoningengine]>=1.77.0"
+   python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1
    ```
 
 **3. Execute the Load Test:**

src/deployment_targets/agent_engine/tests/load_test/load_test.py

@@ -91,15 +91,25 @@ class ChatStreamUser(HttpUser):
                 events = []
                 for line in response.iter_lines():
                     if line:
-                        event = json.loads(line)
-                        events.append(event)
+                        line_str = line.decode("utf-8")
+                        events.append(line_str)
+
+                        if "429 Too Many Requests" in line_str:
+                            self.environment.events.request.fire(
+                                request_type="POST",
+                                name=f"{url_path} rate_limited 429s",
+                                response_time=0,
+                                response_length=len(line),
+                                response=response,
+                                context={},
+                            )
                 end_time = time.time()
                 total_time = end_time - start_time
                 self.environment.events.request.fire(
                     request_type="POST",
                     name="/stream_messages end",
                     response_time=total_time * 1000,  # Convert to milliseconds
-                    response_length=len(json.dumps(events)),
+                    response_length=len(events),
                     response=response,
                     context={},
                 )

src/deployment_targets/cloud_run/Dockerfile

@@ -24,6 +24,9 @@ COPY ./app ./app
 
 RUN uv sync --frozen
 
+ARG COMMIT_SHA=""
+ENV COMMIT_SHA=${COMMIT_SHA}
+
 EXPOSE 8080
 
 CMD ["uv", "run", "uvicorn", "app.server:app", "--host", "0.0.0.0", "--port", "8080"]

src/deployment_targets/cloud_run/app/server.py

@@ -43,11 +43,29 @@ provider.add_span_processor(processor)
 trace.set_tracer_provider(provider)
 
 AGENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+{%- if cookiecutter.session_type == "alloydb" %}
+# AlloyDB session configuration
+db_user = os.environ.get("DB_USER", "postgres")
+db_name = os.environ.get("DB_NAME", "postgres")
+db_pass = os.environ.get("DB_PASS")
+db_host = os.environ.get("DB_HOST")
+
+# Set session_service_uri if database credentials are available
+session_service_uri = None
+if db_host and db_pass:
+    session_service_uri = f"postgresql://{db_user}:{db_pass}@{db_host}:5432/{db_name}"
+{%- else %}
+# In-memory session configuration - no persistent storage
+session_service_uri = None
+{%- endif %}
+
 app: FastAPI = get_fast_api_app(
     agents_dir=AGENT_DIR,
     web=True,
     artifact_service_uri=bucket_name,
     allow_origins=allow_origins,
+    session_service_uri=session_service_uri,
 )
 app.title = "{{cookiecutter.project_name}}"
 app.description = "API for interacting with the Agent {{cookiecutter.project_name}}"

src/deployment_targets/cloud_run/deployment/terraform/dev/service.tf

@@ -0,0 +1,231 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Get project information to access the project number
+data "google_project" "project" {
+  project_id = var.dev_project_id
+}
+
+
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+
+# VPC Network for AlloyDB
+resource "google_compute_network" "default" {
+  name                    = "${var.project_name}-alloydb-network"
+  project                 = var.dev_project_id
+  auto_create_subnetworks = false
+  depends_on = [resource.google_project_service.services]
+}
+
+# Subnet for AlloyDB
+resource "google_compute_subnetwork" "default" {
+  name          = "${var.project_name}-alloydb-network"
+  ip_cidr_range = "10.0.0.0/24"
+  region        = var.region
+  network       = google_compute_network.default.id
+  project       = var.dev_project_id
+
+  # This is required for Cloud Run VPC connectors
+  purpose       = "PRIVATE"
+
+  private_ip_google_access = true
+}
+
+# Private IP allocation for AlloyDB
+resource "google_compute_global_address" "private_ip_alloc" {
+  name          = "${var.project_name}-private-ip"
+  project       = var.dev_project_id
+  address_type  = "INTERNAL"
+  purpose       = "VPC_PEERING"
+  prefix_length = 16
+  network       = google_compute_network.default.id
+
+  depends_on = [resource.google_project_service.services]
+}
+
+# VPC connection for AlloyDB
+resource "google_service_networking_connection" "vpc_connection" {
+  network                 = google_compute_network.default.id
+  service                 = "servicenetworking.googleapis.com"
+  reserved_peering_ranges = [google_compute_global_address.private_ip_alloc.name]
+}
+
+# AlloyDB Cluster
+resource "google_alloydb_cluster" "session_db_cluster" {
+  project    = var.dev_project_id
+  cluster_id = "${var.project_name}-alloydb-cluster"
+  location   = var.region
+
+  network_config {
+    network = google_compute_network.default.id
+  }
+
+  depends_on = [
+    google_service_networking_connection.vpc_connection
+  ]
+}
+
+# AlloyDB Instance
+resource "google_alloydb_instance" "session_db_instance" {
+  cluster       = google_alloydb_cluster.session_db_cluster.name
+  instance_id   = "${var.project_name}-alloydb-instance"
+  instance_type = "PRIMARY"
+
+  availability_type = "REGIONAL" # Regional redundancy
+
+  machine_config {
+    cpu_count = 2
+  }
+}
+
+# Generate a random password for the database user
+resource "random_password" "db_password" {
+  length           = 16
+  special          = true
+  override_special = "!#$%&*()-_=+[]{}<>:?"
+}
+
+# Store the password in Secret Manager
+resource "google_secret_manager_secret" "db_password" {
+  project   = var.dev_project_id
+  secret_id = "${var.project_name}-db-password"
+
+  replication {
+    auto {}
+  }
+
+  depends_on = [resource.google_project_service.services]
+}
+
+resource "google_secret_manager_secret_version" "db_password" {
+  secret      = google_secret_manager_secret.db_password.id
+  secret_data = random_password.db_password.result
+}
+
+resource "google_alloydb_user" "db_user" {
+  cluster        = google_alloydb_cluster.session_db_cluster.name
+  user_id        = "postgres"
+  user_type      = "ALLOYDB_BUILT_IN"
+  password       = random_password.db_password.result
+  database_roles = ["alloydbsuperuser"]
+
+  depends_on = [google_alloydb_instance.session_db_instance]
+}
+
+{%- endif %}
+
+
+resource "google_cloud_run_v2_service" "app" {
+  name                = var.project_name
+  location            = var.region
+  project             = var.dev_project_id
+  deletion_protection = false
+  ingress             = "INGRESS_TRAFFIC_ALL"
+
+  template {
+    containers {
+      image = "us-docker.pkg.dev/cloudrun/container/hello"
+
+      resources {
+        limits = {
+          cpu    = "4"
+          memory = "8Gi"
+        }
+      }
+{%- if cookiecutter.data_ingestion %}
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+
+      env {
+        name  = "DATA_STORE_ID"
+        value = resource.google_discovery_engine_data_store.data_store_dev.data_store_id
+      }
+
+      env {
+        name  = "DATA_STORE_REGION"
+        value = var.data_store_region
+      }
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+      env {
+        name  = "VECTOR_SEARCH_INDEX"
+        value = resource.google_vertex_ai_index.vector_search_index.id
+      }
+
+      env {
+        name  = "VECTOR_SEARCH_INDEX_ENDPOINT"
+        value = resource.google_vertex_ai_index_endpoint.vector_search_index_endpoint.id
+      }
+
+      env {
+        name  = "VECTOR_SEARCH_BUCKET"
+        value = "gs://${resource.google_storage_bucket.data_ingestion_PIPELINE_GCS_ROOT.name}"
+      }
+{%- endif %}
+{%- endif %}
+
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+
+      env {
+        name  = "DB_HOST"
+        value = google_alloydb_instance.session_db_instance.ip_address
+      }
+
+      env {
+        name = "DB_PASS"
+        value_source {
+          secret_key_ref {
+            secret  = google_secret_manager_secret.db_password.secret_id
+            version = "latest"
+          }
+        }
+      }
+{%- endif %}
+    }
+
+    service_account = google_service_account.cloud_run_app_sa.email
+    max_instance_request_concurrency = 40
+
+    scaling {
+      min_instance_count = 1
+      max_instance_count = 10
+    }
+
+    session_affinity = true
+
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+    # VPC access for AlloyDB connectivity
+    vpc_access {
+      network_interfaces {
+        network    = google_compute_network.default.id
+        subnetwork = google_compute_subnetwork.default.id
+      }
+    }
+{%- endif %}
+  }
+
+  traffic {
+    type    = "TRAFFIC_TARGET_ALLOCATION_TYPE_LATEST"
+    percent = 100
+  }
+
+  # This lifecycle block prevents Terraform from overwriting the container image when it's
+  # updated by Cloud Run deployments outside of Terraform (e.g., via CI/CD pipelines)
+  lifecycle {
+    ignore_changes = [
+      template[0].containers[0].image,
+    ]
+  }
+
+  # Make dependencies conditional to avoid errors.
+  depends_on = [resource.google_project_service.services]
+}