fix(rag): correct corpus payload to vectorDbConfig.ragManagedDb; verify engine config actually switched before creating

- RagCorpus REST body uses vectorDbConfig.ragManagedDb (not ragManagedDbConfig)
- RagManagedDb has no tier field - empty object {} = serverless
- Add GET+verify step after PATCH to confirm spanner key is gone before proceeding
- Print actual engine config state so failures are immediately visible

agents/rag/setup_corpus.py

@@ -3,17 +3,21 @@
 setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
 Project: propane-will-491900-m5
 
-SDK 1.153.1 has a bug where backend_config crashes when provided,
+SDK 1.153.1 has a bug where backend_config crashes when provided.
-and defaults to Spanner when omitted. We bypass it by calling the
+We call the REST API directly for corpus creation.
-REST API directly for corpus creation, then use the SDK for everything else.
 
-Key design decisions:
+v1beta1 REST schema facts (verified from Google docs):
-- ragEngineConfig PATCH always targets us-central1 — it is a project-level
+  RagCorpus.backend_config is a union field:
-  control-plane endpoint that does NOT exist in other regions.
+    vectorDbConfig: RagVectorDbConfig
-- Corpus creation uses ragManagedDbConfig.basic in the RagCorpus body.
+      .ragManagedDb: RagManagedDb  <- serverless managed DB, no tier field
-  The v1beta1 API uses ragManagedDbConfig (not vectorDbConfig) on the corpus
+      .vertexVectorSearch, .pinecone, .weaviate  <- alternatives
-  resource. Sending it explicitly with {"basic":{}} forces serverless mode
+  RagVectorDbConfig has NO 'ragManagedDbConfig' field.
-  and overrides any project-level Spanner default.
+  RagManagedDb has NO 'tier' field — empty {} means serverless/basic.
+
+  ragEngineConfig (project-level) controls the DEFAULT backend when
+  vectorDbConfig is omitted. We must verify it is actually in 'basic'
+  state (not 'spanner') before creating a corpus without an explicit
+  vectorDbConfig, OR we can pass vectorDbConfig.ragManagedDb explicitly.
 """
 
 import os
@@ -44,41 +48,50 @@ def get_token() -> str:
     ).strip()
 
 
-def ensure_serverless_engine_config() -> None:
+def _engine_config_url() -> str:
-    """
+    return (
-    Set project-level RAG Engine Config to basic (serverless) tier.
-    Always targets us-central1. The endpoint does not exist in other regions.
-
-    The v1beta1 ragEngineConfig schema:
-      { "ragManagedDbConfig": { "basic": {} } }   <- serverless
-      { "ragManagedDbConfig": { "scaled": {} } }  <- Spanner / managed DB
-    """
-    print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
-    endpoint = (
         f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
         f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
     )
-    # Explicitly set basic and unset spanner by sending only basic.
+
-    payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
+
+def get_engine_config() -> dict:
+    token = get_token()
+    r = subprocess.run(
+        ["curl", "-s", "-H", f"Authorization: Bearer {token}", _engine_config_url()],
+        capture_output=True, text=True,
+    )
+    return json.loads(r.stdout)
+
+
+def ensure_serverless_engine_config() -> None:
+    """
+    Set project-level RAG Engine Config to basic (serverless) tier.
+    Always targets us-central1.
+
+    The PATCH is a merge by default. We verify the result with a GET
+    to confirm 'spanner' is no longer the active mode before proceeding.
+    """
+    print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
     token = get_token()
 
     result = subprocess.run(
         ["curl", "-s", "-X", "PATCH",
          "-H", f"Authorization: Bearer {token}",
          "-H", "Content-Type: application/json",
-         endpoint, "-d", payload],
+         _engine_config_url(),
+         "-d", json.dumps({"ragManagedDbConfig": {"basic": {}}})],
         capture_output=True, text=True,
     )
     resp = json.loads(result.stdout)
 
     if "error" in resp:
         err = resp["error"]
-        print(f"  \u26a0\ufe0f  Engine config warning ({err.get('code')}): {err.get('message')} \u2014 continuing.")
+        print(f"  \u26a0\ufe0f  Engine config PATCH warning ({err.get('code')}): {err.get('message')}")
-        return
 
     op_name = resp.get("name", "")
     if "/operations/" in op_name and not resp.get("done"):
-        print(f"  Polling operation {op_name.split('/')[-1]}...")
+        print(f"  Polling operation...")
         op_url = (
             f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
         )
@@ -93,10 +106,21 @@ def ensure_serverless_engine_config() -> None:
             if op.get("done"):
                 break
         else:
-            print("  \u26a0\ufe0f  Operation timed out \u2014 continuing anyway.")
+            print("  \u26a0\ufe0f  Operation timed out.")
-            return
 
-    print("  \u2713 RAG Engine Config set to basic tier.")
+    # Verify actual state
+    cfg = get_engine_config()
+    db_cfg = cfg.get("ragManagedDbConfig", {})
+    print(f"  Current ragManagedDbConfig: {json.dumps(db_cfg)}")
+    if "spanner" in db_cfg and "basic" not in db_cfg:
+        raise RuntimeError(
+            "Engine config is still in Spanner mode. "
+            "Cannot create a serverless corpus. "
+            "Run: curl -s -X PATCH -H 'Authorization: Bearer $(gcloud auth print-access-token)' "
+            f"-H 'Content-Type: application/json' {_engine_config_url()} "
+            "-d '{\"ragManagedDbConfig\":{\"basic\":{}}}'"
+        )
+    print("  \u2713 RAG Engine Config verified as basic tier.")
 
 
 def create_corpus_rest() -> str:
@@ -104,10 +128,10 @@ def create_corpus_rest() -> str:
     Create corpus via REST API directly, bypassing SDK backend_config bug.
     Returns the corpus resource name.
 
-    Uses ragManagedDbConfig.basic explicitly in the corpus body.
+    Per v1beta1 docs, RagCorpus.backend_config is a union:
-    The v1beta1 RagCorpus resource accepts ragManagedDbConfig directly
+      vectorDbConfig.ragManagedDb: {} -> serverless managed DB (no tier field)
-    (not vectorDbConfig). Sending {"basic":{}} forces serverless mode
+    This explicitly requests serverless at corpus level regardless of
-    regardless of the project-level engine config default.
+    the project-level engine config default.
     """
     url = (
         f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
@@ -120,11 +144,11 @@ def create_corpus_rest() -> str:
                 "model": "publishers/google/models/text-embedding-004"
             }
         },
