{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ur8xi4C7S06n"
   },
   "outputs": [],
   "source": [
    "# Copyright 2025 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": [
    "# Get started with Vertex AI Memory Bank - ADK\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagent-engine%2Fmemory%2Fget_started_with_memory_bank_adk.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
    "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "</table>\n",
    "\n",
    "<div style=\"clear: both;\"></div>\n",
    "\n",
    "<b>Share to:</b>\n",
    "\n",
    "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84f0f73a0f76"
   },
   "source": [
    "| Author(s) |\n",
    "| --- |\n",
    "| [ Kimberly Milam, Ivan Nardini](https://github.com/inardini)|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This tutorial demonstrates how to build an agent with long-term memory using the Google Agent Development Kit (ADK) and Vertex AI Memory Bank.\n",
    "\n",
    "Vertex AI Memory Bank enables you to dynamically generate and store long-term memories from user conversations. This allows an agent to access personalized information across multiple sessions for a particular user, leading to more contextual and continuous interactions.\n",
    "\n",
    "### What you'll learn\n",
    "\n",
    "By the end of this tutorial, you will be able to:\n",
    "* Understand how long-term memory enhances agent capabilities.\n",
    "* Create a Vertex AI Agent Engine instance to use Memory Bank.\n",
    "* Build an ADK agent that uses the `PreloadMemoryTool` tool to retrieve information.\n",
    "* See how the agent stores conversation history and recalls relevant facts from past sessions.\n",
    "* Build more sophisticated, context-aware Agents.\n",
    "\n",
    "### Why this is important\n",
    "\n",
    "Traditional LLM-based agents often lack the ability to recall information from previous interactions, treating each conversation as a new one. This \"amnesia\" prevents them from maintaining context over time. Long-term memory, like that provided by Vertex AI Memory Bank, allows agents to:\n",
    "\n",
    "* Maintain context over extended periods.\n",
    "* Personalize interactions based on user history.\n",
    "* Build a deeper understanding of user preferences and needs.\n",
    "\n",
    "Google's ADK provides a robust framework for building agents, and by extending it with services like Memory Bank, you can unlock a new level of sophistication.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61RBz8LLbxCR"
   },
   "source": [
    "## Get started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "No17Cw5hgx12"
   },
   "source": [
    "### Install Google Gen AI SDK and other required packages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFy3H3aPgx12"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet \"google-cloud-aiplatform>=1.100.0\" \"google-adk>=1.5.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dmWOrTJ3gx13"
   },
   "source": [
    "### Authenticate your notebook environment (Colab only)\n",
    "\n",
    "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
   ]
  },
  {
   "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": [
    "# Use the environment variable if the user doesn't provide Project ID.\n",
    "import os\n",
    "import uuid\n",
    "\n",
    "import vertexai\n",
    "\n",
    "# Project configuration\n",
    "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
    "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
    "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
    "    if not PROJECT_ID:\n",
    "        raise ValueError(\"Project ID not found. Please set it in the form above or as the GOOGLE_CLOUD_PROJECT environment variable.\")\n",
    "\n",
    "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
    "\n",
    "# Set environment variables required for ADK\n",
    "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"TRUE\"\n",
    "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
    "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
    "\n",
    "# Agent configuration\n",
    "MODEL_NAME = \"gemini-2.5-flash\"\n",
    "USER_ID = f\"user_{uuid.uuid4()}\"\n",
    "\n",
    "print(f\"Project: {PROJECT_ID}\")\n",
    "print(f\"Location: {LOCATION}\")\n",
    "print(f\"Session User ID: {USER_ID}\")\n",
    "\n",
    "# Initialize Vertex AI client\n",
    "_ = vertexai.Client(\n",
    "    project=PROJECT_ID,\n",
    "    location=LOCATION,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5303c05f7aa6"
   },
   "source": [
    "### Import libraries\n",
    "\n",
    "Import the Python libraries required for this tutorial, including components from the ADK, Vertex AI SDK, and standard libraries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fc324893334"
   },
   "outputs": [],
   "source": [
    "import uuid\n",
    "import vertexai\n",
    "from vertexai import agent_engines\n",
    "from google import adk\n",
    "from google.adk.memory import VertexAiMemoryBankService\n",
    "from google.adk.sessions import VertexAiSessionService\n",
    "from google.genai import types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qwEDAkpMqyR4"
   },
   "source": [
    "### Helper functions\n",
    "\n",
    "We will define a helper function to simplify running the agent and capturing its response. This will make our interaction loop cleaner and easier to read.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PlnHiogXq0JP"
   },
   "outputs": [],
   "source": [
    "def run_single_turn(query, session, user_id):\n",
    "    \"\"\"Run a single conversation turn.\"\"\"\n",
    "    content = types.Content(role=\"user\", parts=[types.Part(text=query)])\n",
    "    events = runner.run(user_id=user_id, session_id=session, new_message=content)\n",
    "\n",
    "    response_content = None\n",
    "    for event in events:\n",
    "        if event.is_final_response():\n",
    "            response_content = event.content.parts[0].text\n",
    "    return response_content\n",
    "\n",
    "\n",
    "def chat_loop(session, user_id) -> None:\n",
    "    \"\"\"Main chat interface loop.\"\"\"\n",
    "    print(\"\\nStarting chat. Type 'exit' or 'quit' to end.\")\n",
    "    print(\"Every message will be automatically stored in memory.\\n\")\n",
    "\n",
    "    while True:\n",
    "        user_input = input(\"\\nYou: \")\n",
    "        if user_input.lower() in [\"quit\", \"exit\", \"bye\"]:\n",
    "            print(\"\\nAssistant: Thank you for chatting. Have a great day!\")\n",
    "            break\n",
    "\n",
    "        response = run_single_turn(user_input, session, user_id)\n",
    "        if response:\n",
    "            print(f\"\\nAssistant: {response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EdvJRUWRNGHE"
   },
   "source": [
    "## Creating the Agent Engine with Memory Bank\n",
    "\n",
    "Vertex AI Memory Bank is a component of the Vertex AI Agent Engine, a managed service that allows developers to deploy, manage, and scale AI agents.\n",
    "\n",
    "To use Memory Bank, you first need to create or get an existing Agent Engine instance. This provides the necessary APIs to store and retrieve memories associated with specific users.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x5tUDbo8YVhe"
   },
   "outputs": [],
   "source": [
    "agent_engine = agent_engines.create()\n",
    "print(f\"Created Agent Engine: {agent_engine.resource_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NXtRpE0KXvY6"
   },
   "source": [
    "## Building an Agent with Memory Bank using ADK\n",
    "\n",
    "Now we will construct the core components of our memory-enabled agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q0hfdld4X9yg"
   },
   "source": [
    "### Define the Agent\n",
    "\n",
    "Define the ADK agent with a name, instructions (prompt), and specify the `load_memory` tool. This built-in ADK tool enables the agent to call our `VertexAiMemoryBankService` to search for information. The agent's prompt is crucial, as it instructs the LLM when and how to use this tool to recall user-specific context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qz6Zmy_L6HwY"
   },
   "outputs": [],
   "source": [
    "agent = adk.Agent(\n",
    "    model=MODEL_NAME,\n",
    "    name=\"helpful_assistant\",\n",
    "    instruction=\"\"\"You are a helpful assistant with perfect memory.\n",
    "        Instructions:\n",
    "        - Use the context to personalize responses\n",
    "        - Naturally reference past conversations when relevant\n",
    "        - Build upon previous knowledge about the user\n",
    "        - If using semantic search, the memories shown are the most relevant to the current query\"\"\",\n",
    "    tools=[adk.tools.preload_memory_tool.PreloadMemoryTool()],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TnmlrAZkYJZ2"
   },
   "source": [
    "### Configure the ADK Runner\n",
    "\n",
    "The `Runner` orchestrates the interaction between the user, the agent, and the various services. We will configure it with our `agent` and the `VertexAiSessionService` and `VertexAiMemoryBankService` to manage session state and long-term memory.\n",
    "\n",
    "> Notice that the `VertexAiMemoryBankService` released in ADK, you can also use a local session service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wPEy9Sgp7YLV"
   },
   "outputs": [],
   "source": [
    "app_name = \"my_agent_\" + str(uuid.uuid4())[:6]\n",
    "agent_engine_id = agent_engine.name\n",
    "\n",
    "memory_bank_service = VertexAiMemoryBankService(\n",
    "    project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id\n",
    ")\n",
    "\n",
    "\n",
    "session_service = VertexAiSessionService(\n",
    "    project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id\n",
    ")\n",
    "\n",
    "runner = adk.Runner(\n",
    "    agent=agent,\n",
    "    app_name=app_name,\n",
    "    session_service=session_service,\n",
    "    memory_service=memory_bank_service,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PMfVSv9EYQu_"
   },
   "source": [
    "## Interacting with the Agent - Information gathering session\n",
    "\n",
    "Let's begin our first conversation. In this session, we will provide the agent with specific pieces of information that we expect it to remember later.\n",
    "\n",
    "The process is as follows:\n",
    "\n",
    "1.  Create a new ADK session for the user.\n",
    "2.  Send a series of messages to the agent.\n",
    "3.  After the conversation, retrieve the completed session data.\n",
    "4.  Explicitly add the session to our Memory Bank, which triggers the memory generation process.\n",
    "\n",
    "In a production application, adding the session to memory might be handled by a background process or through hooks in a custom `Runner` or `SessionService`. For this tutorial, we perform this step explicitly to clearly illustrate the process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tRqZ0Rz2YdXL"
   },
   "outputs": [],
   "source": [
    "session1 = await runner.session_service.create_session(\n",
    "    app_name=app_name,\n",
    "    user_id=USER_ID,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c4WVP3nuaIV-"
   },
   "source": [
    "Try these example phrases:\n",
    "\n",
    "```\n",
    "You: Hi, I work as an agent engineer\n",
    "You: I love hiking and have a dog named Max\n",
    "You: I'm working on a recommendation system project\n",
    "You: Bye\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-tqSiFUzaKBC"
   },
   "outputs": [],
   "source": [
    "chat_loop(session1.id, USER_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bzfm-l9caJkf"
   },
   "source": [
    "Add the session to memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "55_1rmY2QufY"
   },
   "outputs": [],
   "source": [
    "completed_session = await runner.session_service.get_session(\n",
    "    app_name=app_name, user_id=USER_ID, session_id=session1.id\n",
    ")\n",
    "await memory_bank_service.add_session_to_memory(completed_session)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gF0AX5rzZYXd"
   },
   "source": [
    "## Interacting with the Agent - Memory Recall session\n",
    "\n",
    "Now, we will start a new session with the same user. This time, we will ask questions that require the agent to recall information from the first session. This is where the `load_memory` tool and our `VertexAIMemoryBankService.search_memory()` method will be invoked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vzBwJL8JZlwa"
   },
   "outputs": [],
   "source": [
    "session2 = await runner.session_service.create_session(\n",
    "    app_name=agent_engine.name,\n",
    "    user_id=USER_ID,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rzdQxPYuZuZb"
   },
   "source": [
    "Try these example phrases:\n",
    "\n",
    "```\n",
    "You: What do you remember about me?\n",
    "You: What is my dog's name?\n",
    "You: Bye\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rbbqJEmDZqvv"
   },
   "outputs": [],
   "source": [
    "chat_loop(session2.id, USER_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c7jVCYtvaxbD"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To avoid charges, delete the agent engine when you're done experimenting.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CAtT65gEcAFB"
   },
   "outputs": [],
   "source": [
    "delete_engine = True\n",
    "\n",
    "if delete_engine:\n",
    "    agent_engines.delete(resource_name=agent_engine.resource_name, force=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "get_started_with_memory_bank_adk.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "workbench-notebooks.m130",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
