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:
parent
2838534296
commit
db3807bff9
|
|
@ -10,9 +10,10 @@ REST API directly for corpus creation, then use the SDK for everything else.
|
||||||
Key design decisions:
|
Key design decisions:
|
||||||
- ragEngineConfig PATCH always targets us-central1 — it is a project-level
|
- ragEngineConfig PATCH always targets us-central1 — it is a project-level
|
||||||
control-plane endpoint that does NOT exist in other regions.
|
control-plane endpoint that does NOT exist in other regions.
|
||||||
- Corpus creation payload omits vectorDbConfig entirely. Sending
|
- Corpus creation uses ragManagedDbConfig.basic in the RagCorpus body.
|
||||||
vectorDbConfig.ragManagedDb forces Spanner mode which is allowlist-restricted
|
The v1beta1 API uses ragManagedDbConfig (not vectorDbConfig) on the corpus
|
||||||
for new projects. Omitting it honours the project-level basic (serverless) config.
|
resource. Sending it explicitly with {"basic":{}} forces serverless mode
|
||||||
|
and overrides any project-level Spanner default.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
|
@ -46,15 +47,18 @@ def get_token() -> str:
|
||||||
def ensure_serverless_engine_config() -> None:
|
def ensure_serverless_engine_config() -> None:
|
||||||
"""
|
"""
|
||||||
Set project-level RAG Engine Config to basic (serverless) tier.
|
Set project-level RAG Engine Config to basic (serverless) tier.
|
||||||
Always targets us-central1 — the ragEngineConfig endpoint is global/control-plane
|
Always targets us-central1. The endpoint does not exist in other regions.
|
||||||
and does not exist in other regions (patching europe-west1 etc. returns 'Invalid
|
|
||||||
endpoint name').
|
The v1beta1 ragEngineConfig schema:
|
||||||
|
{ "ragManagedDbConfig": { "basic": {} } } <- serverless
|
||||||
|
{ "ragManagedDbConfig": { "scaled": {} } } <- Spanner / managed DB
|
||||||
"""
|
"""
|
||||||
print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
|
print("Ensuring RAG Engine Config is set to serverless (basic) tier...")
|
||||||
endpoint = (
|
endpoint = (
|
||||||
f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
|
f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
|
||||||
f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
|
f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
|
||||||
)
|
)
|
||||||
|
# Explicitly set basic and unset spanner by sending only basic.
|
||||||
payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
|
payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
|
||||||
token = get_token()
|
token = get_token()
|
||||||
|
|
||||||
|
|
@ -69,7 +73,7 @@ def ensure_serverless_engine_config() -> None:
|
||||||
|
|
||||||
if "error" in resp:
|
if "error" in resp:
|
||||||
err = resp["error"]
|
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
|
return
|
||||||
|
|
||||||
op_name = resp.get("name", "")
|
op_name = resp.get("name", "")
|
||||||
|
|
@ -89,10 +93,10 @@ def ensure_serverless_engine_config() -> None:
|
||||||
if op.get("done"):
|
if op.get("done"):
|
||||||
break
|
break
|
||||||
else:
|
else:
|
||||||
print(" ⚠️ Operation timed out — continuing anyway.")
|
print(" \u26a0\ufe0f Operation timed out \u2014 continuing anyway.")
|
||||||
return
|
return
|
||||||
|
|
||||||
print(" ✓ RAG Engine Config set to basic tier.")
|
print(" \u2713 RAG Engine Config set to basic tier.")
|
||||||
|
|
||||||
|
|
||||||
def create_corpus_rest() -> str:
|
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.
|
Create corpus via REST API directly, bypassing SDK backend_config bug.
|
||||||
Returns the corpus resource name.
|
Returns the corpus resource name.
|
||||||
|
|
||||||
IMPORTANT: vectorDbConfig is intentionally omitted from the payload.
|
Uses ragManagedDbConfig.basic explicitly in the corpus body.
|
||||||
Sending vectorDbConfig.ragManagedDb explicitly requests Spanner mode,
|
The v1beta1 RagCorpus resource accepts ragManagedDbConfig directly
|
||||||
which is allowlist-restricted for new projects. Omitting it causes the
|
(not vectorDbConfig). Sending {"basic":{}} forces serverless mode
|
||||||
API to honour the project-level ragEngineConfig (basic/serverless).
|
regardless of the project-level engine config default.
|
||||||
"""
|
"""
|
||||||
url = (
|
url = (
|
||||||
f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
|
f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1"
|
||||||
|
|
@ -115,8 +119,13 @@ def create_corpus_rest() -> str:
|
||||||
"vertexPredictionEndpoint": {
|
"vertexPredictionEndpoint": {
|
||||||
"model": "publishers/google/models/text-embedding-004"
|
"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()
|
token = get_token()
|
||||||
result = subprocess.run(
|
result = subprocess.run(
|
||||||
|
|
@ -163,7 +172,7 @@ def get_or_create_corpus() -> RagCorpus:
|
||||||
|
|
||||||
print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
|
print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
|
||||||
corpus_name = create_corpus_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
|
# Re-fetch via SDK so we get a proper RagCorpus object
|
||||||
for c in rag.list_corpora():
|
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:
|
def test_retrieval(corpus: RagCorpus) -> None:
|
||||||
|
|
@ -195,7 +204,7 @@ def test_retrieval(corpus: RagCorpus) -> None:
|
||||||
text="test query",
|
text="test query",
|
||||||
rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
|
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:
|
def main() -> None:
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue
Block a user