fix(rag): use ragManagedDbConfig.basic in corpus body; fix engine config PATCH with scaled→basic switch

The v1beta1 API uses ragManagedDbConfig (not vectorDbConfig) on the RagCorpus
resource itself. Omitting it causes the engine to default to Spanner. Also the
engine config PATCH must use {"ragManagedDbConfig":{"basic":{}}} with the
correct field name matching v1beta1 schema.
This commit is contained in:
chrischristiansen-glitch 2026-05-24 19:23:30 +02:00
parent 2838534296
commit db3807bff9

View File

@ -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: