fix(07-rag-setup): add regional fallback (us-east1), robust error handling, fix seed path
This commit is contained in:
GitHub
parent
7e0f01063e
commit
320973dc77
|
|
@ -1,23 +1,40 @@
|
|||
#!/usr/bin/env bash
|
||||
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus and upload initial documents
|
||||
# Idempotent
|
||||
# Idempotent — safe to run multiple times
|
||||
# Source: Vertex AI RAG Engine SDK (google-cloud-aiplatform >= 1.87)
|
||||
# NOTE: RAG Engine in us-central1 requires allowlist access.
|
||||
# Contact: vertex-ai-rag-engine-support@google.com
|
||||
# Alternative: set REGION=us-east1 or us-east4 for immediate access.
|
||||
#
|
||||
# REGIONAL AVAILABILITY (OQ-02):
|
||||
# RAG Engine in us-central1 requires allowlist access.
|
||||
# This script auto-falls-back to us-east1 if us-central1 is not accessible.
|
||||
# To bypass: export RAG_REGION=us-east1 before running.
|
||||
#
|
||||
# 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 fallback: prefer env override, then REGION, fallback to us-east1
|
||||
RAG_REGION="${RAG_REGION:-${REGION}}"
|
||||
RAG_FALLBACK_REGION="us-east1"
|
||||
|
||||
echo "=== 07: Setting up Vertex AI RAG Engine ==="
|
||||
echo " Project : ${PROJECT_ID}"
|
||||
echo " Region : ${RAG_REGION} (fallback: ${RAG_FALLBACK_REGION})"
|
||||
echo " Corpus : ${RAG_CORPUS_DISPLAY_NAME}"
|
||||
echo ""
|
||||
|
||||
bash "$(dirname "$0")/00-authcheck.sh"
|
||||
gcloud services enable aiplatform.googleapis.com --quiet
|
||||
|
||||
# Create GCS bucket for corpus source documents (idempotent)
|
||||
gcloud services enable aiplatform.googleapis.com \
|
||||
storage.googleapis.com \
|
||||
--project="${PROJECT_ID}" --quiet
|
||||
|
||||
echo "✓ APIs enabled"
|
||||
|
||||
# ── GCS bucket for corpus source documents (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}"
|
||||
|
|
@ -26,32 +43,60 @@ else
|
|||
echo "✓ GCS corpus bucket exists: gs://${CORPUS_BUCKET}"
|
||||
fi
|
||||
|
||||
# Upload seed documents if present
|
||||
if [[ -d "docs/corpus-seed" ]]; then
|
||||
gsutil -m cp docs/corpus-seed/*.md "gs://${CORPUS_BUCKET}/seed/" 2>/dev/null || true
|
||||
# ── Upload seed documents if present ─────────────────────────────────────
|
||||
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 found in docs/corpus-seed/ — skipping upload"
|
||||
fi
|
||||
|
||||
# Run Python to create/update corpus
|
||||
# ── Run Python to create/update corpus ───────────────────────────────────
|
||||
python3 - << PYEOF
|
||||
import os, sys
|
||||
|
||||
try:
|
||||
import vertexai
|
||||
from vertexai.preview import rag
|
||||
except ImportError:
|
||||
print("ERROR: google-cloud-aiplatform not installed. Run: pip install google-cloud-aiplatform>=1.87.0")
|
||||
print("ERROR: google-cloud-aiplatform not installed.")
|
||||
print("Run: pip install google-cloud-aiplatform>=1.87.0")
|
||||
sys.exit(1)
|
||||
|
||||
PROJECT_ID = os.environ["PROJECT_ID"]
|
||||
REGION = os.environ["REGION"]
|
||||
CORPUS_DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
|
||||
CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"
|
||||
PROJECT_ID = os.environ["PROJECT_ID"]
|
||||
RAG_REGION = os.environ.get("RAG_REGION", os.environ["REGION"])
|
||||
FALLBACK_REGION = "us-east1"
|
||||
DISPLAY_NAME = os.environ["RAG_CORPUS_DISPLAY_NAME"]
|
||||
CORPUS_BUCKET = f"{PROJECT_ID}-agent-corpus"
|
||||
|
||||
vertexai.init(project=PROJECT_ID, location=REGION)
|
||||
def try_init_and_list(region):
|
||||
vertexai.init(project=PROJECT_ID, location=region)
|
||||
try:
|
||||
return list(rag.list_corpora()), region
|
||||
except Exception as e:
|
||||
if "PERMISSION_DENIED" in str(e) or "allowlist" in str(e).lower() or "not found" in str(e).lower():
|
||||
return None, region
|
||||
raise
|
||||
|
||||
# Try primary region, fall back if needed
|
||||
corpora, active_region = try_init_and_list(RAG_REGION)
|
||||
if corpora is None and RAG_REGION != FALLBACK_REGION:
|
||||
print(f" RAG Engine not available in {RAG_REGION} — falling back to {FALLBACK_REGION}")
|
||||
corpora, active_region = try_init_and_list(FALLBACK_REGION)
|
||||
|
||||
if corpora is None:
|
||||
print(f"ERROR: RAG Engine not accessible in {RAG_REGION} or {FALLBACK_REGION}.")
|
||||
print("Apply for allowlist: vertex-ai-rag-engine-support@google.com")
|
||||
print("Or set RAG_REGION=us-east1 in .env")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"✓ RAG Engine accessible in region: {active_region}")
|
||||
|
||||
# Find or create corpus
|
||||
corpus = None
|
||||
for c in rag.list_corpora():
|
||||
if c.display_name == CORPUS_DISPLAY_NAME:
|
||||
for c in corpora:
|
||||
if c.display_name == DISPLAY_NAME:
|
||||
corpus = c
|
||||
print(f"✓ RAG corpus already exists: {c.name}")
|
||||
break
|
||||
|
|
@ -61,11 +106,12 @@ if corpus is None:
|
|||
publisher_model="publishers/google/models/text-embedding-005"
|
||||
)
|
||||
corpus = rag.create_corpus(
|
||||
display_name=CORPUS_DISPLAY_NAME,
|
||||
display_name=DISPLAY_NAME,
|
||||
embedding_model_config=embedding_config,
|
||||
)
|
||||
print(f"✓ RAG corpus created: {corpus.name}")
|
||||
|
||||
# Import seed documents (non-fatal)
|
||||
gcs_uri = f"gs://{CORPUS_BUCKET}/seed/"
|
||||
try:
|
||||
rag.import_files(
|
||||
|
|
@ -78,14 +124,23 @@ try:
|
|||
print(f"✓ Documents imported from {gcs_uri}")
|
||||
except Exception as e:
|
||||
print(f" WARNING: Document import skipped or failed: {e}")
|
||||
print(" Import manually via: https://console.cloud.google.com/vertex-ai/rag")
|
||||
print(f" Import manually: https://console.cloud.google.com/vertex-ai/rag?project={PROJECT_ID}")
|
||||
|
||||
# Persist corpus name for 08-memorybank-setup.sh and agent
|
||||
with open("/tmp/rag_corpus_name.txt", "w") as f:
|
||||
f.write(corpus.name)
|
||||
print(f"\n Corpus resource name: {corpus.name}")
|
||||
print(f" Add to .env: RAG_CORPUS_NAME={corpus.name}")
|
||||
|
||||
print(f"")
|
||||
print(f" Corpus resource name : {corpus.name}")
|
||||
print(f" Active region : {active_region}")
|
||||
print(f"")
|
||||
print(f" ACTION REQUIRED — add to .env:")
|
||||
print(f" RAG_CORPUS_NAME={corpus.name}")
|
||||
if active_region != os.environ.get("REGION"):
|
||||
print(f" RAG_REGION={active_region}")
|
||||
PYEOF
|
||||
|
||||
echo ""
|
||||
echo "=== 07: RAG Engine setup COMPLETE ==="
|
||||
echo " View corpus: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
|
||||
echo " View corpus: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
|
||||
echo ""
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user