fix: 07-rag-setup.sh — correct Serverless syntax, try RagManagedDb then plain fallback

infrastructure/07-rag-setup.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
-# Serverless mode does NOT require allowlist (unlike Spanner mode)
+# Serverless = RagManagedDb (no Spanner, no allowlist needed)
 # Idempotent — safe to run multiple times
 # Source .env before running: source .env
 
@@ -25,7 +25,7 @@ gcloud services enable aiplatform.googleapis.com \
   --project="${PROJECT_ID}" --quiet
 echo "✓ APIs enabled"
 
-# ── GCS bucket for corpus source documents (idempotent) ──────────────────
+# ── GCS bucket (idempotent) ──────────────────────────────────────────────
 CORPUS_BUCKET="${PROJECT_ID}-agent-corpus"
 if ! gsutil ls -b "gs://${CORPUS_BUCKET}" &>/dev/null; then
   gsutil mb -l "${REGION}" -b on "gs://${CORPUS_BUCKET}"
@@ -34,17 +34,17 @@ else
   echo "✓ GCS corpus bucket exists: gs://${CORPUS_BUCKET}"
 fi
 
-# ── Upload seed documents if present ─────────────────────────────────────
+# ── Upload seed docs ──────────────────────────────────────────────────────────
 SEED_DIR="$(dirname "$0")/../docs/corpus-seed"
 if [[ -d "${SEED_DIR}" ]] && ls "${SEED_DIR}"/*.md &>/dev/null; then
   gsutil -m cp "${SEED_DIR}"/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
   echo "✓ Seed documents uploaded to gs://${CORPUS_BUCKET}/seed/"
 else
-  echo "  No seed documents found in docs/corpus-seed/ — skipping upload"
+  echo "  No seed documents in docs/corpus-seed/ — skipping"
 fi
 
-# ── Create corpus in Serverless mode ─────────────────────────────────────
-python3 - << PYEOF
+# ── Python: create corpus ───────────────────────────────────────────────────────
+python3 - << 'PYEOF'
 import os, sys
 
 try:
@@ -52,17 +52,16 @@ try:
     from vertexai.preview import rag
 except ImportError:
     print("ERROR: google-cloud-aiplatform not installed.")
-    print("Run: pip install google-cloud-aiplatform>=1.87.0")
     sys.exit(1)
 
-PROJECT_ID   = os.environ["PROJECT_ID"]
-RAG_REGION   = os.environ.get("RAG_REGION", os.environ["REGION"])
-DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
+PROJECT_ID    = os.environ["PROJECT_ID"]
+RAG_REGION    = os.environ.get("RAG_REGION", os.environ["REGION"])
+DISPLAY_NAME  = os.environ["RAG_CORPUS_DISPLAY_NAME"]
 CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"
 
 vertexai.init(project=PROJECT_ID, location=RAG_REGION)
 
-# Check for existing corpus
+# Check if corpus already exists
 corpus = None
 try:
     for c in rag.list_corpora():
@@ -74,20 +73,43 @@ 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}...")
+    print(f"  Creating corpus '{DISPLAY_NAME}' 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 (Serverless): {corpus.name}")
+
+    # Try Serverless mode (RagManagedDb) first, fall back to plain create
+    created = False
+    for attempt in ["serverless", "plain"]:
+        try:
+            if attempt == "serverless":
+                # SDK >= 1.87: RagVectorDbConfig with rag_managed_db
+                try:
+                    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,
+                    )
+                except TypeError:
+                    # Older SDK: RagManagedDb not a kwarg — skip to plain
+                    raise
+            else:
+                # Plain create — lets Google pick default (Serverless on new projects)
+                corpus = rag.create_corpus(
+                    display_name=DISPLAY_NAME,
+                    embedding_model_config=embedding_config,
+                )
+            print(f"✓ RAG corpus created [{attempt}]: {corpus.name}")
+            created = True
+            break
+        except Exception as e:
+            if attempt == "plain":
+                print(f"ERROR: Could not create corpus: {e}")
+                sys.exit(1)
+            print(f"  [{attempt}] failed: {e} — retrying with plain...")
 
 # Import seed documents (non-fatal)
 gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"
@@ -111,7 +133,7 @@ print(f"  Corpus resource name : {corpus.name}")
 print(f"  Region               : {RAG_REGION}")
 print(f"")
 print(f"  ACTION REQUIRED — add to .env:")
-print(f"    export RAG_CORPUS_NAME=\"{corpus.name}\"")
+print(f'    export RAG_CORPUS_NAME="{corpus.name}"')
 PYEOF
 
 echo ""