diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index 3035dcf..55bc6f0 100644 --- a/agents/rag/setup_corpus.py +++ b/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, -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. +SDK 1.153.1 has a bug where backend_config crashes when provided. +We call the REST API directly for corpus creation. -Key design decisions: -- ragEngineConfig PATCH always targets us-central1 — it is a project-level - control-plane endpoint that does NOT exist in other regions. -- Corpus creation uses ragManagedDbConfig.basic in the RagCorpus body. - The v1beta1 API uses ragManagedDbConfig (not vectorDbConfig) on the corpus - resource. Sending it explicitly with {"basic":{}} forces serverless mode - and overrides any project-level Spanner default. +v1beta1 REST schema facts (verified from Google docs): + RagCorpus.backend_config is a union field: + vectorDbConfig: RagVectorDbConfig + .ragManagedDb: RagManagedDb <- serverless managed DB, no tier field + .vertexVectorSearch, .pinecone, .weaviate <- alternatives + RagVectorDbConfig has NO 'ragManagedDbConfig' field. + 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: - """ - 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 = ( +def _engine_config_url() -> str: + return ( 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.") - return + print(f" \u26a0\ufe0f Engine config PATCH warning ({err.get('code')}): {err.get('message')}") 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.") - return + print(" \u26a0\ufe0f Operation timed out.") - 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. - The v1beta1 RagCorpus resource accepts ragManagedDbConfig directly - (not vectorDbConfig). Sending {"basic":{}} forces serverless mode - regardless of the project-level engine config default. + Per v1beta1 docs, RagCorpus.backend_config is a union: + vectorDbConfig.ragManagedDb: {} -> serverless managed DB (no tier field) + This explicitly requests serverless at corpus level regardless of + 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. - # The v1beta1 API uses ragManagedDbConfig on the RagCorpus body, - # NOT vectorDbConfig. This overrides any Spanner project default. - "ragManagedDbConfig": { - "basic": {} + # vectorDbConfig.ragManagedDb = serverless RAG Managed DB. + # RagManagedDb has no fields (no tier, no tier enum). + # This is the correct v1beta1 field name — not 'ragManagedDbConfig'. + "vectorDbConfig": { + "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