{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "ur8xi4C7S06n" }, "outputs": [], "source": [ "# Copyright 2026 Google LLC\n", "#\n", "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "JAPoU8Sm5E6e" }, "source": [ "# Persisting LangChain History with Vertex AI Session Service\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Vertex AI Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
\n", "\n", "

\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", "\n", "

" ] }, { "cell_type": "markdown", "metadata": { "id": "tvgnzT1CKxrO" }, "source": [ "## Overview\n", "\n", "This notebook will demonstrate how to use the Vertex AI Session Service to persist the conversational history with LangChain agents.\n", "\n", "You will lean how to:\n", "- Create a Session with the Vertex AI Session Service\n", "- Store conversation turns within the sessions you created\n", "- Store tool calls and tool results within the sessions you created\n", "- Access the stored conversational history" ] }, { "cell_type": "markdown", "metadata": { "id": "61RBz8LLbxCR" }, "source": [ "## Get started" ] }, { "cell_type": "markdown", "metadata": { "id": "No17Cw5hgx12" }, "source": [ "### Install Vertex AI SDK and LangChain for Google\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tFy3H3aPgx12" }, "outputs": [], "source": [ "%pip install \"google-cloud-aiplatform[agent_engines]\" --force-reinstall --quiet\n", "# Install the Google Generative AI integration for LangChain\n", "%pip install -U \"langchain-google-genai\" --quiet" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "r-ES3MMrMlxV" }, "outputs": [], "source": [ "# Restart the session to ensure packages are updated\n", "import os\n", "\n", "os._exit(0)" ] }, { "cell_type": "markdown", "metadata": { "id": "dmWOrTJ3gx13" }, "source": [ "### Authenticate your notebook environment\n", "\n", "If you are running this notebook in **Google Colab**, run the cell below to authenticate your account." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "NyKGtVQjgx13" }, "outputs": [], "source": [ "import sys\n", "\n", "if \"google.colab\" in sys.modules:\n", " from google.colab import auth\n", "\n", " auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": { "id": "DF4l8DTdWgPY" }, "source": [ "### Set Google Cloud project information\n", "\n", "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", "\n", "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Nqwi-5ufWp_B" }, "outputs": [], "source": [ "import os\n", "\n", "# fmt: off\n", "PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n", "LOCATION = \"us-central1\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n", "# fmt: on\n", "\n", "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n", " PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n", "if not LOCATION:\n", " LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\")\n", "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"1\"" ] }, { "cell_type": "markdown", "metadata": { "id": "5303c05f7aa6" }, "source": [ "### Import libraries" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6fc324893334" }, "outputs": [], "source": [ "import datetime\n", "\n", "import requests\n", "from langchain_core.chat_history import BaseChatMessageHistory\n", "from langchain_core.messages import (\n", " BaseMessage,\n", " HumanMessage,\n", " message_to_dict,\n", " messages_from_dict,\n", ")\n", "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n", "from langchain_core.runnables.history import RunnableWithMessageHistory\n", "from langchain_core.tools import tool\n", "from langchain_google_genai import ChatGoogleGenerativeAI\n", "from vertexai import Client" ] }, { "cell_type": "markdown", "metadata": { "id": "EdvJRUWRNGHE" }, "source": [ "## Create Vertex AI Chat Message History\n" ] }, { "cell_type": "markdown", "metadata": { "id": "ydzH42O6O-U5" }, "source": [ "LangChain uses `BaseChatMessageHistory` to persist the chat history of a LangChain agent. It has 3 basic methods:\n", "- messages: a property of the history, listing previous message history\n", "- add_messages: add one or multiple messages to the conversation history\n", "- clear: delete all messages in the conversation history\n", "\n", "We will create the VertexAISessionChatMessageHistory that extends BaseChatMessageHistory, which will use the Vertex AI Session service to implement the messages and add_messages methods." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "m814rnX7O7D5" }, "outputs": [], "source": [ "class VertexAISessionChatMessageHistory(BaseChatMessageHistory):\n", " \"\"\"A LangChain message history backend that defers storage to\n", " Vertex AI Agent Engine's Session Service via the vertexai Client.\n", " \"\"\"\n", "\n", " def __init__(self, client: Client, session_name: str, author: str = \"user\"):\n", " self.client = client\n", " self.session_name = session_name\n", " self.author = author\n", "\n", " @property\n", " def messages(self) -> list[BaseMessage]:\n", " # Call ListEvents from the Vertex AI Session service to return the list of messages.\n", " events = self.client.agent_engines.sessions.events.list(name=self.session_name)\n", "\n", " lc_messages = []\n", " for event in events:\n", " # Convert directly from the LangChain message stored in raw_event\n", " if event.raw_event and \"langchain_message\" in event.raw_event:\n", " # `messages_from_dict` expects a list of serialized messages\n", " extracted_message = messages_from_dict(\n", " [event.raw_event[\"langchain_message\"]]\n", " )[0]\n", " lc_messages.append(extracted_message)\n", " continue\n", " return lc_messages\n", "\n", " def add_message(self, message: BaseMessage) -> None:\n", " is_human = isinstance(message, HumanMessage)\n", " author = self.author if is_human else \"agent\"\n", "\n", " # Dump the entire LangChain message structure (id, tool_calls, kwargs, response_metadata)\n", " # into a JSON serializable dict\n", " serialized_lc_message = message_to_dict(message)\n", "\n", " self.client.agent_engines.sessions.events.append(\n", " name=self.session_name,\n", " author=author,\n", " invocation_id=\"langchain-invocation\",\n", " timestamp=datetime.datetime.now(datetime.timezone.utc),\n", " config={\n", " # Store the full LangChain payload in the raw_event field\n", " \"raw_event\": {\"langchain_message\": serialized_lc_message}\n", " },\n", " )\n", "\n", " def clear(self) -> None:\n", " \"\"\"Clears the session. Vertex AI Session Service currently doesn't support\n", " easily clearing individual events; typical workarounds include\n", " deleting the session natively via `self.client.agent_engines.sessions.delete(...)`.\n", " \"\"\"\n", " raise NotImplementedError(\n", " \"Clear is not natively supported without deleting the entire session via SDK.\"\n", " )" ] }, { "cell_type": "markdown", "metadata": { "id": "FRl1IPZPSXZi" }, "source": [ "## Initialize Vertex AI\n", "\n", "We will use the Vertex AI SDK to create a session instance.\n", "\n", "Vertex AI Sessions are sub resources of Agent Engines, which is the managed solution for deploying agents to Google Cloud. For more details, see [link]. This means all sessions must be associated with an Agent Engine.\n", "\n", "For this demo, we don't need to deploy the agent. We will just use agent_engines.create to create an empty Agent Engine to store our sessions." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "GhiU7PzERpOE" }, "outputs": [], "source": [ "import vertexai\n", "\n", "# Initialize the Vertex AI SDK Client with your project and location\n", "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)\n", "\n", "# Create an agent engine to hold the session\n", "agent_engine = client.agent_engines.create(config={\"display_name\": \"langchain_test\"})\n", "agent_engine_path = agent_engine.api_resource.name" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mvZ4-TjoRrAh" }, "outputs": [], "source": [ "# Create a Vertex AI session resource\n", "session_operation = client.agent_engines.sessions.create(\n", " name=agent_engine_path, user_id=\"langchain_user\"\n", ")\n", "\n", "# Extract the full session resource name from the finished operation response\n", "session_resource_name = session_operation.response.name" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "jfT0Zt0pRsce" }, "outputs": [], "source": [ "# Define a factory function that LangChain will call per session ID\n", "def get_vertex_session_history(session_id: str) -> BaseChatMessageHistory:\n", " return VertexAISessionChatMessageHistory(client=client, session_name=session_id)" ] }, { "cell_type": "markdown", "metadata": { "id": "JIhMzITHTNia" }, "source": [ "## Start a conversation\n", "Since you already signed into your Google Cloud project, we will use Gemini as the model for our agent.\n", "\n", "Here, we will just create a simple conversational agent. When we send messages to the LangChain agent, the messages will be stored in our Session using the `add_messages` method we defined in `VertexAISessionChatMessageHistory`" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ydyerW-3RuXt" }, "outputs": [], "source": [ "# Create a standard LangChain Model and Prompt\n", "model = ChatGoogleGenerativeAI(\n", " model=\"gemini-2.5-flash\", project=PROJECT_ID, location=LOCATION\n", ")\n", "\n", "prompt = ChatPromptTemplate.from_messages(\n", " [\n", " (\"system\", \"You are a helpful assistant.\"),\n", " MessagesPlaceholder(variable_name=\"history\"),\n", " (\"human\", \"{input}\"),\n", " ]\n", ")\n", "\n", "chain = prompt | model\n", "\n", "# Wrap the chain with RunnableWithMessageHistory\n", "chain_with_history = RunnableWithMessageHistory(\n", " chain,\n", " get_vertex_session_history,\n", " input_messages_key=\"input\",\n", " history_messages_key=\"history\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "JWVTDPUXV9Pq" }, "outputs": [], "source": [ "# Invoke the chain using the session resource name we generated\n", "input = \"Hi! My name is John Doe.\"\n", "print(input)\n", "response = chain_with_history.invoke(\n", " {\"input\": input}, config={\"configurable\": {\"session_id\": session_resource_name}}\n", ")\n", "print(\"Gemini: \", response.content)\n", "\n", "input = \"What is my name?\"\n", "print(input)\n", "response = chain_with_history.invoke(\n", " {\"input\": input}, config={\"configurable\": {\"session_id\": session_resource_name}}\n", ")\n", "print(\"Gemini: \", response.content)" ] }, { "cell_type": "markdown", "metadata": { "id": "3XQbxqL-T4-4" }, "source": [ "We can view the entire chat history using the `messages` property we defined in `VertexAISessionChatMessageHistory`" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "iLK18ollRv-r" }, "outputs": [], "source": [ "# View the conversation history\n", "history = get_vertex_session_history(session_resource_name)\n", "for message in history.messages:\n", " print(f\"{message.type.capitalize()}: {message.content}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "tZKASjWbSTaJ" }, "source": [ "## Tool Call\n", "\n", "Now that we have used a simple agent, we will show how the Vertex AI Session service can store more advanced conversations from agents with Tools.\n", "\n", "In this example, we will use a tool called `get_weather` to allow the agent to check the weather in any city we ask." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "gBB4xZoORx14" }, "outputs": [], "source": [ "# Create a new Vertex AI session resource\n", "session_operation = client.agent_engines.sessions.create(\n", " name=agent_engine_path, user_id=\"langchain_user\"\n", ")\n", "\n", "# Extract the full session resource name from the finished operation response\n", "session_resource_name = session_operation.response.name" ] }, { "cell_type": "markdown", "metadata": { "id": "xjeed8KhVcuz" }, "source": [ "## Define the get_weather tool\n", "We will use a basic weather API for our tool. This allows the agent to check the weather for a given city." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "EdSStEwTU1Tq" }, "outputs": [], "source": [ "# Define the tool\n", "@tool\n", "def get_weather(location: str) -> str:\n", " \"\"\"Returns the current weather for a given city name.\"\"\"\n", " # format=3 gives a simple string like: \"London: ⛅️ +15°C\"\n", " try:\n", " response = requests.get(f\"https://wttr.in/{location}?format=3\")\n", " return response.text.strip()\n", " except Exception:\n", " return \"Sorry, I couldn't fetch the weather right now.\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Ro3tWwbARzLC" }, "outputs": [], "source": [ "# Bind the tool to the model\n", "# ChatVertexAI natively supports LangChain tool integration\n", "model_with_tools = ChatGoogleGenerativeAI(\n", " model=\"gemini-2.5-flash\", project=PROJECT_ID, location=LOCATION\n", ").bind_tools([get_weather])\n", "\n", "# Create a Prompt Template with STRICT system instructions\n", "prompt = ChatPromptTemplate.from_messages(\n", " [\n", " (\n", " \"system\",\n", " \"You are a specialized weather assistant. You MUST use the `get_weather` tool to answer any questions about the weather. Do not use your internal knowledge.\",\n", " ),\n", " MessagesPlaceholder(variable_name=\"history\"),\n", " (\"human\", \"{input}\"),\n", " ]\n", ")\n", "\n", "# Combine the prompt and the model into a standard runnable chain\n", "chain = prompt | model_with_tools\n", "\n", "# Wrap the chain with RunnableWithMessageHistory\n", "chain_with_history = RunnableWithMessageHistory(\n", " chain,\n", " get_vertex_session_history,\n", " input_messages_key=\"input\",\n", " history_messages_key=\"history\",\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "7VkQ6H_VXe5I" }, "source": [ "Now that we registered the tool with the model, lets try a basic conversation.\n", "\n", "In LangChain, the application logic is responsible for calling the tool. Here, if the model tells us to call get_weather, we call the function then provide the response to the model." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "fkMq12VgWGuJ" }, "outputs": [], "source": [ "# --- Turn 1: AI Uses the Tool ---\n", "input = \"What is the weather in Paris?\"\n", "print(input)\n", "response = chain_with_history.invoke(\n", " {\"input\": input}, config={\"configurable\": {\"session_id\": session_resource_name}}\n", ")\n", "\n", "# If the system instruction worked, the model will output a tool call instead of text\n", "print(\"(Gemini Tool Request):\", response.tool_calls)\n", "\n", "# --- Turn 2: Providing the tool result ---\n", "if response.tool_calls:\n", " tool_call = response.tool_calls[0]\n", " # Execute the python function manually\n", " tool_result = get_weather.invoke(tool_call)\n", " print(\"(Tool Response):\", tool_result)\n", " request_payload = {\n", " # Pass the ToolMessage containing the tool_call_id directly into the input\n", " \"input\": [tool_result]\n", " }\n", " final_response = chain_with_history.invoke(\n", " request_payload, config={\"configurable\": {\"session_id\": session_resource_name}}\n", " )\n", " print(\"(Gemini Final Answer):\", final_response.content)" ] }, { "cell_type": "markdown", "metadata": { "id": "WH6c62lvXymu" }, "source": [ "We can view the entire chat history using the `messages` property we defined in `VertexAISessionChatMessageHistory`" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "bRp0xp88R1AS" }, "outputs": [], "source": [ "history = get_vertex_session_history(session_resource_name)\n", "for message in history.messages:\n", " print(f\"{message.type.capitalize()}: {message.content}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "2a4e033321ad" }, "source": [ "## Cleaning up\n", "\n", "Now that the demo is done, we can delete the Agent Engine instance to delete the sessions we created." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "3ADJzvTWX9hB" }, "outputs": [], "source": [ "client.agent_engines.delete(name=agent_engine_path, force=True)" ] } ], "metadata": { "colab": { "name": "langchain_vertex_ai_session_service.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }