fix: 07-rag-setup.sh — upgrade SDK + multi-attempt serverless corpus creation

infrastructure/07-rag-setup.sh

@@ -25,6 +25,11 @@ gcloud services enable aiplatform.googleapis.com \
   --project="${PROJECT_ID}" --quiet
 echo "✓ APIs enabled"
 
+# ── Upgrade SDK to ensure RagManagedDb support ───────────────────────────────
+echo "  Upgrading google-cloud-aiplatform SDK..."
+pip install --quiet --upgrade google-cloud-aiplatform
+echo "✓ SDK upgraded"
+
 # ── GCS bucket (idempotent) ──────────────────────────────────────────────
 CORPUS_BUCKET="${PROJECT_ID}-agent-corpus"
 if ! gsutil ls -b "gs://${CORPUS_BUCKET}" &>/dev/null; then
@@ -43,7 +48,7 @@ else
   echo "  No seed documents in docs/corpus-seed/ — skipping"
 fi
 
-# ── Python: create corpus ───────────────────────────────────────────────────────
+# ── Python: create corpus ────────────────────────────────────────────────────
 python3 - << 'PYEOF'
 import os, sys
 
@@ -78,38 +83,75 @@ if corpus is None:
         publisher_model="publishers/google/models/text-embedding-005"
     )
 
-    # Try Serverless mode (RagManagedDb) first, fall back to plain create
+    # Attempt order:
-    created = False
+    # 1. SDK >= 1.87 : RagVectorDbConfig(rag_managed_db=RagManagedDb())
-    for attempt in ["serverless", "plain"]:
+    # 2. Older SDK   : RagVectorDbConfig(rag_managed_db=RagManagedDbConfig())
+    # 3. REST fallback: gapic-style with vector_db proto dict
+    attempts = ["new_sdk", "old_sdk", "proto_dict"]
+    for attempt in attempts:
         try:
-            if attempt == "serverless":
+            if attempt == "new_sdk":
-                # SDK >= 1.87: RagVectorDbConfig with rag_managed_db
+                vector_db = rag.RagVectorDbConfig(
-                try:
+                    rag_managed_db=rag.RagManagedDb()
-                    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,
+                    vector_db=vector_db,
                 )
+
+            elif attempt == "old_sdk":
+                # Some SDK versions use RagManagedDbConfig instead
+                RagManagedDbConfig = getattr(rag, "RagManagedDbConfig", None)
+                if RagManagedDbConfig is None:
+                    raise AttributeError("RagManagedDbConfig not in this SDK version")
+                vector_db = rag.RagVectorDbConfig(
+                    rag_managed_db=RagManagedDbConfig()
+                )
+                corpus = rag.create_corpus(
+                    display_name=DISPLAY_NAME,
+                    embedding_model_config=embedding_config,
+                    vector_db=vector_db,
+                )
+
+            elif attempt == "proto_dict":
+                # Last resort: pass vector_db as a plain dict understood by gapic
+                from google.cloud.aiplatform_v1beta1.types import (
+                    RagCorpus, RagVectorDbConfig, RagManagedDb,
+                    RagEmbeddingModelConfig
+                )
+                from google.cloud import aiplatform_v1beta1 as aip
+
+                client = aip.VertexRagDataServiceClient(
+                    client_options={"api_endpoint": f"{RAG_REGION}-aiplatform.googleapis.com"}
+                )
+                parent = f"projects/{PROJECT_ID}/locations/{RAG_REGION}"
+                rag_corpus = RagCorpus(
+                    display_name=DISPLAY_NAME,
+                    rag_embedding_model_config=RagEmbeddingModelConfig(
+                        vertex_prediction_endpoint=RagEmbeddingModelConfig.VertexPredictionEndpoint(
+                            publisher_model="publishers/google/models/text-embedding-005"
+                        )
+                    ),
+                    rag_vector_db_config=RagVectorDbConfig(
+                        rag_managed_db=RagManagedDb()
+                    ),
+                )
+                op = client.create_rag_corpus(parent=parent, rag_corpus=rag_corpus)
+                result = op.result()
+                # Wrap in a simple namespace so rest of script works
+                class _C:
+                    name = result.name
+                corpus = _C()
+
             print(f"✓ RAG corpus created [{attempt}]: {corpus.name}")
-            created = True
             break
+
         except Exception as e:
-            if attempt == "plain":
+            if attempt == attempts[-1]:
-                print(f"ERROR: Could not create corpus: {e}")
+                print(f"ERROR: All attempts failed. Last error: {e}")
                 sys.exit(1)
-            print(f"  [{attempt}] failed: {e} — retrying with plain...")
+            print(f"  [{attempt}] failed: {e} — trying next method...")
 
 # Import seed documents (non-fatal)
 gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"