feat: add Vertex AI Search and Vector Search data connectors for agentic_rag (#791)

* feat: replace Vertex AI Pipeline data ingestion with GCS Data Connector for vertex_ai_search

- Switch agentic_rag agent from LangChain retriever to native ADK VertexAiSearchTool
  for vertex_ai_search datastore type
- Replace Vertex AI Pipeline-based ingestion with GCS Data Connector approach using
  Discovery Engine setUpDataConnectorV2 API
- Add Terraform null_resource + external data source pattern for managing data
  connectors and retrieving auto-created data store IDs
- Add shell scripts for connector setup, data store ID retrieval, and on-demand sync
- Add Makefile sync-data target for triggering manual data syncs
- Skip data_ingestion pipeline file copy for vertex_ai_search projects
- Update CI/CD pipelines to only include pipeline steps for vertex_ai_vector_search
- Update deployment target service.tf files to use external data source references

* feat: add Vertex AI Search and Vector Search data connectors for agentic_rag

- Migrate data ingestion from shared location to agent-specific `agents/agentic_rag/data_ingestion/`
- Replace shell scripts with Python scripts for data connector setup and management
- Add Vector Search 2.0 Collections API support alongside Vertex AI Search
- Add `setup-datastore` CLI command and sample data for both datastore types
- Use dedicated asp-rag GCP projects for agentic_rag e2e tests
- Conditionally enable vectorsearch API only when datastore type requires it
- Add --wait flag to connector run for blocking sync support

Author: eliasecchig

.cloudbuild/terraform/apis.tf

@@ -25,6 +25,19 @@ resource "google_project_service" "cloud_resource_manager_api" {
   disable_on_destroy = false
 }
 
+# Enable Cloud Resource Manager API for each project in e2e_rag_project_mapping
+resource "google_project_service" "cloud_resource_manager_api_rag" {
+  for_each = {
+    "dev"     = var.e2e_rag_project_mapping.dev
+    "staging" = var.e2e_rag_project_mapping.staging
+    "prod"    = var.e2e_rag_project_mapping.prod
+  }
+
+  project            = each.value
+  service            = "cloudresourcemanager.googleapis.com"
+  disable_on_destroy = false
+}
+
 # Enable Cloud Scheduler API for scheduled cleanup jobs
 resource "google_project_service" "cloud_scheduler_api" {
   project            = var.cicd_runner_project_id

.cloudbuild/terraform/build_triggers.tf

@@ -33,7 +33,7 @@ locals {
     "agent_starter_pack/agents/**",
     "agent_starter_pack/cli/**",
     "tests/**",
-    "agent_starter_pack/data_ingestion/**",
+    "agent_starter_pack/agents/agentic_rag/data_ingestion/**",
     "pyproject.toml",
     "uv.lock",
     ".cloudbuild/**",
@@ -42,7 +42,7 @@ locals {
   lint_templated_agents_included_files = [
     "agent_starter_pack/cli/**",
     "agent_starter_pack/base_templates/**",
-    "agent_starter_pack/data_ingestion/**",
+    "agent_starter_pack/agents/agentic_rag/data_ingestion/**",
     "agent_starter_pack/deployment_targets/**",
     "tests/integration/test_template_linting.py",
     "tests/integration/test_templated_patterns.py",
@@ -82,11 +82,11 @@ locals {
     },
     {
       name  = "agentic_rag-agent_engine-vertex_ai_search"
-      value = "agentic_rag,agent_engine,--include-data-ingestion,--datastore,vertex_ai_search"
+      value = "agentic_rag,agent_engine,--datastore,vertex_ai_search"
     },
     {
       name  = "agentic_rag-cloud_run-vertex_ai_vector_search"
-      value = "agentic_rag,cloud_run,--include-data-ingestion,--datastore,vertex_ai_vector_search"
+      value = "agentic_rag,cloud_run,--datastore,vertex_ai_vector_search"
     },
     {
       name  = "adk_live-agent_engine"
@@ -237,7 +237,7 @@ locals {
     },
     {
       name  = "agentic_rag-agent_engine-vertex_ai_search-github"
-      value = "agentic_rag,agent_engine,--include-data-ingestion,--datastore,vertex_ai_search,--cicd-runner,github_actions"
+      value = "agentic_rag,agent_engine,--datastore,vertex_ai_search,--cicd-runner,github_actions"
     },
     {
       name  = "adk_live-agent_engine-github"
@@ -257,11 +257,11 @@ locals {
     },
     {
       name  = "agentic_rag-agent_engine-vertex_ai_search"
-      value = "agentic_rag,agent_engine,--include-data-ingestion,--datastore,vertex_ai_search"
+      value = "agentic_rag,agent_engine,--datastore,vertex_ai_search"
     },
     {
       name  = "agentic_rag-cloud_run-vertex_ai_vector_search"
-      value = "agentic_rag,cloud_run,--include-data-ingestion,--datastore,vertex_ai_vector_search"
+      value = "agentic_rag,cloud_run,--datastore,vertex_ai_vector_search"
     },
     {
       name  = "adk_live-agent_engine"
@@ -388,7 +388,7 @@ locals {
         "pyproject.toml",
         ] : substr(combo.name, 0, 11) == "agentic_rag" ? [
         "agent_starter_pack/agents/agentic_rag/**",
-        "agent_starter_pack/data_ingestion/**",
+        "agent_starter_pack/agents/agentic_rag/data_ingestion/**",
         "pyproject.toml",
         ] : substr(combo.name, 0, 8) == "adk_live" ? [
         "agent_starter_pack/agents/adk_live/**",
@@ -401,7 +401,7 @@ locals {
         # Shared and Python base templates only (not Go/Java/TypeScript)
         "agent_starter_pack/base_templates/_shared/**",
         "agent_starter_pack/base_templates/python/**",
-        "agent_starter_pack/data_ingestion/**",
+        "agent_starter_pack/agents/agentic_rag/data_ingestion/**",
         # Python deployment targets only (not Go/Java/TypeScript)
         "agent_starter_pack/deployment_targets/agent_engine/_shared/**",
         "agent_starter_pack/deployment_targets/agent_engine/python/**",
@@ -562,9 +562,9 @@ resource "google_cloudbuild_trigger" "main_e2e_deployment_test" {
 
   substitutions = {
     _TEST_AGENT_COMBINATION = each.value.value
-    _E2E_DEV_PROJECT        = var.e2e_test_project_mapping.dev
-    _E2E_STAGING_PROJECT    = var.e2e_test_project_mapping.staging
-    _E2E_PROD_PROJECT       = var.e2e_test_project_mapping.prod
+    _E2E_DEV_PROJECT        = startswith(each.key, "agentic_rag") ? var.e2e_rag_project_mapping.dev : var.e2e_test_project_mapping.dev
+    _E2E_STAGING_PROJECT    = startswith(each.key, "agentic_rag") ? var.e2e_rag_project_mapping.staging : var.e2e_test_project_mapping.staging
+    _E2E_PROD_PROJECT       = startswith(each.key, "agentic_rag") ? var.e2e_rag_project_mapping.prod : var.e2e_test_project_mapping.prod
     _SECRETS_PROJECT_ID     = "asp-e2e-vars"
     _COMMIT_MESSAGE         = "$(push.head_commit.message)"
   }

.cloudbuild/terraform/service_account.tf

@@ -60,6 +60,25 @@ resource "google_project_iam_member" "cicd_runner_e2e_project_roles" {
   member  = "serviceAccount:${google_service_account.cicd_runner_sa.email}"
 }
 
+# Grant permissions to the service account for RAG E2E project environments
+resource "google_project_iam_member" "cicd_runner_rag_project_roles" {
+  for_each = {
+    for idx, proj_role in flatten([
+      for env, project_id in var.e2e_rag_project_mapping : [
+        for role in local.e2e_project_roles : {
+          project = project_id
+          env     = env
+          role    = role
+        }
+      ]
+    ]) : "${proj_role.env}-${proj_role.role}" => proj_role
+  }
+
+  project = each.value.project
+  role    = each.value.role
+  member  = "serviceAccount:${google_service_account.cicd_runner_sa.email}"
+}
+
 # Grant owner permissions to the service account for all cleanup projects
 resource "google_project_iam_member" "cicd_runner_cleanup_project_roles" {
   for_each = toset(var.cleanup_project_ids)

.cloudbuild/terraform/variables.tf

@@ -42,6 +42,15 @@ variable "e2e_test_project_mapping" {
   })
 }
 
+variable "e2e_rag_project_mapping" {
+  description = "Mapping of project IDs for agentic_rag E2E tests (separate projects with datastore resources)"
+  type = object({
+    dev     = string
+    prod    = string
+    staging = string
+  })
+}
+
 variable "cleanup_project_ids" {
   description = "List of all project IDs that need cleanup (for scheduled cleanup job)"
   type        = list(string)
@@ -51,6 +60,9 @@ variable "cleanup_project_ids" {
     "asp-e2e-prd",
     "asp-test-dev",
     "asp-test-prd",
-    "asp-test-stg"
+    "asp-test-stg",
+    "asp-rag-dev",
+    "asp-rag-stg",
+    "asp-rag-prd"
   ]
 }

.cloudbuild/terraform/vars/env.tfvars

@@ -11,3 +11,9 @@ e2e_test_project_mapping = {
   staging = "asp-e2e-stg"
   prod    = "asp-e2e-prd"
 }
+
+e2e_rag_project_mapping = {
+  dev     = "asp-rag-dev"
+  staging = "asp-rag-stg"
+  prod    = "asp-rag-prd"
+}

GEMINI.md

@@ -707,7 +707,6 @@ asp_version = "0.25.0"
 deployment_target = "cloud_run"
 session_type = "in_memory"
 cicd_runner = "skip"
-include_data_ingestion = false
 ```
 
 The `create_params` section enables the `enhance` command to recreate identical scaffolding with the locked ASP version.

agent_starter_pack/agents/agentic_rag/.template/templateconfig.yaml

@@ -20,6 +20,7 @@ settings:
   deployment_targets: ["agent_engine", "cloud_run", "none"]
   extra_dependencies: [
     "google-adk>=1.15.0,<2.0.0",
+    "google-cloud-vectorsearch",
     "langchain-google-vertexai~=2.0.7",
     "langchain~=0.3.24",
     "langchain-core~=0.3.55",

agent_starter_pack/agents/agentic_rag/app/agent.py

@@ -29,13 +29,15 @@
 )
 from google.cloud import bigquery
 {%- endif %}
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+from google.adk.tools import VertexAiSearchTool
+{%- endif %}
 from google.genai import types
-from langchain_google_vertexai import VertexAIEmbeddings
+{%- if cookiecutter.datastore_type == "vertex_ai_vector_search" %}
 
-from {{cookiecutter.agent_directory}}.retrievers import get_compressor, get_retriever
-from {{cookiecutter.agent_directory}}.templates import format_docs
+from {{cookiecutter.agent_directory}}.retrievers import search_collection
+{%- endif %}
 
-EMBEDDING_MODEL = "text-embedding-005"
 LLM_LOCATION = "global"
 LOCATION = "us-central1"
 LLM = "gemini-3-flash-preview"
@@ -46,47 +48,22 @@
 os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
 
 vertexai.init(project=project_id, location=LOCATION)
-embedding = VertexAIEmbeddings(
-    project=project_id, location=LOCATION, model_name=EMBEDDING_MODEL
-)
 
 {% if cookiecutter.datastore_type == "vertex_ai_search" %}
-EMBEDDING_COLUMN = "embedding"
-TOP_K = 5
-
-data_store_region = os.getenv("DATA_STORE_REGION", "us")
-data_store_id = os.getenv("DATA_STORE_ID", "{{cookiecutter.project_name}}-datastore")
-
-retriever = get_retriever(
-    project_id=project_id,
-    data_store_id=data_store_id,
-    data_store_region=data_store_region,
-    embedding=embedding,
-    embedding_column=EMBEDDING_COLUMN,
-    max_documents=10,
+data_store_region = os.getenv("DATA_STORE_REGION", "global")
+data_store_id = os.getenv(
+    "DATA_STORE_ID", "{{cookiecutter.project_name}}-collection_documents"
 )
-{% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
-vector_search_index = os.getenv(
-    "VECTOR_SEARCH_INDEX", "{{cookiecutter.project_name}}-vector-search"
-)
-vector_search_index_endpoint = os.getenv(
-    "VECTOR_SEARCH_INDEX_ENDPOINT", "{{cookiecutter.project_name}}-vector-search-endpoint"
-)
-vector_search_bucket = os.getenv(
-    "VECTOR_SEARCH_BUCKET", f"{project_id}-{{cookiecutter.project_name}}-vs"
+data_store_path = (
+    f"projects/{project_id}/locations/{data_store_region}"
+    f"/collections/default_collection/dataStores/{data_store_id}"
 )
 
-retriever = get_retriever(
-    project_id=project_id,
-    region=LOCATION,
-    vector_search_bucket=vector_search_bucket,
-    vector_search_index=vector_search_index,
-    vector_search_index_endpoint=vector_search_index_endpoint,
-    embedding=embedding,
-)
-{% endif %}
-compressor = get_compressor(
-    project_id=project_id,
+vertex_search_tool = VertexAiSearchTool(data_store_id=data_store_path)
+{% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
+vector_search_collection = os.getenv(
+    "VECTOR_SEARCH_COLLECTION",
+    f"projects/{project_id}/locations/{LOCATION}/collections/{{cookiecutter.project_name}}-collection",
 )
 
 
@@ -99,22 +76,19 @@ def retrieve_docs(query: str) -> str:
         query (str): The user's question or search query.
 
     Returns:
-        str: Formatted string containing relevant document content retrieved and ranked based on the query.
+        str: Formatted string containing relevant document content.
     """
     try:
-        # Use the retriever to fetch relevant documents based on the query
-        retrieved_docs = retriever.invoke(query)
-        # Re-rank docs with Vertex AI Rank for better relevance
-        ranked_docs = compressor.compress_documents(
-            documents=retrieved_docs, query=query
+        return search_collection(
+            query=query,
+            collection_path=vector_search_collection,
         )
-        # Format ranked documents into a consistent structure for LLM consumption
-        formatted_docs = format_docs.format(docs=ranked_docs)
     except Exception as e:
-        return f"Calling retrieval tool with query:\n\n{query}\n\nraised the following error:\n\n{type(e)}: {e}"
-
-    return formatted_docs
-
+        return (
+            f"Calling retrieval tool with query:\n\n{query}\n\n"
+            f"raised the following error:\n\n{type(e)}: {e}"
+        )
+{% endif %}
 
 instruction = """You are an AI assistant for question-answering tasks.
 Answer to the best of your ability using the context provided.
@@ -129,7 +103,11 @@ def retrieve_docs(query: str) -> str:
         retry_options=types.HttpRetryOptions(attempts=3),
     ),
     instruction=instruction,
+{%- if cookiecutter.datastore_type == "vertex_ai_search" %}
+    tools=[vertex_search_tool],
+{%- elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
     tools=[retrieve_docs],
+{%- endif %}
 )
 
 {%- if cookiecutter.bq_analytics %}