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