fix: explicitly set serverless RAG backend + improve error handling

- Pass RagVectorDbConfig with RagManagedDbConfig basic tier explicitly
  to avoid Spanner allowlist errors on new projects
- Add pre-flight ensure_serverless_engine_config() via REST PATCH
- Cleaner error messages with actionable hints
- Idempotent: skips creation if corpus already exists

agents/rag/setup_corpus.py

@@ -3,40 +3,91 @@
 setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
 Project: propane-will-491900-m5
 
-Serverless mode is the default when no backend_config is specified.
-Region: us-central1
+Explicitly uses Serverless (RagManagedDbConfig basic tier) to avoid
+Spanner allowlist restrictions on new projects in us-central1.
 """
 
 import os
+import json
+import subprocess
+
 import vertexai
 from vertexai import rag
+from vertexai.rag import RagCorpus
 
 PROJECT_ID = "propane-will-491900-m5"
 LOCATION = os.environ.get("RAG_LOCATION", "us-central1")
-CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "oavauco-knowledge-base")
+CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "osvauco-knowledge-base")
 GCS_SOURCE = os.environ.get(
     "RAG_GCS_SOURCE",
-    f"gs://{PROJECT_ID}-agent-staging/rag-docs/"
+    f"gs://{PROJECT_ID}-agent-staging/rag-docs/",
 )
 
-def main():
-    vertexai.init(project=PROJECT_ID, location=LOCATION)
 
-    # Check if corpus already exists
-    existing = list(rag.list_corpora())
-    for c in existing:
+def ensure_serverless_engine_config() -> None:
+    """
+    Pre-flight: set project-level RAG Engine Config to basic (serverless) tier.
+    Idempotent — safe to re-run.
+    """
+    print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
+    endpoint = (
+        f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
+        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
+
+    result = subprocess.run(
+        [
+            "curl", "-s", "-X", "PATCH",
+            "-H", f"Authorization: Bearer {token}",
+            "-H", "Content-Type: application/json",
+            endpoint, "-d", payload,
+        ],
+        capture_output=True, text=True,
+    )
+
+    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.")
+    else:
+        print("  ✓ RAG Engine Config set to basic tier.")
+
+
+def get_or_create_corpus() -> RagCorpus:
+    """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}")
-            corpus = c
-            break
-    else:
-        print(f"Creating RAG corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (Serverless)...")
-        # No backend_config = Serverless mode (default)
+            return c
+
+    print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (serverless)...")
     corpus = rag.create_corpus(
         display_name=CORPUS_DISPLAY_NAME,
+        # Explicitly request serverless / basic managed DB tier
+        backend_config=rag.RagVectorDbConfig(
+            rag_managed_db=rag.RagManagedDbConfig(
+                tier=rag.RagManagedDbConfig.Tier.BASIC,
             )
-        print(f"Corpus created: {corpus.name}")
+        ),
+    )
+    print(f"✓ Corpus created: {corpus.name}")
+    return corpus
 
+
+def import_documents(corpus: RagCorpus) -> None:
     print(f"Importing files from {GCS_SOURCE}...")
     rag.import_files(
         corpus.name,
@@ -48,17 +99,33 @@ def main():
             )
         ),
     )
-    print("Import complete.")
-    print(f"\nRAG_CORPUS={corpus.name}")
-    print("Set this as an environment variable or Secret Manager entry.")
+    print("✓ Import complete.")
 
-    # Test retrieval
+
+def test_retrieval(corpus: RagCorpus) -> None:
+    print("Running test retrieval query...")
     response = rag.retrieval_query(
         rag_resources=[rag.RagResource(rag_corpus=corpus.name)],
         text="test query",
         rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
     )
-    print(f"Test retrieval returned {len(response.contexts.contexts)} chunks.")
+    print(f"✓ Test retrieval returned {len(response.contexts.contexts)} chunk(s).")
+
+
+def main() -> None:
+    vertexai.init(project=PROJECT_ID, location=LOCATION)
+
+    ensure_serverless_engine_config()
+
+    corpus = get_or_create_corpus()
+
+    import_documents(corpus)
+
+    print(f"\nRAG_CORPUS={corpus.name}")
+    print("Add this to Secret Manager:")
+    print(f"  gcloud secrets create rag-corpus-name --data-file=- <<<'{corpus.name}'")
+
+    test_retrieval(corpus)
 
 
 if __name__ == "__main__":