143 lines
5.1 KiB
Bash
143 lines
5.1 KiB
Bash
#!/usr/bin/env bash
|
|
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
|
|
# Serverless = RagManagedDb (no Spanner, no allowlist needed)
|
|
# Idempotent — safe to run multiple times
|
|
# Source .env before running: source .env
|
|
|
|
set -euo pipefail
|
|
|
|
: "${PROJECT_ID:?Set PROJECT_ID}"
|
|
: "${REGION:?Set REGION}"
|
|
: "${RAG_CORPUS_DISPLAY_NAME:?Set RAG_CORPUS_DISPLAY_NAME}"
|
|
|
|
RAG_REGION="${RAG_REGION:-${REGION}}"
|
|
|
|
echo "=== 07: Setting up Vertex AI RAG Engine (Serverless mode) ==="
|
|
echo " Project : ${PROJECT_ID}"
|
|
echo " Region : ${RAG_REGION}"
|
|
echo " Corpus : ${RAG_CORPUS_DISPLAY_NAME}"
|
|
echo ""
|
|
|
|
bash "$(dirname "$0")/00-authcheck.sh"
|
|
|
|
gcloud services enable aiplatform.googleapis.com \
|
|
storage.googleapis.com \
|
|
--project="${PROJECT_ID}" --quiet
|
|
echo "✓ APIs enabled"
|
|
|
|
# ── GCS bucket (idempotent) ──────────────────────────────────────────────
|
|
CORPUS_BUCKET="${PROJECT_ID}-agent-corpus"
|
|
if ! gsutil ls -b "gs://${CORPUS_BUCKET}" &>/dev/null; then
|
|
gsutil mb -l "${REGION}" -b on "gs://${CORPUS_BUCKET}"
|
|
echo "✓ GCS corpus bucket created: gs://${CORPUS_BUCKET}"
|
|
else
|
|
echo "✓ GCS corpus bucket exists: gs://${CORPUS_BUCKET}"
|
|
fi
|
|
|
|
# ── Upload seed docs ──────────────────────────────────────────────────────────
|
|
SEED_DIR="$(dirname "$0")/../docs/corpus-seed"
|
|
if [[ -d "${SEED_DIR}" ]] && ls "${SEED_DIR}"/*.md &>/dev/null; then
|
|
gsutil -m cp "${SEED_DIR}"/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
|
|
echo "✓ Seed documents uploaded to gs://${CORPUS_BUCKET}/seed/"
|
|
else
|
|
echo " No seed documents in docs/corpus-seed/ — skipping"
|
|
fi
|
|
|
|
# ── Python: create corpus ───────────────────────────────────────────────────────
|
|
python3 - << 'PYEOF'
|
|
import os, sys
|
|
|
|
try:
|
|
import vertexai
|
|
from vertexai.preview import rag
|
|
except ImportError:
|
|
print("ERROR: google-cloud-aiplatform not installed.")
|
|
sys.exit(1)
|
|
|
|
PROJECT_ID = os.environ["PROJECT_ID"]
|
|
RAG_REGION = os.environ.get("RAG_REGION", os.environ["REGION"])
|
|
DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
|
|
CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"
|
|
|
|
vertexai.init(project=PROJECT_ID, location=RAG_REGION)
|
|
|
|
# Check if corpus already exists
|
|
corpus = None
|
|
try:
|
|
for c in rag.list_corpora():
|
|
if c.display_name == DISPLAY_NAME:
|
|
corpus = c
|
|
print(f"✓ RAG corpus already exists: {c.name}")
|
|
break
|
|
except Exception as e:
|
|
print(f"WARNING: Could not list corpora: {e}")
|
|
|
|
if corpus is None:
|
|
print(f" Creating corpus '{DISPLAY_NAME}' in {RAG_REGION}...")
|
|
embedding_config = rag.EmbeddingModelConfig(
|
|
publisher_model="publishers/google/models/text-embedding-005"
|
|
)
|
|
|
|
# Try Serverless mode (RagManagedDb) first, fall back to plain create
|
|
created = False
|
|
for attempt in ["serverless", "plain"]:
|
|
try:
|
|
if attempt == "serverless":
|
|
# SDK >= 1.87: RagVectorDbConfig with rag_managed_db
|
|
try:
|
|
vector_db = rag.RagVectorDbConfig(
|
|
rag_managed_db=rag.RagManagedDb()
|
|
)
|
|
corpus = rag.create_corpus(
|
|
display_name=DISPLAY_NAME,
|
|
embedding_model_config=embedding_config,
|
|
vector_db=vector_db,
|
|
)
|
|
except TypeError:
|
|
# Older SDK: RagManagedDb not a kwarg — skip to plain
|
|
raise
|
|
else:
|
|
# Plain create — lets Google pick default (Serverless on new projects)
|
|
corpus = rag.create_corpus(
|
|
display_name=DISPLAY_NAME,
|
|
embedding_model_config=embedding_config,
|
|
)
|
|
print(f"✓ RAG corpus created [{attempt}]: {corpus.name}")
|
|
created = True
|
|
break
|
|
except Exception as e:
|
|
if attempt == "plain":
|
|
print(f"ERROR: Could not create corpus: {e}")
|
|
sys.exit(1)
|
|
print(f" [{attempt}] failed: {e} — retrying with plain...")
|
|
|
|
# Import seed documents (non-fatal)
|
|
gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"
|
|
try:
|
|
rag.import_files(
|
|
corpus_name=corpus.name,
|
|
paths=[gcs_uri],
|
|
chunk_size=512,
|
|
chunk_overlap=50,
|
|
max_embedding_requests_per_min=900,
|
|
)
|
|
print(f"✓ Documents imported from {gcs_uri}")
|
|
except Exception as e:
|
|
print(f" WARNING: Document import skipped: {e}")
|
|
|
|
with open("/tmp/rag_corpus_name.txt", "w") as f:
|
|
f.write(corpus.name)
|
|
|
|
print(f"")
|
|
print(f" Corpus resource name : {corpus.name}")
|
|
print(f" Region : {RAG_REGION}")
|
|
print(f"")
|
|
print(f" ACTION REQUIRED — add to .env:")
|
|
print(f' export RAG_CORPUS_NAME="{corpus.name}"')
|
|
PYEOF
|
|
|
|
echo ""
|
|
echo "=== 07: RAG Engine setup COMPLETE ==="
|
|
echo " View: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
|
|
echo ""
|