136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
|
|
Project: propane-will-491900-m5
|
|
|
|
Explicitly uses Serverless (RagManagedDb) to avoid Spanner allowlist
|
|
restrictions on new projects in us-central1.
|
|
|
|
Tested against vertexai SDK 1.153.1:
|
|
RagVectorDbConfig(vector_db=RagManagedDb())
|
|
RagManagedDb is a no-arg marker class — selecting it is sufficient for serverless.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import subprocess
|
|
|
|
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/",
|
|
)
|
|
|
|
|
|
def ensure_serverless_engine_config() -> None:
|
|
"""
|
|
Pre-flight: set project-level RAG Engine Config to basic (serverless) tier.
|
|
Idempotent — safe to re-run.
|
|
"""
|
|
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"
|
|
)
|
|
payload = json.dumps({"ragManagedDbConfig": {"basic": {}}})
|
|
|
|
try:
|
|
token = subprocess.check_output(
|
|
["gcloud", "auth", "print-access-token"], text=True
|
|
).strip()
|
|
except subprocess.CalledProcessError:
|
|
print(" ⚠️ Could not get gcloud token — skipping pre-flight.")
|
|
return
|
|
|
|
result = subprocess.run(
|
|
[
|
|
"curl", "-s", "-X", "PATCH",
|
|
"-H", f"Authorization: Bearer {token}",
|
|
"-H", "Content-Type: application/json",
|
|
endpoint, "-d", payload,
|
|
],
|
|
capture_output=True, text=True,
|
|
)
|
|
|
|
resp = result.stdout
|
|
if '"error"' in resp:
|
|
err = json.loads(resp).get("error", {})
|
|
code, msg = err.get("code"), err.get("message", "")
|
|
if code == 400 and "already" in msg.lower():
|
|
print(f" ✓ Already configured: {msg}")
|
|
else:
|
|
print(f" ⚠️ Engine config warning ({code}): {msg} — continuing anyway.")
|
|
else:
|
|
print(" ✓ RAG Engine Config set to basic tier.")
|
|
|
|
|
|
def get_or_create_corpus() -> RagCorpus:
|
|
"""Return existing corpus by display name, or create a new serverless one."""
|
|
for c in rag.list_corpora():
|
|
if c.display_name == CORPUS_DISPLAY_NAME:
|
|
print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists: {c.name}")
|
|
return c
|
|
|
|
print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (serverless)...")
|
|
corpus = rag.create_corpus(
|
|
display_name=CORPUS_DISPLAY_NAME,
|
|
# RagManagedDb() is a no-arg marker class in SDK 1.153.1
|
|
# Selecting it routes to serverless managed DB (basic tier)
|
|
backend_config=rag.RagVectorDbConfig(
|
|
vector_db=rag.RagManagedDb()
|
|
),
|
|
)
|
|
print(f"✓ Corpus created: {corpus.name}")
|
|
return corpus
|
|
|
|
|
|
def import_documents(corpus: RagCorpus) -> None:
|
|
print(f"Importing files from {GCS_SOURCE}...")
|
|
rag.import_files(
|
|
corpus.name,
|
|
paths=[GCS_SOURCE],
|
|
transformation_config=rag.TransformationConfig(
|
|
chunking_config=rag.ChunkingConfig(
|
|
chunk_size=512,
|
|
chunk_overlap=100,
|
|
)
|
|
),
|
|
)
|
|
print("✓ Import complete.")
|
|
|
|
|
|
def test_retrieval(corpus: RagCorpus) -> None:
|
|
print("Running test retrieval query...")
|
|
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"✓ Test retrieval returned {len(response.contexts.contexts)} chunk(s).")
|
|
|
|
|
|
def main() -> None:
|
|
vertexai.init(project=PROJECT_ID, location=LOCATION)
|
|
|
|
ensure_serverless_engine_config()
|
|
|
|
corpus = get_or_create_corpus()
|
|
|
|
import_documents(corpus)
|
|
|
|
print(f"\nRAG_CORPUS={corpus.name}")
|
|
print("Add this to Secret Manager:")
|
|
print(f" gcloud secrets create rag-corpus-name --data-file=- <<<'{corpus.name}'")
|
|
|
|
test_retrieval(corpus)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|