Get started with Vertex AI Memory Bank - LangGraph

Source notebook

Repo path: gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.ipynb · Open on GitHub · intermediate

Builds a LangGraph chatbot with Vertex AI Memory Bank for long-term personalized memory.

Summary

This notebook teaches how to provision a Vertex AI Agent Engine and use its Memory Bank API with a LangGraph conversational agent. The workflow creates an Agent Engine, optionally seeds user memory, defines LangGraph state, retrieves relevant memories, injects them into a system prompt, calls Gemini through ChatVertexAI, stores each turn back to Memory Bank, inspects stored memories, and deletes the Agent Engine.

Key code patterns

Initialize Vertex AI client

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

Sets project and region before creating Agent Engine and Memory Bank resources.

Create Agent Engine

agent_engine = client.agent_engines.create()
print(agent_engine.api_resource.name)

Provisions the managed Agent Engine resource used by Memory Bank.

Seed user memory

client.agent_engines.create_memory(
    name=agent_engine.api_resource.name,
    fact=f"The user's name is {USER_ID} and they have a passion for learning new things.",
    scope={"user_id": USER_ID},
)

Shows how to pre-load user-specific facts before the first chat turn.

Retrieve relevant memories

memories = client.agent_engines.retrieve_memories(
    name=agent_engine_name,
    scope={"user_id": user_id},
    similarity_search_params={
        "search_query": user_message,
        "top_k": 10,
    },
)

Uses semantic search to fetch memories relevant to the current user message.

Inject memory into prompt

system_message = SystemMessage(
    content=f"""You are a helpful assistant with perfect memory.
 
    {memory_context}
 
    Instructions:
    - Use the context to personalize responses"""
)

Provides retrieved facts as system context for personalized Gemini responses.

Call Gemini with LangChain

messages = [system_message] + state["messages"]
llm = ChatVertexAI(model=MODEL_NAME, project=PROJECT_ID, location=LOCATION)
response = llm.invoke(messages)

Uses ChatVertexAI to generate a response from the memory-augmented conversation state.

Store each turn

client.agent_engines.generate_memories(
    name=agent_engine_name,
    direct_contents_source={"events": events},
    scope={"user_id": user_id},
    config={"wait_for_completion": True},
)

Persists user and model messages so later turns can use long-term memory.

Compile LangGraph flow

graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot_with_memory)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
compiled_graph = graph_builder.compile()

Defines a single-node graph that processes each conversation turn.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Vertex AI Agent Engine, Vertex AI Memory Bank
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langgraph, langchain-core, langchain-google-vertexai

When to use this

Use this pattern when building LangGraph agents that need user-scoped long-term memory across chat sessions.

Gotchas & caveats

  • The notebook requires an existing Google Cloud project with the Vertex AI API enabled.
  • Colab users may need to run google.colab auth before using Google Cloud resources.
  • PROJECT_ID must be set directly or through GOOGLE_CLOUD_PROJECT.
  • LOCATION defaults to us-central1 from GOOGLE_CLOUD_REGION when not provided.
  • The package requirement includes google-cloud-aiplatform>=1.100.0.
  • User messages must be strings or the chatbot node raises ValueError.
  • The notebook deletes the Agent Engine only when delete_engine is True to avoid charges.

Best practices

  • Scope memories by user_id so retrieved facts are user-specific.
  • Retrieve memories before generation and inject them into the system prompt.
  • Use semantic search with top_k to limit retrieved memories to relevant facts.
  • Store both user and model messages after each turn to continuously build memory.
  • Use wait_for_completion when generating memories so persistence completes before continuing.
  • Delete the Agent Engine after experiments to avoid charges.