fix(rag): correct corpus payload to vectorDbConfig.ragManagedDb; verify engine config actually switched before creating
- RagCorpus REST body uses vectorDbConfig.ragManagedDb (not ragManagedDbConfig)
- RagManagedDb has no tier field - empty object {} = serverless
- Add GET+verify step after PATCH to confirm spanner key is gone before proceeding
- Print actual engine config state so failures are immediately visible
This commit is contained in:
parent
db3807bff9
commit
68ac4fde35
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user