73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
setup_corpus.py — Create a Vertex AI RAG Engine corpus and import documents.
|
|
Project: propane-will-491900-m5
|
|
|
|
IMPORTANT: RAG Engine in us-central1 requires allowlist access.
|
|
Contact: vertex-ai-rag-engine-support@google.com
|
|
Alternative: Use us-east1 or us-east4 for immediate access.
|
|
"""
|
|
|
|
import os
|
|
import vertexai
|
|
from vertexai import rag
|
|
|
|
PROJECT_ID = "propane-will-491900-m5"
|
|
LOCATION = os.environ.get("RAG_LOCATION", "us-east1") # change to us-central1 after allowlist
|
|
CORPUS_DISPLAY_NAME = os.environ.get("RAG_CORPUS_NAME", "oavauco-knowledge-base")
|
|
GCS_SOURCE = os.environ.get(
|
|
"RAG_GCS_SOURCE",
|
|
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:
|
|
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}...")
|
|
corpus = rag.create_corpus(
|
|
display_name=CORPUS_DISPLAY_NAME,
|
|
backend_config=rag.RagVectorDbConfig(
|
|
rag_embedding_model_config=rag.RagEmbeddingModelConfig(
|
|
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
|
|
publisher_model="publishers/google/models/text-embedding-005"
|
|
)
|
|
)
|
|
),
|
|
)
|
|
print(f"Corpus created: {corpus.name}")
|
|
|
|
print(f"Importing files from {GCS_SOURCE}...")
|
|
rag.import_files(
|
|
corpus.name,
|
|
paths=[GCS_SOURCE],
|
|
transformation_config=rag.TransformationConfig(
|
|
chunking_config=rag.ChunkingConfig(
|
|
chunk_size=512,
|
|
chunk_overlap=100,
|
|
)
|
|
),
|
|
)
|
|
print("Import complete.")
|
|
print(f"\nRAG_CORPUS={corpus.name}")
|
|
print("Set this as an environment variable or Secret Manager entry.")
|
|
|
|
# Test retrieval
|
|
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.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|