Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-llamaindex-rag

4 min read

Leverage LlamaIndex with Vertex AI Vector Search to perform question answering RAG

Source notebook

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

Builds LlamaIndex RAG on Vertex AI Vector Search, compares prompts, and adds multi-document agents.

Summary

This notebook builds a RAG question-answering system over PDF documents using LlamaIndex, Vertex AI Vector Search, and Gemini. It provisions Cloud Storage and Vector Search resources, loads sample power-grid PDFs, embeds them with text-embedding-005, and queries them with gemini-2.5-flash. It compares LlamaIndex built-in, LangChain, and custom grounded few-shot prompts, then extends the workflow to multi-document ReAct agents with vector and summary tools.

Key code patterns

Vector Search setup

vertexai.init(project=PROJECT_ID, location="us-central1")
aiplatform.init(project=PROJECT_ID, location=REGION)
vs_index = create_vector_search_index(VS_INDEX_NAME, 768)
vs_endpoint = create_vector_search_endpoint(VS_INDEX_ENDPOINT_NAME)
deployed_endpoint = deploy_vector_search_endpoint(vs_index, vs_endpoint, VS_INDEX_NAME)

Creates or reuses the cloud resources needed before ingestion.

LlamaIndex vector store

vector_store = VertexAIVectorStore(
    project_id=PROJECT_ID,
    region=REGION,
    index_id=vs_index.resource_name,
    endpoint_id=vs_endpoint.resource_name,
    gcs_bucket_name=GCS_BUCKET,
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

Connects LlamaIndex storage to the deployed Vertex AI Vector Search index.

Vertex models

Settings.embed_model = VertexTextEmbedding("text-embedding-005")
Settings.llm = Vertex("gemini-2.5-flash")

Sets the embedding model and generation model used by LlamaIndex.

Index and query

docs = SimpleDirectoryReader(DOC_FOLDER).load_data()
vector_index = VectorStoreIndex.from_documents(
    docs,
    storage_context=storage_context,
)
query_engine = vector_index.as_query_engine()
response = query_engine.query("what is minimum reserve rate of power?")

Loads local documents, embeds them into Vector Search, and runs RAG queries.

LangChain prompt swap

langchain_prompt = hub.pull("rlm/rag-prompt")
template = LangchainPromptTemplate(
    template=langchain_prompt,
    template_var_mappings={"query_str": "question", "context_str": "context"},
)
query_engine.update_prompts({"response_synthesizer:text_qa_template": template})

Replaces the default LlamaIndex QA prompt with a LangChain RAG template.

Custom grounded prompt

qa_prompt_custom_string = """Context information is below.
{context_str}
Given the context information and not prior knowledge...
Please output your answer in the following JSON format...
Query: {query_str}
Answer:"""
custom_RAG_template = PromptTemplate(template=qa_prompt_custom_string)
query_engine.update_prompts({"response_synthesizer:text_qa_template": custom_RAG_template})

Uses context-only instructions, few-shot examples, JSON output, and justification.

Document agents

vector_query_engine = vector_index.as_query_engine()
summary_query_engine = summary_index.as_query_engine()
query_engine_tools = [
    QueryEngineTool(query_engine=vector_query_engine, metadata=ToolMetadata(name="vector_tool")),
    QueryEngineTool(query_engine=summary_query_engine, metadata=ToolMetadata(name="summary_tool")),
]
agent = ReActAgent.from_tools(query_engine_tools, llm=Vertex("gemini-2.5-flash"), verbose=True)

Builds per-document agents that can choose semantic search or summarization.

Top-level agent

obj_index = ObjectIndex.from_objects(all_tools, index_cls=VectorStoreIndex)
top_agent = ReActAgent.from_tools(
    tool_retriever=obj_index.as_retriever(similarity_top_k=3),
    system_prompt="Please always use the tools provided to answer a question.",
    verbose=True,
)

Routes multi-document questions to the most relevant document-agent tools.

Models & APIs used

  • Models: text-embedding-005, gemini-2.5-flash
  • APIs / services: Vertex AI, Vertex AI Vector Search, Cloud Storage
  • SDKs / libraries: vertexai, google-cloud-aiplatform, google-cloud-storage, llama-index, llama-index-embeddings-vertex, llama-index-llms-vertex, llama-index-vector_stores-vertexaivectorsearch, langchain, langchainhub

When to use this

Use this pattern when you need grounded QA and summaries over private documents using LlamaIndex, Vertex AI Vector Search, and Gemini.

Gotchas & caveats

  • Requires an initialized Google Cloud project, Vertex AI API enabled, and an existing VPC/Subnet.
  • Defaults to us-central1 and user-provided project, bucket, index, and endpoint names.
  • Vector Search index creation can take up to 30 minutes.
  • Uses billable Vertex AI and Cloud Storage resources.
  • The notebook requires a runtime restart after package installation.
  • Colab runs auth.authenticate_user() separately.
  • Vector Search dimensions are set to 768 for the configured text embeddings.
  • The setup code uses both GCS_BUCKET_URI and GCS_BUCKET, so bucket URI versus bucket name format must be kept consistent.

Best practices

  • Check for existing buckets, indexes, and endpoints before creating resources.
  • Set LlamaIndex embed_model and llm settings before building indexes.
  • Use SimpleDirectoryReader and SentenceSplitter to parse documents into documents and nodes.
  • Display prompt templates before running RAG queries.
  • Inspect source text, relevance score, file name, page label, and file path with each response.
  • Compare built-in, LangChain, and custom few-shot prompts on the same query engine.
  • Instruct custom prompts and agents to use retrieved context or tools instead of prior knowledge.
  • Pair vector search tools with summary tools for document-level agents.

Graph View

Backlinks