#!/usr/bin/env python3 """ 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. We call the REST API directly for corpus creation. 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 import json import time 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", "osvauco-knowledge-base") GCS_SOURCE = os.environ.get( "RAG_GCS_SOURCE", f"gs://{PROJECT_ID}-agent-staging/rag-docs/", ) # ragEngineConfig is a project-level control-plane endpoint. # It only exists in us-central1 regardless of where the corpus lives. _ENGINE_CONFIG_LOCATION = "us-central1" def get_token() -> str: return subprocess.check_output( ["gcloud", "auth", "print-access-token"], text=True ).strip() def _engine_config_url() -> str: return ( f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1" f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig" ) 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", _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 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_url = ( f"https://{_ENGINE_CONFIG_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(" \u26a0\ufe0f Operation timed out.") # 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: """ Create corpus via REST API directly, bypassing SDK backend_config bug. Returns the corpus resource name. 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" 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 = 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() 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, ) 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']}") op_name = resp.get("name", "") if "/operations/" not in op_name: 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}" 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 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} via REST...") corpus_name = create_corpus_rest() print(f"\u2713 Corpus created: {corpus_name}") 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: print(f"Importing files from {GCS_SOURCE}...") rag.import_files( corpus.name, paths=[GCS_SOURCE], transformation_config=rag.TransformationConfig( chunking_config=rag.ChunkingConfig( chunk_size=512, chunk_overlap=100, ) ), ) print("\u2713 Import complete.") 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"\u2713 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__": main()