133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
sync_corpus.py — Synchronize RAG Engine corpus with GCS and perform health checks.
|
|
Project: propane-will-491900-m5
|
|
|
|
Capabilities:
|
|
- Locates existing corpus by display name.
|
|
- Lists current files in the corpus.
|
|
- Triggers incremental import from GCS.
|
|
- Performs a health-check retrieval query.
|
|
- Reports status and chunk counts.
|
|
|
|
Graceful degradation:
|
|
- If corpus is not found (e.g. due to Spanner restriction), logs a warning and exits cleanly.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import time
|
|
import subprocess
|
|
import sys
|
|
|
|
import vertexai
|
|
from vertexai import rag
|
|
from vertexai.rag import RagCorpus
|
|
|
|
PROJECT_ID = "propane-will-491900-m5"
|
|
LOCATION = os.environ.get("RAG_LOCATION", "us-central1")
|
|
CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "osvauco-knowledge-base")
|
|
GCS_SOURCE = os.environ.get(
|
|
"RAG_GCS_SOURCE",
|
|
f"gs://{PROJECT_ID}-agent-staging/rag-docs/",
|
|
)
|
|
|
|
def get_token() -> str:
|
|
return subprocess.check_output(
|
|
["gcloud", "auth", "print-access-token"], text=True
|
|
).strip()
|
|
|
|
def find_corpus() -> RagCorpus:
|
|
"""Search for the corpus in the configured location and fallback location."""
|
|
locations_to_try = [LOCATION]
|
|
if LOCATION != "europe-west4":
|
|
locations_to_try.append("europe-west4")
|
|
|
|
for loc in locations_to_try:
|
|
print(f"Searching for corpus '{CORPUS_DISPLAY_NAME}' in {loc}...")
|
|
vertexai.init(project=PROJECT_ID, location=loc)
|
|
try:
|
|
for c in rag.list_corpora():
|
|
if c.display_name == CORPUS_DISPLAY_NAME:
|
|
print(f" \u2713 Found: {c.name}")
|
|
return c
|
|
except Exception as e:
|
|
print(f" \u26a0\ufe0f Notice: Could not list corpora in {loc}: {e}")
|
|
|
|
return None
|
|
|
|
def list_files(corpus: RagCorpus):
|
|
"""List files currently in the corpus."""
|
|
print(f"Listing files in corpus...")
|
|
files = list(rag.list_files(corpus_name=corpus.name))
|
|
if not files:
|
|
print(" (Corpus is empty)")
|
|
return []
|
|
|
|
for f in files:
|
|
print(f" - {f.display_name} (state: {f.state.name})")
|
|
return files
|
|
|
|
def sync_from_gcs(corpus: RagCorpus):
|
|
"""Trigger incremental import from GCS."""
|
|
print(f"Syncing from {GCS_SOURCE}...")
|
|
# import_files in SDK is essentially an incremental trigger
|
|
rag.import_files(
|
|
corpus.name,
|
|
paths=[GCS_SOURCE],
|
|
transformation_config=rag.TransformationConfig(
|
|
chunking_config=rag.ChunkingConfig(
|
|
chunk_size=512,
|
|
chunk_overlap=100,
|
|
)
|
|
),
|
|
)
|
|
print(" \u2713 Sync job submitted.")
|
|
|
|
def run_health_check(corpus: RagCorpus):
|
|
"""Perform a retrieval query and report results."""
|
|
print("Running health check retrieval...")
|
|
try:
|
|
response = rag.retrieval_query(
|
|
rag_resources=[rag.RagResource(rag_corpus=corpus.name)],
|
|
text="How does OSVauco handle operator experience?",
|
|
rag_retrieval_config=rag.RagRetrievalConfig(top_k=3),
|
|
)
|
|
contexts = response.contexts.contexts
|
|
print(f" \u2713 Health check returned {len(contexts)} contexts.")
|
|
for i, ctx in enumerate(contexts):
|
|
print(f" [{i+1}] Source: {ctx.source_uri} (Score: {ctx.score:.4f})")
|
|
except Exception as e:
|
|
print(f" \u2718 Health check failed: {e}")
|
|
|
|
def main():
|
|
print(f"--- RAG Sync & Maintenance Tool ---")
|
|
|
|
corpus = find_corpus()
|
|
if not corpus:
|
|
print(f"\n\u26a0\ufe0f WARN: Corpus '{CORPUS_DISPLAY_NAME}' not found.")
|
|
print(" RAG Engine might be restricted for this project or setup_corpus.py has not been run.")
|
|
print(" Skipping sync and health check.")
|
|
sys.exit(0)
|
|
|
|
# Update global LOCATION if we found it in a different one
|
|
found_loc = corpus.name.split('/')[3]
|
|
print(f"Active Location: {found_loc}")
|
|
|
|
list_files(corpus)
|
|
|
|
# Always trigger a sync to ensure new files are picked up
|
|
sync_from_gcs(corpus)
|
|
|
|
print("\nWaiting 5s for sync initialization...")
|
|
time.sleep(5)
|
|
|
|
run_health_check(corpus)
|
|
|
|
print(f"\nCOMPLETED")
|
|
print(f"RAG_LOCATION={found_loc}")
|
|
print(f"RAG_CORPUS={corpus.name}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|