Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-rag-qna-with-bq-and-featurestore

4 min read

Building a Gen AI RAG application with Vertex AI Feature Store and BigQuery

Source notebook

Repo path: gemini/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.ipynb · Open on GitHub · intermediate

Builds a LangChain RAG Q&A app using BigQuery Vector Search and Vertex AI Feature Store.

Summary

This notebook teaches a RAG workflow that loads a car manual PDF from Cloud Storage, chunks it, creates text embeddings, and stores them in BigQuery Vector Search. It then builds a LangChain RetrievalQA chain with Gemini and shows how to move the same BigQuery-backed data into Vertex AI Feature Store for low-latency online retrieval. It also compares BigQuery batch search with Feature Store single-request serving and covers metadata filtering, document lookup, deletion, MMR, and cleanup.

Key code patterns

Initialize Vertex AI and embeddings

PROJECT_ID = "[your-project-id]"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)
embedding_model = VertexAIEmbeddings(
    model_name="text-embedding-005", project=PROJECT_ID
)

Sets the project, region, and embedding model used by both vector stores.

Load and chunk PDF corpus

loader = PyPDFLoader("cymbal-starlight-2024.pdf")
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=50,
    separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
)
doc_splits = text_splitter.split_documents(documents)
for idx, split in enumerate(doc_splits):
    split.metadata["chunk"] = idx

Creates context-sized chunks and stores chunk IDs in metadata.

Ingest into BigQuery vector store

bq_store = BigQueryVectorStore(
    project_id=PROJECT_ID,
    location=LOCATION,
    dataset_name=DATASET,
    table_name=TABLE,
    embedding=embedding_model,
)
doc_ids = bq_store.add_documents(doc_splits)
bq_store.similarity_search(search_query)

Uses BigQuery Vector Search to store embeddings and retrieve similar chunks.

Build RetrievalQA with Gemini

set_debug(True)
llm = VertexAI(model_name="gemini-2.0-flash")
langchain_retriever = bq_store.as_retriever()
retrieval_qa = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff", retriever=langchain_retriever
)
response = retrieval_qa.invoke(search_query)
print(response["result"])

Connects retrieved chunks to Gemini through a LangChain RetrievalQA chain.

Move to Feature Store serving

vertex_fs = bq_store.to_vertex_fs_vector_store()
vertex_fs = VertexFSVectorStore(
    project_id=PROJECT_ID, location=LOCATION,
    dataset_name=DATASET, table_name=TABLE,
    embedding=embedding_model,
)
vertex_fs.sync_data()
langchain_retriever = vertex_fs.as_retriever()

Synchronizes BigQuery data to Feature Online Store for low-latency retrieval.

Batch and vector latency tests

my_embedding = embedding_model.embed(search_query)[0]
results = bq_store.batch_search(
    embeddings=None,
    queries=["search_query", "search_query"],
)
fake_embeddings = [my_embedding] * 10000
results = bq_store.batch_search(embeddings=fake_embeddings)
vertex_fs.similarity_search_by_vector(my_embedding)

Separates embedding latency from retrieval latency and shows BigQuery batch search.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI API, BigQuery, BigQuery Vector Search, Vertex AI Feature Store, Feature Online Store, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langchain, langchain-google-vertexai, langchain-google-community, pypdf

When to use this

Use this pattern when building a RAG app that starts with BigQuery-based vector search and later needs low-latency online retrieval through Vertex AI Feature Store.

Gotchas & caveats

  • After package installation, the notebook restarts the Jupyter runtime before continuing.
  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • A Google Cloud project is required and the Vertex AI API must be enabled.
  • The notebook uses LOCATION = “us-central1”.
  • The first Feature Store synchronization can take around 20 minutes because the Feature Online Store is created.
  • Feature Store metadata filtering during search requires filter_columns when the feature view is first created.
  • Only string fields are supported for Feature Store string_filters at the time referenced by the notebook.
  • Server-side latency should be checked in the Feature Store Serving Latency dashboard.

Best practices

  • Use BigQueryVectorStore for prototyping because it requires no infrastructure startup time.
  • Split documents so a few chunks can fit within the LLM context length.
  • Store document source, document name, and chunk number in metadata.
  • Use add_texts_with_embeddings when embeddings are already precomputed.
  • Use VertexFSVectorStore for production-ready user-facing Gen AI applications that need low-latency retrieval.
  • Use BigQueryVectorStore batch_search for batch searches such as retrieval evaluation.
  • Precompute embeddings when comparing retrieval latency.
  • Delete the feature view and online store during cleanup.

Graph View

Backlinks