diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index 8fd5bb4..3035dcf 100644 --- a/agents/rag/setup_corpus.py +++ b/agents/rag/setup_corpus.py @@ -10,9 +10,10 @@ REST API directly for corpus creation, then use the SDK for everything else. 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 payload omits vectorDbConfig entirely. Sending - vectorDbConfig.ragManagedDb forces Spanner mode which is allowlist-restricted - for new projects. Omitting it honours the project-level basic (serverless) config. +- 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. """ import os @@ -46,15 +47,18 @@ def get_token() -> str: def ensure_serverless_engine_config() -> None: """ Set project-level RAG Engine Config to basic (serverless) tier. - Always targets us-central1 — the ragEngineConfig endpoint is global/control-plane - and does not exist in other regions (patching europe-west1 etc. returns 'Invalid - endpoint name'). + 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 = ( 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": {}}}) token = get_token() @@ -69,7 +73,7 @@ def ensure_serverless_engine_config() -> None: if "error" in resp: err = resp["error"] - print(f" ⚠️ Engine config warning ({err.get('code')}): {err.get('message')} — continuing.") + print(f" \u26a0\ufe0f Engine config warning ({err.get('code')}): {err.get('message')} \u2014 continuing.") return op_name = resp.get("name", "") @@ -89,10 +93,10 @@ def ensure_serverless_engine_config() -> None: if op.get("done"): break else: - print(" ⚠️ Operation timed out — continuing anyway.") + print(" \u26a0\ufe0f Operation timed out \u2014 continuing anyway.") return - print(" ✓ RAG Engine Config set to basic tier.") + print(" \u2713 RAG Engine Config set to basic tier.") def create_corpus_rest() -> str: @@ -100,10 +104,10 @@ def create_corpus_rest() -> str: Create corpus via REST API directly, bypassing SDK backend_config bug. Returns the corpus resource name. - IMPORTANT: vectorDbConfig is intentionally omitted from the payload. - Sending vectorDbConfig.ragManagedDb explicitly requests Spanner mode, - which is allowlist-restricted for new projects. Omitting it causes the - API to honour the project-level ragEngineConfig (basic/serverless). + 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. """ url = ( f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1" @@ -115,8 +119,13 @@ def create_corpus_rest() -> str: "vertexPredictionEndpoint": { "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 intentionally omitted — see docstring above. }) token = get_token() result = subprocess.run( @@ -163,7 +172,7 @@ def get_or_create_corpus() -> RagCorpus: print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...") corpus_name = create_corpus_rest() - print(f"✓ Corpus created: {corpus_name}") + print(f"\u2713 Corpus created: {corpus_name}") # Re-fetch via SDK so we get a proper RagCorpus object for c in rag.list_corpora(): @@ -185,7 +194,7 @@ def import_documents(corpus: RagCorpus) -> None: ) ), ) - print("✓ Import complete.") + print("\u2713 Import complete.") def test_retrieval(corpus: RagCorpus) -> None: @@ -195,7 +204,7 @@ def test_retrieval(corpus: RagCorpus) -> None: text="test query", rag_retrieval_config=rag.RagRetrievalConfig(top_k=3), ) - print(f"✓ Test retrieval returned {len(response.contexts.contexts)} chunk(s).") + print(f"\u2713 Test retrieval returned {len(response.contexts.contexts)} chunk(s).") def main() -> None: