diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index 37ad29f..8fd5bb4 100644 --- a/agents/rag/setup_corpus.py +++ b/agents/rag/setup_corpus.py @@ -6,6 +6,13 @@ 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. + +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. """ import os @@ -25,6 +32,10 @@ GCS_SOURCE = os.environ.get( 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( @@ -35,12 +46,14 @@ def get_token() -> str: def ensure_serverless_engine_config() -> None: """ Set project-level RAG Engine Config to basic (serverless) tier. - Polls the operation until done. + 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'). """ print("Ensuring RAG Engine Config is set to serverless (basic) tier...") endpoint = ( - f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1" - f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragEngineConfig" + f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1" + f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig" ) payload = json.dumps({"ragManagedDbConfig": {"basic": {}}}) token = get_token() @@ -62,7 +75,9 @@ def ensure_serverless_engine_config() -> None: op_name = resp.get("name", "") if "/operations/" in op_name and not resp.get("done"): print(f" Polling operation {op_name.split('/')[-1]}...") - op_url = f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}" + op_url = ( + f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}" + ) for _ in range(20): time.sleep(3) token = get_token() @@ -84,6 +99,11 @@ 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). """ url = ( f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1" @@ -95,10 +115,8 @@ def create_corpus_rest() -> str: "vertexPredictionEndpoint": { "model": "publishers/google/models/text-embedding-004" } - }, - "vectorDbConfig": { - "ragManagedDb": {} } + # vectorDbConfig intentionally omitted — see docstring above. }) token = get_token() result = subprocess.run(