agent-starter-pack 0.7.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

src/base_template/deployment/terraform/build_triggers.tf

@@ -39,7 +39,7 @@ resource "google_cloudbuild_trigger" "pr_checks" {
   {% endif %}
   ]
   include_build_logs = "INCLUDE_BUILD_LOGS_WITH_STATUS"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 # b. Create CD pipeline trigger
@@ -85,13 +85,13 @@ resource "google_cloudbuild_trigger" "cd_pipeline" {
 {% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
     _VECTOR_SEARCH_INDEX           = resource.google_vertex_ai_index.vector_search_index_staging.id
     _VECTOR_SEARCH_INDEX_ENDPOINT  = resource.google_vertex_ai_index_endpoint.vector_search_index_endpoint_staging.id
-    _VECTOR_SEARCH_BUCKET          = "gs://${resource.google_storage_bucket.vector_search_data_bucket["staging"].name}"
+    _VECTOR_SEARCH_BUCKET          = resource.google_storage_bucket.vector_search_data_bucket["staging"].url
 
 {% endif %}
 {% endif %}
     # Your other CD Pipeline substitutions
   }
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 
 }
 
@@ -128,11 +128,11 @@ resource "google_cloudbuild_trigger" "deploy_to_prod_pipeline" {
 {% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
     _VECTOR_SEARCH_INDEX           = resource.google_vertex_ai_index.vector_search_index_prod.id
     _VECTOR_SEARCH_INDEX_ENDPOINT  = resource.google_vertex_ai_index_endpoint.vector_search_index_endpoint_prod.id
-    _VECTOR_SEARCH_BUCKET          = "gs://${resource.google_storage_bucket.vector_search_data_bucket["prod"].name}"
+    _VECTOR_SEARCH_BUCKET          = resource.google_storage_bucket.vector_search_data_bucket["prod"].url
 {% endif %}
 {% endif %}
     # Your other Deploy to Prod Pipeline substitutions
   }
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 
 }

src/base_template/deployment/terraform/dev/apis.tf

@@ -24,7 +24,14 @@ locals {
     "bigquery.googleapis.com",
     "serviceusage.googleapis.com",
     "logging.googleapis.com",
-    "cloudtrace.googleapis.com"
+    "cloudtrace.googleapis.com",
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+    "compute.googleapis.com",
+    "servicenetworking.googleapis.com",
+    "alloydb.googleapis.com",
+    "secretmanager.googleapis.com",
+    "dns.googleapis.com"
+{%- endif %}
   ]
 }
 

src/base_template/deployment/terraform/dev/variables.tf

