Generative AI Google Brain

gemini-rag-engine-rag-engine-feature-store

3 min read

Vertex AI RAG Engine with Vertex AI Feature Store

Source notebook

Repo path: gemini/rag-engine/rag_engine_feature_store.ipynb · Open on GitHub · advanced

Builds a Vertex AI RAG Engine corpus backed by Vertex AI Feature Store and queries it with Gemini.

Summary

This notebook teaches how to use Vertex AI Feature Store as a vector database for Vertex AI RAG Engine. It creates a BigQuery source table, provisions an optimized FeatureOnlineStore and FeatureView, creates a RAG corpus backed by that FeatureView, imports files from Cloud Storage, syncs the FeatureView index, and uses the corpus as retrieval context for Gemini generation. It also shows direct retrieval with rag.retrieval_query for use with other generation APIs.

Key code patterns

Initialize Vertex AI and GenAI clients

vertexai.init(project=PROJECT_ID, location="us-east1")
client = genai.Client(vertexai=True, project=PROJECT_ID, location="global")

Sets up Vertex AI SDK for RAG resources and Google GenAI SDK for Gemini generation.

Create RAG source table

schema = [
    bigquery.SchemaField("corpus_id", "STRING", mode="REQUIRED"),
    bigquery.SchemaField("file_id", "STRING", mode="REQUIRED"),
    bigquery.SchemaField("chunk_id", "STRING", mode="REQUIRED"),
    bigquery.SchemaField("embeddings", "FLOAT64", mode="REPEATED"),
]
table = bq_client.create_table(bigquery.Table(table_ref, schema=schema))

Defines the BigQuery schema required by Feature Store as the RAG data source.

Provision Feature Store vector backend

fos = feature_store.FeatureOnlineStore.create_optimized_store(FEATURE_ONLINE_STORE_ID)
fv = fos.create_feature_view(
    name=FEATURE_VIEW_ID,
    source=feature_store.utils.FeatureViewVertexRagSource(uri=BIGQUERY_TABLE),
)

Connects the BigQuery table to an optimized FeatureOnlineStore through a FeatureView.

Create RAG corpus using Feature Store

vector_db = rag.VertexFeatureStore(resource_name=fv.resource_name)
rag_corpus = rag.create_corpus(
    display_name="Feature Store Corpus",
    vector_db=vector_db,
)

Associates a RAG corpus with the Vertex AI Feature Store vector database.

Import and sync RAG files

response = rag.import_files(
    corpus_name=rag_corpus.name,
    paths=[GCS_BUCKET],
    chunk_size=512,
    chunk_overlap=50,
)
feature_view_sync = fv.sync()
feature_view_sync.wait()

Embeds imported files into the BigQuery-backed FeatureView and makes them available for online serving.

Use RAG retrieval with Gemini

rag_retrieval_tool = Tool(
    retrieval=Retrieval(
        vertex_rag_store=VertexRagStore(
            rag_corpora=[rag_corpus.name], similarity_top_k=10,
            vector_distance_threshold=0.4,
        )
    )
)
response = client.models.generate_content(
    model=MODEL_ID, contents=GENERATE_CONTENT_PROMPT,
    config=GenerateContentConfig(tools=[rag_retrieval_tool]),
)

Adds retrieved corpus context to Gemini generation through a retrieval tool.

Models & APIs used

  • Models: gemini-3.5-flash
  • APIs / services: Vertex AI, Vertex AI RAG Engine, Vertex AI Feature Store, BigQuery, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, google-genai, google-cloud-bigquery, vertexai

When to use this

Use this pattern when a RAG application needs Vertex AI RAG Engine with Feature Store for scalable, low-latency vector retrieval.

Gotchas & caveats

  • Vertex AI API must be enabled for the Google Cloud project.
  • The notebook initializes Vertex AI in us-east1 and the GenAI client in global.
  • RAG Engine supported regions should be checked before choosing a location.
  • Colab requires authenticate_user before using Google Cloud resources.
  • Installing packages requires a runtime restart.
  • FeatureOnlineStore first-time provisioning might take about five minutes.
  • FeatureView sync might take 20 minutes to complete.
  • The Vertex RAG Data Service Agent needs Viewer access to the Cloud Storage bucket.
  • The notebook states currently only one corpus is allowed in rag.retrieval_query.

Best practices

  • Use the required BigQuery schema fields for the RAG Feature Store source table.
  • Use optimized online serving for FeatureOnlineStore when using vector similarity search.
  • Use chunk_size and chunk_overlap when importing files into the RAG corpus.
  • List imported files after import because processing may take a few seconds.
  • Run FeatureView sync after uploading data to make it available for online serving.
  • Set similarity_top_k and vector_distance_threshold when configuring VertexRagStore retrieval.
  • Use retrieved contexts from rag.retrieval_query with other SDKs or model generation APIs when needed.
  • Include cleanup logic to delete the RAG corpus when desired.

Graph View

Backlinks