diff --git a/infrastructure/07-rag-setup.sh b/infrastructure/07-rag-setup.sh index bc17208..b57b07a 100644 --- a/infrastructure/07-rag-setup.sh +++ b/infrastructure/07-rag-setup.sh @@ -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" +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 +" 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 -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}") + 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": {} + } + }') -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" - ) + # 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) - # 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, - ) + if [[ -z "${OPERATION}" ]]; then + echo "ERROR: No operation returned. Response:" + echo "${RESPONSE}" + exit 1 + fi - 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, - ) + 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 - 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 "${CORPUS_NAME}" ]]; then + echo "ERROR: Corpus creation timed out or failed." + exit 1 + fi +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() +# ── 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)" - print(f"✓ RAG corpus created [{attempt}]: {corpus.name}") - break - - 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 +# ── 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}"