Generative AI Google Brain

gemini-agent-engine-langchain-vertex-ai-session-service

3 min read

Persisting LangChain History with Vertex AI Session Service

Source notebook

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

Persists LangChain chat and tool history in Vertex AI Agent Engine Session Service.

Summary

This notebook teaches how to implement a LangChain BaseChatMessageHistory backend using Vertex AI Agent Engine sessions. It creates an Agent Engine, creates session resources, stores serialized LangChain messages as raw_event payloads, retrieves history, and demonstrates both plain conversation and a weather tool call workflow with Gemini.

Key code patterns

Session-backed chat history

class VertexAISessionChatMessageHistory(BaseChatMessageHistory):
    def __init__(self, client, session_name, author="user"):
        self.client = client
        self.session_name = session_name
        self.author = author

Adapts LangChain’s history interface to Vertex AI Session Service storage.

Restore LangChain messages

events = self.client.agent_engines.sessions.events.list(name=self.session_name)
for event in events:
    if event.raw_event and "langchain_message" in event.raw_event:
        msg = messages_from_dict([event.raw_event["langchain_message"]])[0]
        lc_messages.append(msg)

Reads session events and reconstructs full LangChain message objects.

Append serialized messages

serialized = message_to_dict(message)
self.client.agent_engines.sessions.events.append(
    name=self.session_name,
    author=author,
    invocation_id="langchain-invocation",
    timestamp=datetime.datetime.now(datetime.timezone.utc),
    config={"raw_event": {"langchain_message": serialized}},
)

Preserves LangChain message metadata, tool calls, kwargs, and response metadata.

Runnable history wrapper

chain_with_history = RunnableWithMessageHistory(
    chain,
    get_vertex_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

Connects LangChain runtime history handling to Vertex AI session resources.

Manual tool loop

response = chain_with_history.invoke({"input": input}, config={"configurable": {"session_id": session_resource_name}})
if response.tool_calls:
    tool_call = response.tool_calls[0]
    tool_result = get_weather.invoke(tool_call)
    final_response = chain_with_history.invoke({"input": [tool_result]}, config={"configurable": {"session_id": session_resource_name}})

Shows that application code executes requested tools and feeds tool results back to the model.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Vertex AI Agent Engine Session Service
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langchain, langchain-google-genai, requests

When to use this

Use this pattern when a LangChain agent needs persistent, retrievable conversation and tool-call history on Vertex AI.

Gotchas & caveats

  • The Vertex AI API must be enabled for the selected Google Cloud project.
  • PROJECT_ID and LOCATION must be configured, with LOCATION defaulting to us-central1 in the notebook.
  • The notebook sets GOOGLE_GENAI_USE_VERTEXAI=1 before using ChatGoogleGenerativeAI.
  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • Sessions are subresources of Agent Engines, so the demo creates an empty Agent Engine even without deploying an agent.
  • Vertex AI Session Service does not natively support clearing individual events in this implementation.
  • The LangChain application logic is responsible for executing tool calls and returning tool results.

Best practices

  • Store the full LangChain message payload with message_to_dict instead of only message text.
  • Use messages_from_dict to reconstruct LangChain messages from stored raw_event payloads.
  • Wrap chains with RunnableWithMessageHistory to centralize history injection and persistence.
  • Use strict system instructions when the model must call a tool instead of using internal knowledge.
  • Delete the Agent Engine with force=True after the demo to clean up sessions.

Graph View

Backlinks