Generative AI Google Brain

gemini-rag-engine-rag-engine-cross-corpus-retrieval

2 min read

Vertex AI Rag: Cross-Corpus Retrieval with AskContexts and AsyncRetrieveContexts Demo

Source notebook

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

Demonstrates cross-corpus Vertex AI RAG retrieval with ask_contexts and async_retrieve_contexts.

Summary

This notebook teaches how to query multiple Vertex AI RAG corpora in one request using the Vertex AI Python SDK preview RAG module. It shows installation, Colab authentication, Vertex AI initialization, configuring two existing corpus paths as RagResource objects, and running synchronous and asynchronous context retrieval.

Key code patterns

Initialize Vertex AI

from google.cloud import aiplatform
from vertexai.preview import rag
 
PROJECT_ID = "[your-project-id]"
LOCATION = "us-central1"
ENDPOINT = "us-central1-aiplatform.googleapis.com"
 
aiplatform.init(
    project=PROJECT_ID,
    location=LOCATION,
    api_endpoint=ENDPOINT,
)

Sets the project, region, and regional API endpoint before using Vertex AI RAG.

Build cross-corpus resources

CORPUS_PATH_1 = "[your-corpus-path]"
CORPUS_PATH_2 = "[your-corpus-path]"
 
rag_resources_list = [
    rag.RagResource(rag_corpus=CORPUS_PATH_1),
    rag.RagResource(rag_corpus=CORPUS_PATH_2),
]

Wraps multiple RAG corpus paths so one retrieval call can search across corpora.

Synchronous context retrieval

config = rag.RagRetrievalConfig(top_k=2)
query = "[Your query]"
 
response = rag.ask_contexts(
    text=query,
    rag_resources=rag_resources_list,
    rag_retrieval_config=config,
)

Uses ask_contexts to retrieve relevant contexts from multiple corpora in a blocking flow.

Asynchronous context retrieval

query = "[Your query]"
 
response = await rag.async_retrieve_contexts(
    text=query,
    rag_resources=rag_resources_list,
    rag_retrieval_config=config,
)

Uses async_retrieve_contexts for non-blocking retrieval across multiple corpora.

Models & APIs used

  • APIs / services: Vertex AI, Vertex AI RAG, Vertex AI API
  • SDKs / libraries: google-cloud-aiplatform, vertexai, vertexai.preview.rag

When to use this

Use this pattern when a Vertex AI RAG knowledge base is split across multiple corpora and one query should retrieve contexts across them.

Gotchas & caveats

  • The AskContexts and AsyncRetrieveContexts handlers are in preview and require vertexai.preview.rag.
  • The notebook requires an existing Google Cloud project with the Vertex AI API enabled.
  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • The example uses the us-central1 location and us-central1-aiplatform.googleapis.com endpoint.
  • The notebook assumes existing RAG corpus paths unless you create corpora separately.

Best practices

  • Initialize Vertex AI with an explicit project, location, and API endpoint.
  • Use RagResource objects to pass corpus paths into retrieval calls.
  • Configure retrieval behavior with RagRetrievalConfig, including top_k.
  • Use async_retrieve_contexts when integrating retrieval into an asynchronous application flow.

Graph View

Backlinks