-        # Explicitly request serverless (basic) mode at the corpus level.
+        # vectorDbConfig.ragManagedDb = serverless RAG Managed DB.
-        # The v1beta1 API uses ragManagedDbConfig on the RagCorpus body,
+        # RagManagedDb has no fields (no tier, no tier enum).
-        # NOT vectorDbConfig. This overrides any Spanner project default.
+        # This is the correct v1beta1 field name — not 'ragManagedDbConfig'.
-        "ragManagedDbConfig": {
+        "vectorDbConfig": {
-            "basic": {}
+            "ragManagedDb": {}
         }
     })
     token = get_token()
@@ -135,14 +159,14 @@ def create_corpus_rest() -> str:
          url, "-d", payload],
         capture_output=True, text=True,
     )
+    print(f"  API response: {result.stdout[:300]}")
     resp = json.loads(result.stdout)
     if "error" in resp:
         raise RuntimeError(f"Failed to create corpus: {resp['error']}")
 
-    # REST returns a long-running operation
     op_name = resp.get("name", "")
     if "/operations/" not in op_name:
-        raise RuntimeError(f"Unexpected response: {resp}")
+        raise RuntimeError(f"Unexpected response (no operation): {resp}")
 
     print(f"  Polling corpus creation operation...")
     op_url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
@@ -174,7 +198,6 @@ def get_or_create_corpus() -> RagCorpus:
     corpus_name = create_corpus_rest()
     print(f"\u2713 Corpus created: {corpus_name}")
 
-    # Re-fetch via SDK so we get a proper RagCorpus object
     for c in rag.list_corpora():
         if c.name == corpus_name:
             return c