Generative AI Google Brain

gemini-rag-engine-intro-rag-engine

3 min read

Intro to Building a Scalable and Modular RAG System with RAG Engine in Vertex AI

Source notebook

Repo path: gemini/rag-engine/intro_rag_engine.ipynb · Open on GitHub · intro

Builds a Vertex AI RAG Engine corpus, imports files, retrieves context, and grounds Gemini or Llama responses.

Summary

This notebook introduces Vertex AI RAG Engine as a way to ground LLM responses with private or external data. It walks through project setup, corpus creation with a Google embedding model, local file upload, Cloud Storage and Google Drive imports with chunking, direct retrieval, and generation through RAG retrieval tools. It demonstrates Gemini generation with the Google GenAI SDK and a Llama3 pattern through a self-deployed Vertex AI endpoint.

Key code patterns

Initialize clients

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

Sets the Vertex AI project and locations before using RAG Engine and Gemini generation.

Create RAG corpus

EMBEDDING_MODEL = "publishers/google/models/text-embedding-005"
rag_corpus = rag.create_corpus(
    display_name="my-rag-corpus",
    backend_config=rag.RagVectorDbConfig(
        rag_embedding_model_config=rag.RagEmbeddingModelConfig(
            vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
                publisher_model=EMBEDDING_MODEL
            )
        )
    ),
)

Creates the searchable RAG corpus and binds it to a supported Google embedding model.

Upload local file

rag_file = rag.upload_file(
    corpus_name=rag_corpus.name,
    path="test.md",
    display_name="test.md",
    description="my test file",
)

Adds a local Markdown file to the corpus for retrieval.

Import Cloud Storage files

response = rag.import_files(
    corpus_name=rag_corpus.name,
    paths=[INPUT_GCS_BUCKET],
    transformation_config=rag.TransformationConfig(
        chunking_config=rag.ChunkingConfig(chunk_size=1024, chunk_overlap=100)
    ),
    max_embedding_requests_per_min=900,
)

Imports a GCS folder and controls chunking and embedding request rate.

Direct retrieval query

response = rag.retrieval_query(
    rag_resources=[rag.RagResource(rag_corpus=rag_corpus.name)],
    rag_retrieval_config=rag.RagRetrievalConfig(
        top_k=10,
        filter=rag.Filter(vector_distance_threshold=0.5),
    ),
    text="What is RAG and why it is helpful?",
)

Retrieves matching corpus context directly before or apart from model generation.

Generate with RAG tool

rag_retrieval_tool = Tool(
    retrieval=Retrieval(
        vertex_rag_store=VertexRagStore(
            rag_corpora=[rag_corpus.name],
            similarity_top_k=10,
            vector_distance_threshold=0.5,
        )
    )
)
response = client.models.generate_content(
    model="gemini-3.5-flash",
    contents="What is RAG?",
    config=GenerateContentConfig(tools=[rag_retrieval_tool]),
)

Lets Gemini call the RAG corpus as a retrieval tool during content generation.

Models & APIs used

  • Models: publishers/google/models/text-embedding-005, gemini-3.5-flash
  • APIs / services: Vertex AI, Vertex AI RAG Engine, Cloud Storage, Google Drive
  • SDKs / libraries: google-cloud-aiplatform, google-genai, vertexai

When to use this

Use this pattern when you need a managed Vertex AI RAG corpus that can ingest documents and ground Gemini or endpoint-hosted model answers.

Gotchas & caveats

  • Requires an existing Google Cloud project and the Vertex AI API enabled.
  • Colab users must run auth.authenticate_user().
  • The notebook uses us-east1 for Vertex AI SDK initialization and global for the Google GenAI client.
  • RAG Engine supported regions must be checked before choosing a location.
  • Corpus creation notes that RAG Engine currently supports Google first-party embedding models.
  • Cloud Storage and Google Drive imports require Viewer access for the Vertex RAG Data Service Agent.
  • The Llama3 example requires a self-deployed Vertex AI endpoint.

Best practices

  • Use an explicit embedding model when creating the RAG corpus.
  • Configure chunk_size and chunk_overlap during bulk imports.
  • Use top_k and vector_distance_threshold to tune retrieval scope.
  • Run rag.retrieval_query to inspect retrieved context directly.
  • Pass a RAG retrieval tool into generate_content to ground model responses.
  • Throttle import embedding traffic with max_embedding_requests_per_min when importing from Cloud Storage.

Graph View

Backlinks