335 lines
11 KiB
Python
335 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
|
|
Project: propane-will-491900-m5
|
|
|
|
SDK 1.153.1 has bugs in backend_config and RagManagedDbConfig.
|
|
We use the REST API (v1beta1) for corpus creation and engine configuration.
|
|
|
|
Changes implemented:
|
|
- Graceful degradation: If RAG Engine is restricted (Spanner mode), skip and exit cleanly.
|
|
- ensure_serverless_engine_config() sets RAG_ENGINE_AVAILABLE flag.
|
|
- get_or_create_corpus() returns None if RAG is unavailable.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import time
|
|
import subprocess
|
|
import sys
|
|
|
|
import vertexai
|
|
from vertexai import rag
|
|
from vertexai.rag import RagCorpus
|
|
|
|
PROJECT_ID = "propane-will-491900-m5"
|
|
LOCATION = os.environ.get("RAG_LOCATION", "us-central1")
|
|
CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "osvauco-knowledge-base")
|
|
GCS_SOURCE = os.environ.get(
|
|
"RAG_GCS_SOURCE",
|
|
f"gs://{PROJECT_ID}-agent-staging/rag-docs/",
|
|
)
|
|
|
|
# Global status flag
|
|
RAG_ENGINE_AVAILABLE = True
|
|
|
|
# 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(
|
|
["gcloud", "auth", "print-access-token"], text=True
|
|
).strip()
|
|
|
|
|
|
def _engine_config_url() -> str:
|
|
return (
|
|
f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1"
|
|
f"/projects/{PROJECT_ID}/locations/{_ENGINE_CONFIG_LOCATION}/ragEngineConfig"
|
|
)
|
|
|
|
|
|
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,
|
|
)
|
|
try:
|
|
return json.loads(r.stdout)
|
|
except Exception:
|
|
print(f"Failed to parse engine config: {r.stdout}")
|
|
return {}
|
|
|
|
|
|
def ensure_serverless_engine_config() -> None:
|
|
"""
|
|
Set project-level RAG Engine Config to basic (serverless) tier.
|
|
Only applicable for us-central1, us-east1, us-east4.
|
|
"""
|
|
global RAG_ENGINE_AVAILABLE
|
|
|
|
restricted_regions = {"us-central1", "us-east1", "us-east4"}
|
|
if LOCATION not in restricted_regions:
|
|
print(f"Skipping engine config check for location: {LOCATION} (not restricted)")
|
|
return
|
|
|
|
print(f"Ensuring RAG Engine Config in {_ENGINE_CONFIG_LOCATION} is set to serverless (basic) tier...")
|
|
token = get_token()
|
|
|
|
# Use updateMask=ragManagedDbConfig to target the entire configuration block.
|
|
# Set spanner: null to explicitly clear it.
|
|
payload = {
|
|
"ragManagedDbConfig": {
|
|
"basic": {},
|
|
"spanner": None
|
|
}
|
|
}
|
|
|
|
url = f"{_engine_config_url()}"
|
|
|
|
result = subprocess.run(
|
|
["curl", "-s", "-X", "PATCH",
|
|
"-H", f"Authorization: Bearer {token}",
|
|
"-H", "Content-Type: application/json",
|
|
url,
|
|
"-d", json.dumps(payload)],
|
|
capture_output=True, text=True,
|
|
)
|
|
|
|
try:
|
|
resp = json.loads(result.stdout)
|
|
except Exception:
|
|
print(f" \u26a0\ufe0f Failed to parse PATCH response: {result.stdout}")
|
|
return
|
|
|
|
if "error" in resp:
|
|
err = resp["error"]
|
|
print(f" \u26a0\ufe0f Engine config PATCH error ({err.get('code')}): {err.get('message')}")
|
|
|
|
op_name = resp.get("name", "")
|
|
if "/operations/" in op_name and not resp.get("done"):
|
|
print(f" Polling engine config operation...")
|
|
op_url = (
|
|
f"https://{_ENGINE_CONFIG_LOCATION}-aiplatform.googleapis.com/v1beta1/{op_name}"
|
|
)
|
|
for _ in range(20):
|
|
time.sleep(3)
|
|
token = get_token()
|
|
r = subprocess.run(
|
|
["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
|
|
capture_output=True, text=True,
|
|
)
|
|
op = json.loads(r.stdout)
|
|
if op.get("done"):
|
|
break
|
|
else:
|
|
print(" \u26a0\ufe0f Engine config operation timed out.")
|
|
|
|
# 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:
|
|
print(f"\n\u26a0\ufe0f WARN: Engine config in {_ENGINE_CONFIG_LOCATION} is still in Spanner mode.")
|
|
print(" RAG Engine appears to be restricted for this project.")
|
|
RAG_ENGINE_AVAILABLE = False
|
|
return
|
|
|
|
print(" \u2713 RAG Engine Config verified as basic tier.")
|
|
|
|
|
|
def create_corpus_rest(loc: str) -> str:
|
|
"""
|
|
Create corpus via REST API directly.
|
|
"""
|
|
url = (
|
|
f"https://{loc}-aiplatform.googleapis.com/v1beta1"
|
|
f"/projects/{PROJECT_ID}/locations/{loc}/ragCorpora"
|
|
)
|
|
# vectorDbConfig.ragManagedDb = serverless RAG Managed DB.
|
|
# This is the correct v1beta1 field name.
|
|
payload = json.dumps({
|
|
"displayName": CORPUS_DISPLAY_NAME,
|
|
"ragEmbeddingModelConfig": {
|
|
"vertexPredictionEndpoint": {
|
|
"model": "publishers/google/models/text-embedding-004"
|
|
}
|
|
},
|
|
"vectorDbConfig": {
|
|
"ragManagedDb": {}
|
|
}
|
|
})
|
|
token = get_token()
|
|
print(f" Sending POST to {url}...")
|
|
result = subprocess.run(
|
|
["curl", "-s", "-X", "POST",
|
|
"-H", f"Authorization: Bearer {token}",
|
|
"-H", "Content-Type: application/json",
|
|
url, "-d", payload],
|
|
capture_output=True, text=True,
|
|
)
|
|
|
|
try:
|
|
resp = json.loads(result.stdout)
|
|
except Exception:
|
|
raise RuntimeError(f"Failed to parse API response: {result.stdout}")
|
|
|
|
if "error" in resp:
|
|
# Pass the error object up so get_or_create_corpus can inspect it
|
|
raise RuntimeError(json.dumps(resp["error"]))
|
|
|
|
op_name = resp.get("name", "")
|
|
if "/operations/" not in op_name:
|
|
raise RuntimeError(f"Unexpected response (no operation): {resp}")
|
|
|
|
print(f" Polling corpus creation operation {op_name}...")
|
|
op_url = f"https://{loc}-aiplatform.googleapis.com/v1beta1/{op_name}"
|
|
for _ in range(40):
|
|
time.sleep(5)
|
|
token = get_token()
|
|
r = subprocess.run(
|
|
["curl", "-s", "-H", f"Authorization: Bearer {token}", op_url],
|
|
capture_output=True, text=True,
|
|
)
|
|
op = json.loads(r.stdout)
|
|
if op.get("done"):
|
|
if "error" in op:
|
|
raise RuntimeError(json.dumps(op["error"]))
|
|
corpus_name = op["response"]["name"]
|
|
return corpus_name
|
|
|
|
raise RuntimeError("Corpus creation operation timed out.")
|
|
|
|
|
|
def get_or_create_corpus() -> RagCorpus:
|
|
"""Return existing corpus by display name, or create a new serverless one with fallback."""
|
|
global LOCATION, RAG_ENGINE_AVAILABLE
|
|
|
|
if not RAG_ENGINE_AVAILABLE:
|
|
print("\n\u2139 RAG Engine appears to be in restricted Spanner mode for this project. Skipping corpus creation.")
|
|
return None
|
|
|
|
def _find_in_list():
|
|
try:
|
|
for c in rag.list_corpora():
|
|
if c.display_name == CORPUS_DISPLAY_NAME:
|
|
return c
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
# 1. Try finding in current LOCATION
|
|
vertexai.init(project=PROJECT_ID, location=LOCATION)
|
|
corpus = _find_in_list()
|
|
if corpus:
|
|
print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists in {LOCATION}: {corpus.name}")
|
|
return corpus
|
|
|
|
# 2. Try creating in current LOCATION
|
|
print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} via REST...")
|
|
try:
|
|
corpus_name = create_corpus_rest(LOCATION)
|
|
print(f"\u2713 Corpus created in {LOCATION}: {corpus_name}")
|
|
except RuntimeError as e:
|
|
err_str = str(e)
|
|
# Check for Spanner restriction error message
|
|
if "using Spanner mode with RAG Engine in us-central1, us-east1, and us-east4 is restricted" in err_str:
|
|
if LOCATION != "europe-west4":
|
|
print(f"\u26a0\ufe0f Spanner restriction detected in {LOCATION}. Retrying in europe-west4...")
|
|
LOCATION = "europe-west4"
|
|
vertexai.init(project=PROJECT_ID, location=LOCATION)
|
|
|
|
# Check if it already exists in the fallback location
|
|
corpus = _find_in_list()
|
|
if corpus:
|
|
print(f"Corpus '{CORPUS_DISPLAY_NAME}' found in fallback {LOCATION}: {corpus.name}")
|
|
return corpus
|
|
|
|
try:
|
|
corpus_name = create_corpus_rest(LOCATION)
|
|
print(f"\u2713 Corpus created in fallback {LOCATION}: {corpus_name}")
|
|
except RuntimeError as e2:
|
|
if "using Spanner mode" in str(e2):
|
|
print("\n\u26a0\ufe0f RAG Engine restriction confirmed in fallback region. RAG is effectively unavailable.")
|
|
RAG_ENGINE_AVAILABLE = False
|
|
return None
|
|
raise
|
|
else:
|
|
print("\n\u26a0\ufe0f RAG Engine restriction confirmed in fallback region. RAG is effectively unavailable.")
|
|
RAG_ENGINE_AVAILABLE = False
|
|
return None
|
|
else:
|
|
raise RuntimeError(f"Failed to create corpus in {LOCATION}: {err_str}")
|
|
|
|
# Final retrieval of the corpus object
|
|
for c in rag.list_corpora():
|
|
if c.name == corpus_name:
|
|
return c
|
|
|
|
raise RuntimeError(f"Corpus created but not found in list: {corpus_name}")
|
|
|
|
|
|
def import_documents(corpus: RagCorpus) -> None:
|
|
print(f"Importing files from {GCS_SOURCE} into {corpus.name}...")
|
|
rag.import_files(
|
|
corpus.name,
|
|
paths=[GCS_SOURCE],
|
|
transformation_config=rag.TransformationConfig(
|
|
chunking_config=rag.ChunkingConfig(
|
|
chunk_size=512,
|
|
chunk_overlap=100,
|
|
)
|
|
),
|
|
)
|
|
print("\u2713 Import job submitted.")
|
|
|
|
|
|
def test_retrieval(corpus: RagCorpus) -> None:
|
|
print("Running test retrieval query...")
|
|
try:
|
|
response = rag.retrieval_query(
|
|
rag_resources=[rag.RagResource(rag_corpus=corpus.name)],
|
|
text="test query",
|
|
rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
|
|
)
|
|
print(f"\u2713 Test retrieval returned {len(response.contexts.contexts)} chunk(s).")
|
|
except Exception as e:
|
|
print(f" \u26a0\ufe0f Retrieval test failed (import might still be processing): {e}")
|
|
|
|
|
|
def main() -> None:
|
|
print(f"--- RAG Setup Starting (Target Location: {LOCATION}) ---")
|
|
|
|
ensure_serverless_engine_config()
|
|
|
|
corpus = get_or_create_corpus()
|
|
|
|
if corpus is None:
|
|
print("\n--- RAG AVAILABILITY SUMMARY ---")
|
|
print(f"Project : {PROJECT_ID}")
|
|
print("Status : UNAVAILABLE (Platform restriction: Spanner Mode only)")
|
|
print("Action : Skipping RAG integration. Downstream tools will operate without a corpus.")
|
|
print("---------------------------------")
|
|
sys.exit(0)
|
|
|
|
import_documents(corpus)
|
|
|
|
print(f"\nSUCCESS")
|
|
print(f"RAG_LOCATION={LOCATION}")
|
|
print(f"RAG_CORPUS={corpus.name}")
|
|
print("\nNext steps (set in your environment):")
|
|
print(f"export RAG_LOCATION={LOCATION}")
|
|
print(f"export RAG_CORPUS_NAME={corpus.name}")
|
|
|
|
print("\nWaiting 10s before test retrieval...")
|
|
time.sleep(10)
|
|
test_retrieval(corpus)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|