fix: create corpus via REST API directly, bypassing SDK backend_config bug

SDK 1.153.1 always defaults to Spanner when backend_config is omitted,
and crashes when it's provided. Use REST POST directly with
vectorDbConfig.ragManagedDb set to force serverless.

agents/rag/setup_corpus.py

@@ -3,13 +3,14 @@
 setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
 Project: propane-will-491900-m5
 
-Serverless mode is controlled via the project-level ragEngineConfig PATCH
-(ensure_serverless_engine_config). Passing backend_config to create_corpus
-triggers a SDK 1.153.1 bug and must be omitted.
+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
+REST API directly for corpus creation, then use the SDK for everything else.
 """
 
 import os
 import json
+import time
 import subprocess
 
 import vertexai
@@ -25,11 +26,16 @@ GCS_SOURCE = os.environ.get(
 )
 
 
+def get_token() -> str:
+    return subprocess.check_output(
+        ["gcloud", "auth", "print-access-token"], text=True
+    ).strip()
+
+
 def ensure_serverless_engine_config() -> None:
     """
-    Pre-flight: set project-level RAG Engine Config to basic (serverless) tier.
-    This is what actually controls the backend — backend_config in create_corpus
-    is intentionally omitted due to SDK bug in 1.153.1.
+    Set project-level RAG Engine Config to basic (serverless) tier.
+    Polls the operation until done.
     """
     print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
     endpoint = (
@@ -37,51 +43,116 @@ def ensure_serverless_engine_config() -> None:
         f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragEngineConfig"
     )
     payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
-
-    try:
-        token = subprocess.check_output(
-            ["gcloud", "auth", "print-access-token"], text=True
-        ).strip()
-    except subprocess.CalledProcessError:
-        print("  ⚠️  Could not get gcloud token — skipping pre-flight.")
-        return
+    token = get_token()
 
     result = subprocess.run(
-        [
-            "curl", "-s", "-X", "PATCH",
+        ["curl", "-s", "-X", "PATCH",
          "-H", f"Authorization: Bearer {token}",
          "-H", "Content-Type: application/json",
-            endpoint, "-d", payload,
-        ],
+         endpoint, "-d", payload],
         capture_output=True, text=True,
     )
+    resp = json.loads(result.stdout)
 
-    resp = result.stdout
-    if '"error"' in resp:
-        err = json.loads(resp).get("error", {})
-        code, msg = err.get("code"), err.get("message", "")
-        if code == 400 and "already" in msg.lower():
-            print(f"  ✓ Already configured: {msg}")
-        else:
-            print(f"  ⚠️  Engine config warning ({code}): {msg} — continuing anyway.")
+    if "error" in resp:
+        err = resp["error"]
+        print(f"  ⚠️  Engine config warning ({err.get('code')}): {err.get('message')} — continuing.")
+        return
+
+    op_name = resp.get("name", "")
+    if "/operations/" in op_name and not resp.get("done"):
+        print(f"  Polling operation {op_name.split('/')[-1]}...")
+        op_url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
+        for _ in range(20):
+            time.sleep(3)
+            token = get_token()
+            r = subprocess.run(
+                ["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
+                capture_output=True, text=True,
+            )
+            op = json.loads(r.stdout)
+            if op.get("done"):
+                break
         else:
+            print("  ⚠️  Operation timed out — continuing anyway.")
+            return
+
     print("  ✓ RAG Engine Config set to basic tier.")
 
 
+def create_corpus_rest() -> str:
+    """
+    Create corpus via REST API directly, bypassing SDK backend_config bug.
+    Returns the corpus resource name.
+    """
+    url = (
+        f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
+        f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora"
+    )
+    payload = json.dumps({
+        "displayName": CORPUS_DISPLAY_NAME,
+        "ragEmbeddingModelConfig": {
+            "vertexPredictionEndpoint": {
+                "model": "publishers/google/models/text-embedding-004"
+            }
+        },
+        "vectorDbConfig": {
+            "ragManagedDb": {}
+        }
+    })
+    token = get_token()
+    result = subprocess.run(
+        ["curl", "-s", "-X", "POST",
+         "-H", f"Authorization: Bearer {token}",
+         "-H", "Content-Type: application/json",
+         url, "-d", payload],
+        capture_output=True, text=True,
+    )
+    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}")
+
+    print(f"  Polling corpus creation operation...")
+    op_url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
+    for _ in range(40):
+        time.sleep(5)
+        token = get_token()
+        r = subprocess.run(
+            ["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
+            capture_output=True, text=True,
+        )
+        op = json.loads(r.stdout)
+        if op.get("done"):
+            if "error" in op:
+                raise RuntimeError(f"Corpus creation failed: {op['error']}")
+            corpus_name = op["response"]["name"]
+            return corpus_name
+
+    raise RuntimeError("Corpus creation operation timed out.")
+
+
 def get_or_create_corpus() -> RagCorpus:
-    """Return existing corpus by display name, or create a new one."""
+    """Return existing corpus by display name, or create a new serverless one."""
     for c in rag.list_corpora():
         if c.display_name == CORPUS_DISPLAY_NAME:
             print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists: {c.name}")
             return c
 
-    print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION}...")
-    # NOTE: backend_config is intentionally omitted.
-    # SDK 1.153.1 has a bug where passing any backend_config crashes.
-    # Serverless mode is already guaranteed by ensure_serverless_engine_config().
-    corpus = rag.create_corpus(display_name=CORPUS_DISPLAY_NAME)
-    print(f"✓ Corpus created: {corpus.name}")
-    return corpus
+    print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
+    corpus_name = create_corpus_rest()
+    print(f"✓ 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
+
+    raise RuntimeError(f"Corpus created but not found in list: {corpus_name}")
 
 
 def import_documents(corpus: RagCorpus) -> None: