gemini-rag-engine-rag-engine-vector-search

Vertex AI RAG Engine with Vertex AI Vector Search

Source notebook

Repo path: gemini/rag-engine/rag_engine_vector_search.ipynb · Open on GitHub · intermediate

Builds a Vertex AI RAG Engine corpus backed by Vertex AI Vector Search and queries it with Gemini.

Summary

This notebook teaches how to use Vertex AI RAG Engine with Vertex AI Vector Search as the vector database. It walks through project setup, optional Vector Search index and endpoint creation, RAG corpus creation, file ingestion from local files, Cloud Storage, and Drive, then Gemini generation with a RagStore retrieval tool. It also shows direct retrieval_query usage so retrieved context can be passed to other generation APIs.

Key code patterns

Initialize clients

aiplatform.init(project=PROJECT_ID, location=LOCATION)
client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Initializes Vertex AI SDK and Google GenAI SDK for Vertex AI-backed Gemini calls.

Create stream Vector Search index

my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    dimensions=768,
    distance_measure_type="DOT_PRODUCT_DISTANCE",
    index_update_method="STREAM_UPDATE",
)

RAG Engine requires a stream-update Vector Search index with compatible distance and dimensions.

Create RAG corpus with Vector Search

vector_db = rag.VertexVectorSearch(
    index=my_index.resource_name,
    index_endpoint=my_index_endpoint.resource_name,
)
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME,
    backend_config=rag.RagVectorDbConfig(vector_db=vector_db),
)

Connects a RAG corpus to an existing Vector Search index and endpoint.

Import chunked files

response = rag.import_files(
    corpus_name=rag_corpus.name,
    paths=[GCS_BUCKET],
    transformation_config=rag.TransformationConfig(
        chunking_config=rag.ChunkingConfig(chunk_size=512, chunk_overlap=50)
    ),
)

Shows ingestion with explicit chunk size and overlap for retrieval preparation.

Use RagStore tool with Gemini

rag_retrieval_tool = Tool(retrieval=Retrieval(
    vertex_rag_store=VertexRagStore(
        rag_resources=[VertexRagStoreRagResource(rag_corpus=rag_corpus.name)],
        similarity_top_k=10,
        vector_distance_threshold=0.4,
    )
))

Adds retrieved corpus context to Gemini generation through a Google GenAI tool.

Run direct retrieval query

response = rag.retrieval_query(
    rag_resources=[rag.RagResource(rag_corpus=rag_corpus.name)],
    text=RETRIEVAL_QUERY,
    rag_retrieval_config=rag.RagRetrievalConfig(top_k=10),
)

Retrieves context separately so it can be passed to any SDK or generation API.

Models & APIs used

• Models: gemini-2.5-flash
• APIs / services: Vertex AI, Vertex AI RAG Engine, Vertex AI Vector Search, Cloud Storage, Google Drive
• SDKs / libraries: google-cloud-aiplatform, google-genai, vertexai

When to use this

Use this pattern when a RAG application needs Vertex AI RAG Engine with Vertex AI Vector Search as its vector database.

Gotchas & caveats

• The Vertex AI API must be enabled for the Google Cloud project.
• The notebook installs upgraded packages and requires a runtime restart.
• Colab requires explicit user authentication.
• The Vector Search index must use STREAM_UPDATE.
• The Vector Search distance measure must be DOT_PRODUCT_DISTANCE or COSINE_DISTANCE.
• Vector dimensions must match the embedding model used by the RAG corpus.
• First index deployment to an endpoint can take approximately 30 minutes.
• Cloud Storage and Drive imports require Viewer access for the Vertex RAG Data Service Agent.
• Only one corpus is currently allowed in the shown RagStore and retrieval_query patterns.
• Imported files may take a few seconds to process before listing.

Best practices

• Use environment variables for project and region defaults when notebook parameters are not provided.
• Create a public Vector Search index endpoint because RAG Engine supports public endpoints.
• Tune chunk_size and chunk_overlap during file import.
• Use similarity_top_k and vector_distance_threshold to control retrieved context.
• List files after import to check ingestion progress.
• Clean up the created RAG corpus when it is no longer needed.

Related

• Concepts: RAG & Grounding · Embeddings & Vector Search · Gemini Capabilities
• Entities: Vertex AI · Google GenAI SDK · Vertex AI SDK · Vector Search · Cloud Storage · Gemini
• Area: Gemini Notebooks
• Best practices: RAG & Grounding - Best Practices · Embeddings & Vector Search - Best Practices · Gemini Capabilities - Best Practices