fix: 07-rag-setup.sh — switch to Serverless mode (no allowlist required)

infrastructure/07-rag-setup.sh

@@ -1,13 +1,7 @@
 #!/usr/bin/env bash
-# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus and upload initial documents
+# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
+# Serverless mode does NOT require allowlist (unlike Spanner mode)
 # Idempotent — safe to run multiple times
-# Source: Vertex AI RAG Engine SDK (google-cloud-aiplatform >= 1.87)
-#
-# REGIONAL AVAILABILITY (OQ-02):
-#   RAG Engine in us-central1 requires allowlist access.
-#   This script auto-falls-back to us-east1 if us-central1 is not accessible.
-#   To bypass: export RAG_REGION=us-east1 before running.
-#
 # Source .env before running: source .env
 
 set -euo pipefail
@@ -16,13 +10,11 @@ set -euo pipefail
 : "${REGION:?Set REGION}"
 : "${RAG_CORPUS_DISPLAY_NAME:?Set RAG_CORPUS_DISPLAY_NAME}"
 
-# RAG region fallback: prefer env override, then REGION, fallback to us-east1
 RAG_REGION="${RAG_REGION:-${REGION}}"
-RAG_FALLBACK_REGION="us-east1"
 
-echo "=== 07: Setting up Vertex AI RAG Engine ==="
+echo "=== 07: Setting up Vertex AI RAG Engine (Serverless mode) ==="
 echo "    Project  : ${PROJECT_ID}"
-echo "    Region   : ${RAG_REGION} (fallback: ${RAG_FALLBACK_REGION})"
+echo "    Region   : ${RAG_REGION}"
 echo "    Corpus   : ${RAG_CORPUS_DISPLAY_NAME}"
 echo ""
 
@@ -31,7 +23,6 @@ bash "$(dirname "$0")/00-authcheck.sh"
 gcloud services enable aiplatform.googleapis.com \
   storage.googleapis.com \
   --project="${PROJECT_ID}" --quiet
-
 echo "✓ APIs enabled"
 
 # ── GCS bucket for corpus source documents (idempotent) ──────────────────
@@ -52,7 +43,7 @@ else
   echo "  No seed documents found in docs/corpus-seed/ — skipping upload"
 fi
 
-# ── Run Python to create/update corpus ───────────────────────────────────
+# ── Create corpus in Serverless mode ─────────────────────────────────────
 python3 - << PYEOF
 import os, sys
 
@@ -64,52 +55,39 @@ except ImportError:
     print("Run: pip install google-cloud-aiplatform>=1.87.0")
     sys.exit(1)
 
-PROJECT_ID        = os.environ["PROJECT_ID"]
+PROJECT_ID   = os.environ["PROJECT_ID"]
-RAG_REGION        = os.environ.get("RAG_REGION", os.environ["REGION"])
+RAG_REGION   = os.environ.get("RAG_REGION", os.environ["REGION"])
-FALLBACK_REGION   = "us-east1"
+DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
-DISPLAY_NAME      = os.environ["RAG_CORPUS_DISPLAY_NAME"]
+CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"
-CORPUS_BUCKET     = f"{PROJECT_ID}-agent-corpus"
 
-def try_init_and_list(region):
+vertexai.init(project=PROJECT_ID, location=RAG_REGION)
-    vertexai.init(project=PROJECT_ID, location=region)
-    try:
-        return list(rag.list_corpora()), region
-    except Exception as e:
-        if "PERMISSION_DENIED" in str(e) or "allowlist" in str(e).lower() or "not found" in str(e).lower():
-            return None, region
-        raise
 
-# Try primary region, fall back if needed
+# Check for existing corpus
-corpora, active_region = try_init_and_list(RAG_REGION)
-if corpora is None and RAG_REGION != FALLBACK_REGION:
-    print(f"  RAG Engine not available in {RAG_REGION} — falling back to {FALLBACK_REGION}")
-    corpora, active_region = try_init_and_list(FALLBACK_REGION)
-
-if corpora is None:
-    print(f"ERROR: RAG Engine not accessible in {RAG_REGION} or {FALLBACK_REGION}.")
-    print("Apply for allowlist: vertex-ai-rag-engine-support@google.com")
-    print("Or set RAG_REGION=us-east1 in .env")
-    sys.exit(1)
-
-print(f"✓ RAG Engine accessible in region: {active_region}")
-
-# Find or create corpus
 corpus = None
-for c in corpora:
+try:
-    if c.display_name == DISPLAY_NAME:
+    for c in rag.list_corpora():
-        corpus = c
+        if c.display_name == DISPLAY_NAME:
-        print(f"✓ RAG corpus already exists: {c.name}")
+            corpus = c
-        break
+            print(f"✓ RAG corpus already exists: {c.name}")
+            break
+except Exception as e:
+    print(f"WARNING: Could not list corpora: {e}")
 
 if corpus is None:
+    print(f"  Creating corpus in Serverless mode in {RAG_REGION}...")
     embedding_config = rag.EmbeddingModelConfig(
         publisher_model="publishers/google/models/text-embedding-005"
     )
+    # Serverless mode: RagVectorDbConfig with rag_managed_db (no Spanner)
+    vector_db = rag.RagVectorDbConfig(
+        rag_managed_db=rag.RagManagedDb()
+    )
     corpus = rag.create_corpus(
         display_name=DISPLAY_NAME,
         embedding_model_config=embedding_config,
+        vector_db=vector_db,
     )
-    print(f"✓ RAG corpus created: {corpus.name}")
+    print(f"✓ RAG corpus created (Serverless): {corpus.name}")
 
 # Import seed documents (non-fatal)
 gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"
@@ -123,24 +101,20 @@ try:
     )
     print(f"✓ Documents imported from {gcs_uri}")
 except Exception as e:
-    print(f"  WARNING: Document import skipped or failed: {e}")
+    print(f"  WARNING: Document import skipped: {e}")
-    print(f"  Import manually: https://console.cloud.google.com/vertex-ai/rag?project={PROJECT_ID}")
 
-# Persist corpus name for 08-memorybank-setup.sh and agent
 with open("/tmp/rag_corpus_name.txt", "w") as f:
     f.write(corpus.name)
 
 print(f"")
 print(f"  Corpus resource name : {corpus.name}")
-print(f"  Active region        : {active_region}")
+print(f"  Region               : {RAG_REGION}")
 print(f"")
 print(f"  ACTION REQUIRED — add to .env:")
-print(f"    RAG_CORPUS_NAME={corpus.name}")
+print(f"    export RAG_CORPUS_NAME=\"{corpus.name}\"")
-if active_region != os.environ.get("REGION"):
-    print(f"    RAG_REGION={active_region}")
 PYEOF
 
 echo ""
 echo "=== 07: RAG Engine setup COMPLETE ==="
-echo "    View corpus: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
+echo "    View: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
 echo ""