diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index fb94419..e0005dd 100644 --- a/agents/rag/setup_corpus.py +++ b/agents/rag/setup_corpus.py @@ -3,40 +3,91 @@ setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents. Project: propane-will-491900-m5 -Serverless mode is the default when no backend_config is specified. -Region: us-central1 +Explicitly uses Serverless (RagManagedDbConfig basic tier) to avoid +Spanner allowlist restrictions on new projects in us-central1. """ import os +import json +import subprocess + 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", "oavauco-knowledge-base") +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/" + f"gs://{PROJECT_ID}-agent-staging/rag-docs/", ) -def main(): - vertexai.init(project=PROJECT_ID, location=LOCATION) - # Check if corpus already exists - existing = list(rag.list_corpora()) - for c in existing: +def ensure_serverless_engine_config() -> None: + """ + Pre-flight: set project-level RAG Engine Config to basic (serverless) tier. + Idempotent — safe to re-run. + """ + print("Ensuring RAG Engine Config is set to serverless (basic) tier...") + endpoint = ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1beta1" + f"/projects/{PROJECT_ID}/locations/{LOCATION}/ragEngineConfig" + ) + payload = json.dumps({"ragManagedDbConfig": {"basic": {}}}) + + try: + token = subprocess.check_output( + ["gcloud", "auth", "print-access-token"], text=True + ).strip() + except subprocess.CalledProcessError: + print(" ⚠️ Could not get gcloud token — skipping pre-flight.") + return + + result = subprocess.run( + [ + "curl", "-s", "-X", "PATCH", + "-H", f"Authorization: Bearer {token}", + "-H", "Content-Type: application/json", + endpoint, "-d", payload, + ], + capture_output=True, text=True, + ) + + resp = result.stdout + if '"error"' in resp: + err = json.loads(resp).get("error", {}) + code, msg = err.get("code"), err.get("message", "") + if code == 400 and "already" in msg.lower(): + print(f" ✓ Already configured: {msg}") + else: + print(f" ⚠️ Engine config warning ({code}): {msg} — continuing anyway.") + else: + print(" ✓ RAG Engine Config set to basic tier.") + + +def get_or_create_corpus() -> RagCorpus: + """Return existing corpus by display name, or create a new serverless one.""" + for c in rag.list_corpora(): if c.display_name == CORPUS_DISPLAY_NAME: print(f"Corpus '{CORPUS_DISPLAY_NAME}' already exists: {c.name}") - corpus = c - break - else: - print(f"Creating RAG corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (Serverless)...") - # No backend_config = Serverless mode (default) - corpus = rag.create_corpus( - display_name=CORPUS_DISPLAY_NAME, - ) - print(f"Corpus created: {corpus.name}") + return c + print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (serverless)...") + corpus = rag.create_corpus( + display_name=CORPUS_DISPLAY_NAME, + # Explicitly request serverless / basic managed DB tier + backend_config=rag.RagVectorDbConfig( + rag_managed_db=rag.RagManagedDbConfig( + tier=rag.RagManagedDbConfig.Tier.BASIC, + ) + ), + ) + print(f"✓ Corpus created: {corpus.name}") + return corpus + + +def import_documents(corpus: RagCorpus) -> None: print(f"Importing files from {GCS_SOURCE}...") rag.import_files( corpus.name, @@ -48,17 +99,33 @@ def main(): ) ), ) - print("Import complete.") - print(f"\nRAG_CORPUS={corpus.name}") - print("Set this as an environment variable or Secret Manager entry.") + print("✓ Import complete.") - # Test retrieval + +def test_retrieval(corpus: RagCorpus) -> None: + print("Running test retrieval query...") response = rag.retrieval_query( rag_resources=[rag.RagResource(rag_corpus=corpus.name)], text="test query", rag_retrieval_config=rag.RagRetrievalConfig(top_k=3), ) - print(f"Test retrieval returned {len(response.contexts.contexts)} chunks.") + print(f"✓ Test retrieval returned {len(response.contexts.contexts)} chunk(s).") + + +def main() -> None: + vertexai.init(project=PROJECT_ID, location=LOCATION) + + ensure_serverless_engine_config() + + corpus = get_or_create_corpus() + + import_documents(corpus) + + print(f"\nRAG_CORPUS={corpus.name}") + print("Add this to Secret Manager:") + print(f" gcloud secrets create rag-corpus-name --data-file=- <<<'{corpus.name}'") + + test_retrieval(corpus) if __name__ == "__main__":