@@ -54,6 +54,9 @@ variable "agentengine_sa_roles" {
 {% endif %}
   type        = list(string)
   default = [
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+    "roles/secretmanager.secretAccessor",
+{%- endif %}
     "roles/aiplatform.user",
     "roles/discoveryengine.editor",
     "roles/logging.logWriter",
@@ -61,7 +64,6 @@ variable "agentengine_sa_roles" {
     "roles/storage.admin"
   ]
 }
-
 {% if cookiecutter.data_ingestion %}
 
 variable "pipelines_roles" {

src/base_template/deployment/terraform/iam.tf

@@ -25,7 +25,7 @@ resource "google_project_iam_member" "cicd_project_roles" {
   project    = var.cicd_runner_project_id
   role       = each.value
   member     = "serviceAccount:${resource.google_service_account.cicd_runner_sa.email}"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 
 }
 
@@ -42,7 +42,7 @@ resource "google_project_iam_member" "other_projects_roles" {
   project    = each.value.project_id
   role       = each.value.role
   member     = "serviceAccount:${resource.google_service_account.cicd_runner_sa.email}"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% if cookiecutter.deployment_target == 'cloud_run' %}
 # 3. Allow Cloud Run service SA to pull containers stored in the CICD project
@@ -52,7 +52,7 @@ resource "google_project_iam_member" "cicd_run_invoker_artifact_registry_reader"
 
   role       = "roles/artifactregistry.reader"
   member     = "serviceAccount:service-${data.google_project.projects[each.key].number}@serverless-robot-prod.iam.gserviceaccount.com"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 
 }
 
@@ -69,7 +69,7 @@ resource "google_project_iam_member" "cloud_run_app_sa_roles" {
   project    = each.value.project
   role       = each.value.role
   member     = "serviceAccount:${google_service_account.cloud_run_app_sa[split(",", each.key)[0]].email}"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% elif cookiecutter.deployment_target == 'agent_engine' %}
 resource "google_project_service_identity" "vertex_sa" {
@@ -92,7 +92,7 @@ resource "google_project_iam_member" "vertex_ai_sa_permissions" {
   project     = each.value.project
   role        = each.value.role
   member      = "serviceAccount:service-${data.google_project.projects[split("_", each.key)[0]].number}@gcp-sa-aiplatform-re.iam.gserviceaccount.com"
-  depends_on  = [resource.google_project_service.shared_services, resource.google_project_service_identity.vertex_sa]
+  depends_on  = [resource.google_project_service.deploy_project_services, resource.google_project_service_identity.vertex_sa]
 }
 {% endif %}
 
@@ -101,14 +101,14 @@ resource "google_service_account_iam_member" "cicd_run_invoker_token_creator" {
   service_account_id = google_service_account.cicd_runner_sa.name
   role               = "roles/iam.serviceAccountTokenCreator"
   member             = "serviceAccount:${resource.google_service_account.cicd_runner_sa.email}"
-  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 # Special assignment: Allow the CICD SA to impersonate himself for trigger creation
 resource "google_service_account_iam_member" "cicd_run_invoker_account_user" {
   service_account_id = google_service_account.cicd_runner_sa.name
   role               = "roles/iam.serviceAccountUser"
   member             = "serviceAccount:${resource.google_service_account.cicd_runner_sa.email}"
-  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 {%- if cookiecutter.data_ingestion %}
@@ -125,6 +125,6 @@ resource "google_project_iam_member" "vertexai_pipeline_sa_roles" {
   project    = each.value.project
   role       = each.value.role
   member     = "serviceAccount:${google_service_account.vertexai_pipeline_app_sa[split(",", each.key)[0]].email}"
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {%- endif %}

src/base_template/deployment/terraform/locals.tf

@@ -23,7 +23,7 @@ locals {
     "cloudtrace.googleapis.com"
   ]
 
-  shared_services = [
+  deploy_project_services = [
     "aiplatform.googleapis.com",
     "run.googleapis.com",
     "discoveryengine.googleapis.com",
@@ -32,7 +32,14 @@ locals {
     "bigquery.googleapis.com",
     "serviceusage.googleapis.com",
     "logging.googleapis.com",
-    "cloudtrace.googleapis.com"
+    "cloudtrace.googleapis.com",
+{%- if "adk" in cookiecutter.tags and cookiecutter.session_type == "alloydb" %}
+    "compute.googleapis.com",
+    "servicenetworking.googleapis.com",
+    "alloydb.googleapis.com",
+    "secretmanager.googleapis.com",
+    "dns.googleapis.com"
+{%- endif %}
   ]
 
   deploy_project_ids = {

src/base_template/deployment/terraform/log_sinks.tf

@@ -26,7 +26,7 @@ resource "google_bigquery_dataset" "feedback_dataset" {
   dataset_id    = replace("${var.project_name}_feedback", "-", "_")
   friendly_name = "${var.project_name}_feedback"
   location      = var.region
-  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 resource "google_bigquery_dataset" "telemetry_logs_dataset" {
@@ -35,7 +35,7 @@ resource "google_bigquery_dataset" "telemetry_logs_dataset" {
   dataset_id    = replace("${var.project_name}_telemetry", "-", "_")
   friendly_name = "${var.project_name}_telemetry"
   location      = var.region
-  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 module "log_export_to_bigquery" {

src/base_template/deployment/terraform/service_accounts.tf

@@ -16,7 +16,7 @@ resource "google_service_account" "cicd_runner_sa" {
   account_id   = "${var.project_name}-cb"
   display_name = "CICD Runner SA"
   project      = var.cicd_runner_project_id
-  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% if cookiecutter.deployment_target == 'cloud_run' %}
 resource "google_service_account" "cloud_run_app_sa" {
@@ -25,7 +25,7 @@ resource "google_service_account" "cloud_run_app_sa" {
   account_id   = "${var.project_name}-cr"
   display_name = "Cloud Run Generative AI app SA"
   project      = each.value
-  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% endif %}
 
@@ -37,6 +37,6 @@ resource "google_service_account" "vertexai_pipeline_app_sa" {
   account_id   = "${var.project_name}-rag"
   display_name = "Vertex AI Pipeline app SA"
   project      = each.value
-  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on   = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% endif %}

src/base_template/deployment/terraform/storage.tf

@@ -23,7 +23,7 @@ resource "google_storage_bucket" "bucket_load_test_results" {
   project                     = var.cicd_runner_project_id
   uniform_bucket_level_access = true
   force_destroy               = true
-  depends_on                  = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on                  = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 resource "google_storage_bucket" "logs_data_bucket" {
@@ -34,7 +34,7 @@ resource "google_storage_bucket" "logs_data_bucket" {
   uniform_bucket_level_access = true
   force_destroy               = true
 
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% if cookiecutter.deployment_target == 'cloud_run' %}
 resource "google_artifact_registry_repository" "repo-artifacts-genai" {
@@ -43,7 +43,7 @@ resource "google_artifact_registry_repository" "repo-artifacts-genai" {
   description   = "Repo for Generative AI applications"
   format        = "DOCKER"
   project       = var.cicd_runner_project_id
-  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on    = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 {% endif %}
 
@@ -56,7 +56,7 @@ resource "google_storage_bucket" "data_ingestion_pipeline_gcs_root" {
   uniform_bucket_level_access = true
   force_destroy               = true
 
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 {% if cookiecutter.datastore_type == "vertex_ai_search" %}
@@ -71,7 +71,7 @@ resource "google_discovery_engine_data_store" "data_store_staging" {
   solution_types              = ["SOLUTION_TYPE_SEARCH"]
   create_advanced_site_search = false
   provider                    = google.staging_billing_override
-  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 resource "google_discovery_engine_search_engine" "search_engine_staging" {
@@ -97,7 +97,7 @@ resource "google_discovery_engine_data_store" "data_store_prod" {
   solution_types              = ["SOLUTION_TYPE_SEARCH"]
   create_advanced_site_search = false
   provider                    = google.prod_billing_override
-  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on         = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 resource "google_discovery_engine_search_engine" "search_engine_prod" {
@@ -122,7 +122,7 @@ resource "google_storage_bucket" "vector_search_data_bucket" {
   uniform_bucket_level_access = true
   force_destroy               = true
 
-  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]
+  depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
 }
 
 resource "google_vertex_ai_index" "vector_search_index_staging" {

src/base_template/deployment/terraform/variables.tf

@@ -76,6 +76,7 @@ variable "agentengine_sa_roles" {
   default = [
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
     "roles/run.invoker",
+    "roles/secretmanager.secretAccessor",
 {%- endif %}
     "roles/aiplatform.user",
     "roles/discoveryengine.editor",
@@ -84,6 +85,8 @@ variable "agentengine_sa_roles" {
     "roles/storage.admin"
   ]
 }
+{%- if cookiecutter.deployment_target == 'cloud_run' %}
+{%- endif %}
 
 variable "cicd_roles" {
   description = "List of roles to assign to the CICD runner service account in the CICD project"

src/base_template/pyproject.toml

@@ -16,11 +16,12 @@ dependencies = [
 {%- endif %}
     "google-cloud-logging~=3.11.4",
 {%- if cookiecutter.deployment_target == 'cloud_run' %}
-    "google-cloud-aiplatform[evaluation]~=1.99.0",
+    "google-cloud-aiplatform[evaluation]~=1.101.0",
     "fastapi~=0.115.8",
-    "uvicorn~=0.34.0"
+    "uvicorn~=0.34.0",
+    "psycopg2-binary>=2.9.10",
 {%- elif cookiecutter.deployment_target == 'agent_engine' %}
-    "google-cloud-aiplatform[evaluation,agent-engines]~=1.99.0"
+    "google-cloud-aiplatform[evaluation,agent-engines]~=1.101.0"
 {%- endif %}
 ]
 {% if cookiecutter.deployment_target == 'cloud_run' %}

src/cli/commands/create.py

@@ -32,6 +32,7 @@ from ..utils.template import (
     process_template,
     prompt_datastore_selection,
     prompt_deployment_target,
+    prompt_session_type_selection,
 )
 
 console = Console()
@@ -93,6 +94,11 @@ def normalize_project_name(project_name: str) -> str:
     type=click.Choice(DATASTORE_TYPES),
     help="Type of datastore to use for data ingestion (requires --include-data-ingestion)",
 )
+@click.option(
+    "--session-type",
+    type=click.Choice(["in_memory", "alloydb"]),
+    help="Type of session storage to use",
+)
 @click.option("--debug", is_flag=True, help="Enable debug logging")
 @click.option(
     "--output-dir",
@@ -122,6 +128,7 @@ def create(
     deployment_target: str | None,
     include_data_ingestion: bool,
     datastore: str | None,
+    session_type: str | None,
     debug: bool,
     output_dir: str | None,
     auto_approve: bool,
@@ -232,6 +239,36 @@ def create(
         if debug:
             logging.debug(f"Selected deployment target: {final_deployment}")
 
+        # Session type validation and selection (only for agents that require session management)
+        final_session_type = session_type
+
+        # Check if agent requires session management
+        requires_session = config.get("settings", {}).get("requires_session", False)
+
+        if requires_session:
+            if final_deployment == "agent_engine" and session_type:
+                console.print(
+                    "Error: --session-type cannot be used with agent_engine deployment target. "
+                    "Agent Engine handles session management internally.",
+                    style="bold red",
+                )
+                return
+
+            if final_deployment == "cloud_run" and not session_type:
+                final_session_type = prompt_session_type_selection()
+        else:
+            # Agents that don't require session management always use in-memory sessions
+            final_session_type = "in_memory"
+            if session_type and session_type != "in_memory":
+                console.print(
+                    "Warning: Session type options are only available for agents that require session management. "
+                    "Using in-memory sessions for this agent.",
+                    style="yellow",
+                )
+
+        if debug and final_session_type:
+            logging.debug(f"Selected session type: {final_session_type}")
+
         # Region confirmation (if not explicitly passed)
         if (
             not auto_approve
@@ -289,6 +326,7 @@ def create(
             deployment_target=final_deployment,
             include_data_ingestion=include_data_ingestion,
             datastore=datastore,
+            session_type=final_session_type,
             output_dir=destination_dir,
         )
 

src/cli/commands/setup_cicd.py

@@ -168,8 +168,8 @@ def update_build_triggers(tf_dir: Path) -> None:
 
         # Add repository dependency to all trigger resources
         modified_content = content.replace(
-            "depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services]",
-            "depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.shared_services, google_cloudbuildv2_repository.repo]",
+            "depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]",
+            "depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services, google_cloudbuildv2_repository.repo]",
         )
 
         # Update repository reference in all triggers

src/cli/utils/cicd.py

@@ -212,7 +212,10 @@ def create_github_connection(
                 return secret_id, app_installation_id
             elif status == "PENDING_USER_OAUTH" or status == "PENDING_INSTALL_APP":
                 if attempt < max_retries - 1:  # Don't print waiting on last attempt
-                    console.print("⏳ Waiting for GitHub authorization...")
+                    console.print("⏳ Waiting for authorization...")
+                    print(
+                        f"Console: https://console.cloud.google.com/cloud-build/repositories?project={project_id}"
+                    )
                     # Extract and print the action URI for user authentication
                     try:
                         action_uri = (
@@ -225,9 +228,9 @@ def create_github_connection(
                                 "\n🔑 Authentication Required:", style="bold yellow"
                             )
                             console.print(
-                                f"Please visit [link={action_uri}]this page[/link] to authenticate Cloud Build with GitHub:"
+                                f"Please visit [link={action_uri}][bold blue]this page[/bold blue][/link] to authenticate Cloud Build with GitHub:"
                             )
-                            console.print(f"{action_uri}", style="bold blue")
+                            print(f"\n{action_uri}\n")
                             console.print(
                                 "(Copy and paste the link into your browser if clicking doesn't work)"
                             )

src/cli/utils/template.py

@@ -206,6 +206,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:
@@ -365,6 +395,7 @@ 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,
 ) -> None:
     """Process the template directory and create a new project.
@@ -375,6 +406,8 @@ 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
     """
     logging.debug(f"Processing template from {template_dir}")
@@ -467,15 +500,6 @@ def process_template(
                             agent_folder, project_folder, agent_name, overwrite=True
                         )
 
-            # Copy agent README.md if it exists
-            agent_readme = agent_path / "README.md"
-            if agent_readme.exists():
-                agent_readme_dest = project_template / "agent_README.md"
-                shutil.copy2(agent_readme, agent_readme_dest)
-                logging.debug(
-                    f"Copied agent README from {agent_readme} to {agent_readme_dest}"
-                )
-
             # Load and validate template config first
             template_path = pathlib.Path(template_dir)
             config = load_template_config(template_path)
@@ -536,6 +560,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,

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}}"