fix: 07-rag-setup.sh — REST API corpus creation, RAG_REGION=europe-west4 default (no SDK needed)
This commit is contained in:
parent
eea7e8497f
commit
9928c595b8
|
|
@ -1,6 +1,7 @@
|
|||
#!/usr/bin/env bash
|
||||
# 07-rag-setup.sh — Create Vertex AI RAG Engine corpus in Serverless mode
|
||||
# Serverless = RagManagedDb (no Spanner, no allowlist needed)
|
||||
# Uses REST API directly — no SDK version dependency
|
||||
# RAG_REGION defaults to europe-west4 (Serverless available, no allowlist)
|
||||
# Idempotent — safe to run multiple times
|
||||
# Source .env before running: source .env
|
||||
|
||||
|
|
@ -10,11 +11,12 @@ set -euo pipefail
|
|||
: "${REGION:?Set REGION}"
|
||||
: "${RAG_CORPUS_DISPLAY_NAME:?Set RAG_CORPUS_DISPLAY_NAME}"
|
||||
|
||||
RAG_REGION="${RAG_REGION:-${REGION}}"
|
||||
# europe-west4: Serverless RAG available, no allowlist, gemini-2.0-flash + 2.5-pro present
|
||||
RAG_REGION="${RAG_REGION:-europe-west4}"
|
||||
|
||||
echo "=== 07: Setting up Vertex AI RAG Engine (Serverless mode) ==="
|
||||
echo " Project : ${PROJECT_ID}"
|
||||
echo " Region : ${RAG_REGION}"
|
||||
echo " Region : ${RAG_REGION} (override with RAG_REGION= if needed)"
|
||||
echo " Corpus : ${RAG_CORPUS_DISPLAY_NAME}"
|
||||
echo ""
|
||||
|
||||
|
|
@ -25,11 +27,6 @@ gcloud services enable aiplatform.googleapis.com \
|
|||
--project="${PROJECT_ID}" --quiet
|
||||
echo "✓ APIs enabled"
|
||||
|
||||
# ── Upgrade SDK to ensure RagManagedDb support ───────────────────────────────
|
||||
echo " Upgrading google-cloud-aiplatform SDK..."
|
||||
pip install --quiet --upgrade google-cloud-aiplatform
|
||||
echo "✓ SDK upgraded"
|
||||
|
||||
# ── GCS bucket (idempotent) ──────────────────────────────────────────────
|
||||
CORPUS_BUCKET="${PROJECT_ID}-agent-corpus"
|
||||
if ! gsutil ls -b "gs://${CORPUS_BUCKET}" &>/dev/null; then
|
||||
|
|
@ -48,136 +45,114 @@ else
|
|||
echo " No seed documents in docs/corpus-seed/ — skipping"
|
||||
fi
|
||||
|
||||
# ── Python: create corpus ────────────────────────────────────────────────────
|
||||
python3 - << 'PYEOF'
|
||||
import os, sys
|
||||
# ── REST: check if corpus exists ───────────────────────────────────────────
|
||||
TOKEN=$(gcloud auth print-access-token)
|
||||
API="https://${RAG_REGION}-aiplatform.googleapis.com/v1beta1"
|
||||
PARENT="projects/${PROJECT_ID}/locations/${RAG_REGION}"
|
||||
|
||||
try:
|
||||
import vertexai
|
||||
from vertexai.preview import rag
|
||||
except ImportError:
|
||||
print("ERROR: google-cloud-aiplatform not installed.")
|
||||
sys.exit(1)
|
||||
echo " Checking for existing corpus..."
|
||||
EXISTING=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
|
||||
"${API}/${PARENT}/ragCorpora" 2>/dev/null || echo '{}')
|
||||
|
||||
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}")
|
||||
CORPUS_NAME=$(echo "${EXISTING}" | python3 -c "
|
||||
import sys, json
|
||||
data = json.load(sys.stdin)
|
||||
for c in data.get('ragCorpora', []):
|
||||
if c.get('displayName') == '${RAG_CORPUS_DISPLAY_NAME}':
|
||||
print(c['name'])
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"WARNING: Could not list corpora: {e}")
|
||||
" 2>/dev/null || true)
|
||||
|
||||
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"
|
||||
)
|
||||
if [[ -n "${CORPUS_NAME}" ]]; then
|
||||
echo "✓ RAG corpus already exists: ${CORPUS_NAME}"
|
||||
else
|
||||
echo " Creating corpus '${RAG_CORPUS_DISPLAY_NAME}' in ${RAG_REGION} (Serverless/REST)..."
|
||||
|
||||
# Attempt order:
|
||||
# 1. SDK >= 1.87 : RagVectorDbConfig(rag_managed_db=RagManagedDb())
|
||||
# 2. Older SDK : RagVectorDbConfig(rag_managed_db=RagManagedDbConfig())
|
||||
# 3. REST fallback: gapic-style with vector_db proto dict
|
||||
attempts = ["new_sdk", "old_sdk", "proto_dict"]
|
||||
for attempt in attempts:
|
||||
try:
|
||||
if attempt == "new_sdk":
|
||||
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,
|
||||
)
|
||||
RESPONSE=$(curl -sf -X POST \
|
||||
-H "Authorization: Bearer ${TOKEN}" \
|
||||
-H "Content-Type: application/json" \
|
||||
"${API}/${PARENT}/ragCorpora" \
|
||||
-d '{
|
||||
"displayName": "'"${RAG_CORPUS_DISPLAY_NAME}"'",
|
||||
"ragEmbeddingModelConfig": {
|
||||
"vertexPredictionEndpoint": {
|
||||
"publisherModel": "publishers/google/models/text-embedding-005"
|
||||
}
|
||||
},
|
||||
"ragVectorDbConfig": {
|
||||
"ragManagedDb": {}
|
||||
}
|
||||
}')
|
||||
|
||||
elif attempt == "old_sdk":
|
||||
# Some SDK versions use RagManagedDbConfig instead
|
||||
RagManagedDbConfig = getattr(rag, "RagManagedDbConfig", None)
|
||||
if RagManagedDbConfig is None:
|
||||
raise AttributeError("RagManagedDbConfig not in this SDK version")
|
||||
vector_db = rag.RagVectorDbConfig(
|
||||
rag_managed_db=RagManagedDbConfig()
|
||||
)
|
||||
corpus = rag.create_corpus(
|
||||
display_name=DISPLAY_NAME,
|
||||
embedding_model_config=embedding_config,
|
||||
vector_db=vector_db,
|
||||
)
|
||||
# REST create returns an LRO — poll until done
|
||||
OPERATION=$(echo "${RESPONSE}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('name',''))" 2>/dev/null || true)
|
||||
|
||||
elif attempt == "proto_dict":
|
||||
# Last resort: pass vector_db as a plain dict understood by gapic
|
||||
from google.cloud.aiplatform_v1beta1.types import (
|
||||
RagCorpus, RagVectorDbConfig, RagManagedDb,
|
||||
RagEmbeddingModelConfig
|
||||
)
|
||||
from google.cloud import aiplatform_v1beta1 as aip
|
||||
if [[ -z "${OPERATION}" ]]; then
|
||||
echo "ERROR: No operation returned. Response:"
|
||||
echo "${RESPONSE}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
client = aip.VertexRagDataServiceClient(
|
||||
client_options={"api_endpoint": f"{RAG_REGION}-aiplatform.googleapis.com"}
|
||||
)
|
||||
parent = f"projects/{PROJECT_ID}/locations/{RAG_REGION}"
|
||||
rag_corpus = RagCorpus(
|
||||
display_name=DISPLAY_NAME,
|
||||
rag_embedding_model_config=RagEmbeddingModelConfig(
|
||||
vertex_prediction_endpoint=RagEmbeddingModelConfig.VertexPredictionEndpoint(
|
||||
publisher_model="publishers/google/models/text-embedding-005"
|
||||
)
|
||||
),
|
||||
rag_vector_db_config=RagVectorDbConfig(
|
||||
rag_managed_db=RagManagedDb()
|
||||
),
|
||||
)
|
||||
op = client.create_rag_corpus(parent=parent, rag_corpus=rag_corpus)
|
||||
result = op.result()
|
||||
# Wrap in a simple namespace so rest of script works
|
||||
class _C:
|
||||
name = result.name
|
||||
corpus = _C()
|
||||
|
||||
print(f"✓ RAG corpus created [{attempt}]: {corpus.name}")
|
||||
echo " Waiting for LRO: ${OPERATION}"
|
||||
for i in $(seq 1 30); do
|
||||
sleep 10
|
||||
LRO=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
|
||||
"https://${RAG_REGION}-aiplatform.googleapis.com/v1beta1/${OPERATION}" 2>/dev/null || echo '{}')
|
||||
DONE=$(echo "${LRO}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('done','false'))" 2>/dev/null || echo 'false')
|
||||
if [[ "${DONE}" == "True" ]] || [[ "${DONE}" == "true" ]]; then
|
||||
CORPUS_NAME=$(echo "${LRO}" | python3 -c "
|
||||
import sys, json
|
||||
d = json.load(sys.stdin)
|
||||
print(d.get('response', {}).get('name', '') or d.get('metadata', {}).get('genericMetadata', {}).get('updateTime', ''))
|
||||
" 2>/dev/null || true)
|
||||
# Fallback: re-list to get name
|
||||
if [[ -z "${CORPUS_NAME}" ]] || [[ "${CORPUS_NAME}" != *ragCorpora* ]]; then
|
||||
CORPUS_NAME=$(curl -sf -H "Authorization: Bearer ${TOKEN}" \
|
||||
"${API}/${PARENT}/ragCorpora" 2>/dev/null \
|
||||
| python3 -c "
|
||||
import sys, json
|
||||
data = json.load(sys.stdin)
|
||||
for c in data.get('ragCorpora', []):
|
||||
if c.get('displayName') == '${RAG_CORPUS_DISPLAY_NAME}':
|
||||
print(c['name'])
|
||||
break
|
||||
" 2>/dev/null || true)
|
||||
fi
|
||||
echo "✓ RAG corpus created: ${CORPUS_NAME}"
|
||||
break
|
||||
fi
|
||||
echo " ... still waiting (${i}/30)"
|
||||
done
|
||||
|
||||
except Exception as e:
|
||||
if attempt == attempts[-1]:
|
||||
print(f"ERROR: All attempts failed. Last error: {e}")
|
||||
sys.exit(1)
|
||||
print(f" [{attempt}] failed: {e} — trying next method...")
|
||||
if [[ -z "${CORPUS_NAME}" ]]; then
|
||||
echo "ERROR: Corpus creation timed out or failed."
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# 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}")
|
||||
# ── Import seed docs via REST ─────────────────────────────────────────────
|
||||
GCS_URI="gs://${CORPUS_BUCKET}/seed/"
|
||||
curl -sf -X POST \
|
||||
-H "Authorization: Bearer ${TOKEN}" \
|
||||
-H "Content-Type: application/json" \
|
||||
"${API}/${CORPUS_NAME}:importRagFiles" \
|
||||
-d '{
|
||||
"importRagFilesConfig": {
|
||||
"gcsSource": { "uris": ["'"${GCS_URI}"'"] },
|
||||
"ragFileChunkingConfig": { "chunkSize": 512, "chunkOverlap": 50 }
|
||||
}
|
||||
}' &>/dev/null && echo "✓ Seed import triggered (async)" || echo " WARNING: Seed import skipped (non-fatal)"
|
||||
|
||||
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
|
||||
# ── Write output ──────────────────────────────────────────────────────────────────
|
||||
echo "${CORPUS_NAME}" > /tmp/rag_corpus_name.txt
|
||||
|
||||
echo ""
|
||||
echo " Corpus resource name : ${CORPUS_NAME}"
|
||||
echo " Region : ${RAG_REGION}"
|
||||
echo ""
|
||||
echo " ACTION REQUIRED — add to .env:"
|
||||
echo " export RAG_CORPUS_NAME=\"${CORPUS_NAME}\""
|
||||
echo " export RAG_REGION=\"${RAG_REGION}\""
|
||||
echo ""
|
||||
echo "=== 07: RAG Engine setup COMPLETE ==="
|
||||
echo " View: https://console.cloud.google.com/vertex-ai/rag?project=${PROJECT_ID}"
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user