Generative AI Google Brain

gemini-agent-engine-tutorial-langgraph-rag-agent

4 min read

Building a Multi-Agent RAG Application with LangGraph and Agent Engine

Source notebook

Repo path: gemini/agent-engine/tutorial_langgraph_rag_agent.ipynb · Open on GitHub · advanced

Builds and deploys a LangGraph multi-agent RAG app on Vertex AI Agent Engine with Cloud SQL vector stores.

Summary

This notebook teaches how to combine LangGraph, LangChain tools, Cloud SQL for PostgreSQL vector stores, and Vertex AI Agent Engine into a deployable RAG application. It loads Harry Potter movie and book documents from Cloud Storage, embeds them with text-embedding-005, routes questions to book or movie retrievers, summarizes results with gemini-2.0-flash, tests locally, deploys remotely, and cleans up resources.

Key code patterns

Initialize Vertex AI

PROJECT_ID = "[your-project-id]"
LOCATION = "us-central1"
STAGING_BUCKET = "gs://[your-staging-bucket]"
 
vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)

Sets the project, region, and staging bucket required before creating Agent Engine resources.

Create vector store tables

engine = await PostgresEngine.afrom_instance(
    PROJECT_ID, REGION, INSTANCE, DATABASE, user="postgres", password=PASSWORD
)
await engine.ainit_vectorstore_table(
    table_name=table_name,
    vector_size=768,
)

Creates Cloud SQL PostgreSQL tables sized for text-embedding-005 embeddings.

Load documents from Cloud Storage

storage_client = storage.Client()
bucket = storage_client.bucket(bucket_name)
blob = bucket.blob(f"{gcs_dir}/{blob_name}")
with blob.open("r") as f:
    json_docs = json.loads(f.read())
docs = [Document(**doc["kwargs"]) for doc in json_docs]

Fetches JSON document payloads from GCS and converts them into LangChain Document objects.

Add embeddings to Postgres

vector_store = await PostgresVectorStore.create(
    engine,
    table_name=MOVIE_TABLE_NAME,
    embedding_service=VertexAIEmbeddings(model_name="text-embedding-005", project=PROJECT_ID),
)
ids = [str(uuid.uuid4()) for i in range(len(docs))]
await vector_store.aadd_documents(docs, ids=ids)

Embeds source documents and stores them in PostgreSQL-backed vector tables.

Define retriever tool

@tool
def movie_similarity_search(query: str) -> str:
    vector_store = PostgresVectorStore.create_sync(
        engine,
        table_name=MOVIE_TABLE_NAME,
        embedding_service=VertexAIEmbeddings(model_name="text-embedding-005", project=PROJECT_ID),
    )
    return str([doc for doc in vector_store.as_retriever().invoke(query)])

Exposes semantic retrieval as a LangChain tool callable from the graph.

Route graph by question type

def router(state: list[BaseMessage]) -> Literal["book_similarity_search", "movie_similarity_search", "__end__"]:
    if not state[0].content or len(state[1].tool_calls) == 0:
        return "__end__"
    if "book" in state[0].content:
        return "book_similarity_search"
    if "movie" in state[0].content:
        return "movie_similarity_search"
    return "__end__"

Controls whether the graph retrieves from the book table, movie table, or stops.

Compile LangGraph app

model = ChatVertexAI(model="gemini-2.0-flash")
builder = MessageGraph()
builder.add_node("checker", checker)
builder.set_entry_point("checker")
builder.add_node("tools", model.bind_tools([book_similarity_search, movie_similarity_search]))
builder.add_conditional_edges("tools", router)
self.runnable = builder.compile()

Builds a stateful multi-stage workflow around Gemini, tools, routing, retrieval, and summarization.

Deploy to Agent Engine

remote_app = agent_engines.create(
    MultiStageLangGraphApp(project=PROJECT_ID, location=LOCATION),
    requirements=["google-cloud-aiplatform[agent_engines,langchain]==1.60.0", "langgraph==0.0.51"],
    display_name="Agent Engine with LangGraph RAG Agent",
)

Packages the custom LangGraph application and required dependencies as a managed Agent Engine app.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, Agent Engine, Cloud SQL for PostgreSQL, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langchain, langchain-google-cloud-sql-pg, langchain-google-vertexai, langgraph, google-cloud-storage

When to use this

Use this pattern when you need a deployable multi-agent RAG workflow that routes queries across multiple PostgreSQL vector stores on Vertex AI Agent Engine.

Gotchas & caveats

  • Requires an existing Google Cloud project with billing enabled.
  • The notebook asks for Owner IAM role for completing the tutorial.
  • Vertex AI API must be enabled before initializing and deploying.
  • The notebook states only us-central1 is supported at the time.
  • Cloud SQL for PostgreSQL setup is assumed from a previous tutorial.
  • The runtime must be restarted after installing pinned package versions.
  • The PostgreSQL password is collected with input and reused by retriever tools.
  • The vector table uses vector_size=768 for text-embedding-005.

Best practices

  • Pin deployment requirements to the same package versions used in the notebook.
  • Test the LangGraph app locally before deploying it to Agent Engine.
  • Use separate vector store tables for movie and book data.
  • Use Cloud Storage as the source for reusable JSON document datasets.
  • Clean up Agent Engine apps and Cloud SQL instances after the tutorial to avoid billing.

Graph View

Backlinks