From a2790be421b92bfdc280f5c290e718a89eb54989 Mon Sep 17 00:00:00 2001 From: chrischristiansen-glitch Date: Sun, 24 May 2026 19:11:25 +0200 Subject: [PATCH] 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 | 145 +++++++++++++++++++++++++++---------- 1 file changed, 108 insertions(+), 37 deletions(-) diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index d64f130..37ad29f 100644 --- a/agents/rag/setup_corpus.py +++ b/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", - "-H", f"Authorization: Bearer {token}", - "-H", "Content-Type: application/json", - endpoint, "-d", payload, - ], + ["curl", "-s", "-X", "PATCH", + "-H", f"Authorization: Bearer {token}", + "-H", "Content-Type: application/json", + 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}") + 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(f" ⚠️ Engine config warning ({code}): {msg} — continuing anyway.") - else: - print(" ✓ RAG Engine Config set to basic tier.") + 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: