Generative AI Google Brain

gemini-orchestration-llamaindex-workflows

3 min read

LlamaIndex RAG Workflows using Gemini and Firestore

Source notebook

Repo path: gemini/orchestration/llamaindex_workflows.ipynb · Open on GitHub · advanced

Builds a LlamaIndex RAG workflow with Gemini, Vertex embeddings, and Firestore storage.

Summary

This notebook teaches how to orchestrate an event-driven LlamaIndex workflow for Retrieval Augmented Generation. It demonstrates one branch that ingests local text files into a Firestore-backed docstore and builds a VectorStoreIndex, and another branch that decomposes a complex query into sub-questions, answers them, reranks results, adds citations, and synthesizes a final answer.

Key code patterns

Initialize Vertex AI

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the Google Cloud project and region before using Vertex-hosted models.

Configure Gemini and embeddings

embedding_model = VertexTextEmbedding(
    model_name="text-embedding-005",
    credentials=credentials,
)
llm = Vertex(
    model="gemini-2.5-flash",
    temperature=0.2,
    max_tokens=3000,
    safety_settings=safety_config,
    credentials=credentials,
)
Settings.embed_model = embedding_model
Settings.llm = llm

Connects LlamaIndex global settings to Vertex AI embeddings and Gemini generation.

Persist documents in Firestore

documents = SimpleDirectoryReader(dirname).load_data()
docstore = FirestoreDocumentStore.from_database(
    project=PROJECT_ID,
    database=FIRESTORE_DATABASE_ID,
)
docstore.add_documents(documents)
storage_context = StorageContext.from_defaults(docstore=docstore)
index = VectorStoreIndex.from_documents(
    documents=documents,
    storage_context=storage_context,
)

Loads local files, stores document nodes in Firestore, and builds the retrieval index.

Define workflow steps

class RAGWorkflow(Workflow):
    @step
    async def ingest_data(self, ctx: Context, ev: StartEvent):
        dirname = ev.get("dirname")
        if not dirname:
            return None
        documents = SimpleDirectoryReader(dirname).load_data()
        await ctx.set("documents", documents)
        return FirestoreIndexData(status="First step complete.")

Uses LlamaIndex events and @step methods to route ingestion and query actions.

Run multi-step query

result = await w.run(
    query="What is the significance of the green light?",
    index=index,
    num_steps=NUM_STEPS,
)
display(Markdown(f"{result}"))

Executes the RAG branch with a query, index, and bounded number of generated sub-questions.

Models & APIs used

  • Models: text-embedding-005, gemini-2.5-flash
  • APIs / services: Vertex AI, Firestore
  • SDKs / libraries: vertexai, llama-index, llama-index-embeddings-vertex, llama-index-llms-vertex, llama-index-storage-docstore-firestore, llama-index-utils-workflow, google.auth

When to use this

Use this pattern when you need a structured RAG workflow with ingestion, multi-step querying, reranking, citations, and Firestore-backed storage.

Gotchas & caveats

  • The notebook requires an initialized Google Cloud project, Vertex AI API enabled, an existing VPC/Subnet, and an existing Firestore database.
  • The runtime must be restarted after installing the pinned LlamaIndex packages.
  • Colab authentication is handled separately with google.colab.auth.authenticate_user().
  • PROJECT_ID defaults to GOOGLE_CLOUD_PROJECT only if the placeholder is unchanged.
  • LOCATION defaults to us-central1 when GOOGLE_CLOUD_REGION is unset.
  • The query branch prints a warning and returns None if the index is empty.
  • The Firestore database ID must be supplied in FIRESTORE_DATABASE_ID.

Best practices

  • Pin LlamaIndex package versions used by the workflow.
  • Use google.auth.default with quota_project_id and refresh credentials before model setup.
  • Configure safety settings for dangerous content, harassment, and sexually explicit content.
  • Set Settings.embed_model and Settings.llm so LlamaIndex components share the same Vertex models.
  • Use custom Event classes to make workflow transitions explicit.
  • Limit multi-step query expansion with num_steps.
  • Use citation prompt templates that require source citations and restrict answers to provided sources.
  • Provide a cleanup section for deleting project resources or individual resources.

Graph View

Backlinks