Generative AI Google Brain

gemini-rag-engine-rag-engine-vertex-ai-search

3 min read

Vertex AI RAG Engine with Vertex AI Search

Source notebook

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

Build a Vertex AI RAG Engine corpus backed by Vertex AI Search and query it with Gemini.

Summary

This notebook teaches how to use Vertex AI RAG Engine with Vertex AI Search as the retrieval backend. It optionally creates a Vertex AI Search data store and search engine, imports documents from Cloud Storage, creates a RAG corpus from the search serving config, and queries it with a Gemini retrieval tool. It also shows direct retrieval_query usage so retrieved contexts can be passed to another generation API.

Key code patterns

Initialize Vertex AI

import os
import vertexai
 
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets project and region before using Vertex AI SDK and RAG Engine resources.

Create Vertex AI Search datastore

client = discoveryengine.DataStoreServiceClient(client_options=client_options)
data_store = discoveryengine.DataStore(
    display_name=data_store_name,
    industry_vertical=discoveryengine.IndustryVertical.GENERIC,
    content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED,
)
operation = client.create_data_store(request=discoveryengine.CreateDataStoreRequest(...))

Creates the Discovery Engine data store used as the Vertex AI Search retrieval source.

Import GCS documents

client = discoveryengine.DocumentServiceClient(client_options=client_options)
request = discoveryengine.ImportDocumentsRequest(
    parent=parent,
    gcs_source=discoveryengine.GcsSource(input_uris=[f"{gcs_uri}/*"], data_schema="content"),
    reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
)
operation = client.import_documents(request=request)

Loads Cloud Storage documents into the Vertex AI Search data store.

Create RAG corpus from search engine

vertex_ai_search_config = rag.VertexAiSearchConfig(
    serving_config=f"{ENGINE_NAME}/servingConfigs/default_search",
)
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME,
    vertex_ai_search_config=vertex_ai_search_config,
)

Connects RAG Engine to an existing Vertex AI Search serving config.

Use Gemini with RAG retrieval tool

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(source=rag.VertexRagStore(
        rag_resources=[rag_resource], similarity_top_k=10
    ))
)
rag_model = GenerativeModel("gemini-2.0-flash", tools=[rag_retrieval_tool])
response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)

Lets Gemini call the RAG retrieval tool during GenerateContent.

Retrieve contexts directly

response = rag.retrieval_query(
    rag_resources=[rag_resource],
    text=RETRIEVAL_QUERY,
    similarity_top_k=10,
)
retrieved_context = " ".join([c.text for c in response.contexts.contexts]).replace("\n", "")

Separates retrieval from generation so contexts can be passed to any SDK or model API.

Models & APIs used

  • Models: gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI RAG Engine, Vertex AI Search, Discovery Engine API, Cloud Storage, Gemini GenerateContent API
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-discoveryengine, vertexai

When to use this

Use this pattern when a RAG application needs scalable, low-latency retrieval from Vertex AI Search through Vertex AI RAG Engine.

Gotchas & caveats

  • The notebook requires an existing Google Cloud project with Vertex AI API enabled.
  • Creating Vertex AI Search resources requires the Discovery Engine API to be enabled.
  • Installed packages require a Jupyter runtime restart before continuing.
  • Colab requires explicit auth.authenticate_user(project_id=PROJECT_ID).
  • Vertex AI Search datastore names can only contain lowercase letters, numbers, and hyphens.
  • The Vertex AI Search engine can take a few minutes to become queryable after creation.
  • A recently created engine may return 404 Engine is not found until it is ready.
  • The RAG tool examples note that currently only one corpus is allowed.

Best practices

  • Use environment variables for GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_REGION when notebook parameters are not provided.
  • Use Vertex AI Search as the retrieval backend for large datasets, low-latency retrieval, and improved scalability.
  • Import documents from Cloud Storage using INCREMENTAL reconciliation mode.
  • Check the created corpus with rag.get_corpus before querying.
  • Clean up created RAG resources with rag.delete_corpus when delete_rag_corpus is enabled.

Graph View

Backlinks