{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "li-5uLZMOWbm"
      },
      "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": "iY-d5cexcFzH"
      },
      "source": [
        "# Get started with Memory Bank on 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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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%2Fagents%2Fagent_engine%2Fmemory_bank%2Fget_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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/agents/agent_engine/memory_bank/get_started_with_memory_bank_on_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": "OHNf0Rg1cHqf"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [Kimberly Milam](https://github.com/klmilam) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qlXbr9YbMyzc"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to build ADK agents with Memory, including both `InMemoryMemoryService` and `VertexAiMemoryBankService`.\n",
        "\n",
        "This tutorial will cover:\n",
        "* Comparing `InMemoryMemoryService` vs. `VertexAiMemoryBankService`.\n",
        "* Generating memories with ADK and Agent Engine Memory Bank.\n",
        "* Retrieving memories with ADK and Agent Engine Memory Bank.\n",
        "* Customizing your Memory Bank instance's behavior."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSSuQi5wS-AY"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ngEAiXw8cQRt"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZNg0utat73eG"
      },
      "outputs": [],
      "source": [
        "%pip install google-adk google-cloud-aiplatform --upgrade --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbArJaMNcWVz"
      },
      "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": "oMmt34WM-08o"
      },
      "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": "mbeaDydxcdEo"
      },
      "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": "PKZ_ad_g_AuU"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# fmt: off\n",
        "PROJECT = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "if not PROJECT or PROJECT == \"[your-project-id]\":\n",
        "    PROJECT = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "client = vertexai.Client(project=PROJECT, location=LOCATION)\n",
        "\n",
        "# Set environment variables for ADK.\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"TRUE\"\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT\n",
        "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2fe816fc0d0"
      },
      "source": [
        "## (Optional) Set up logging\n",
        "To surface ADK logs in your notebook, you may need to configure the logging level of ADK's logger."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "91955e5b44ec"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "import sys\n",
        "\n",
        "# Set up a general handler for all logs to be printed in your Colab output.\n",
        "# We'll set the overall level to INFO.\n",
        "logging.basicConfig(\n",
        "    level=logging.INFO,\n",
        "    format=\"%(asctime)s - %(name)s - %(levelname)s - %(message)s\",\n",
        "    stream=sys.stdout,\n",
        ")\n",
        "\n",
        "# Get the specific logger used by ADK.\n",
        "adk_logger = logging.getLogger(\"google_adk\")\n",
        "\n",
        "# Set its level to DEBUG to ensure you capture ALL messages from the SDK,\n",
        "# even if the general level is higher (like INFO).\n",
        "adk_logger.setLevel(logging.DEBUG)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W0n_JRrHNUUX"
      },
      "source": [
        "## Define a simple agent\n",
        "\n",
        "The ADK Runner defines what session service is used by the agent. The session service is responsible for saving the conversation history.\n",
        "\n",
        "In this example, we're using the in-memory session service (`InMemorySessionService`). However, you can use any session service with Memory.\n",
        "\n",
        "If you want to use Google's pre-built Memory tools, you can provide a Memory service to the Runner as well. In this example, we'll start by directly invoking memory so that we can compare different services. So, we're not directly setting `memory_service` in the Runner."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rKxFZmSRAECu"
      },
      "outputs": [],
      "source": [
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.genai.types import Content, Part\n",
        "\n",
        "APP_NAME = \"memory_example_app\"\n",
        "MODEL = \"gemini-2.5-flash\"\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        ")\n",
        "\n",
        "session_service = InMemorySessionService()\n",
        "\n",
        "runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service)\n",
        "\n",
        "\n",
        "def call_agent(query, session, user_id):\n",
        "    content = Content(role=\"user\", parts=[Part(text=query)])\n",
        "    events = runner.run(user_id=user_id, session_id=session, new_message=content)\n",
        "\n",
        "    for event in events:\n",
        "        if event.is_final_response():\n",
        "            final_response = event.content.parts[0].text\n",
        "            print(\"Agent Response: \", final_response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WEAg11vnOyYs"
      },
      "source": [
        "## Introduction to ADK Memory Service\n",
        "ADK memory services offer an interface for creating and searching for memories. The [`BaseMemoryService`](https://github.com/google/adk-python/blob/main/src/google/adk/memory/base_memory_service.py) defines the interface for managing this searchable, long-term knowledge store. Its primary responsibilities are:\n",
        "\n",
        "* Ingesting Information (`add_session_to_memory`): Taking the contents of a (usually completed) Session and adding relevant information to the long-term knowledge store.\n",
        "* Searching Information (`search_memory`): Allowing an agent (typically via a Tool) to query the knowledge store and retrieve relevant snippets or context based on a search query.\n",
        "\n",
        "In this colab, we'll cover two options:\n",
        "\n",
        "* `InMemoryMemoryService`: A in-memory service that allows you to search your conversation history to retrieve the raw dialogue using basic keyword matching.\n",
        "* `VertexAiMemoryBankService`: A remote service that remotely generates memories for \"meaningful\" information from user conversation.\n",
        "\n",
        "\n",
        "You can also interact directly with Memory Bank by using the Agent Engine SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSZQQAPQWBXs"
      },
      "source": [
        "### `InMemoryMemoryService`\n",
        "\n",
        "[`InMemoryMemoryService`](https://github.com/google/adk-python/blob/main/src/google/adk/memory/in_memory_memory_service.py) persists **all** information in-memory. The events are not condensed, and information is not extracted from them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HqXCHxi2Cb-n"
      },
      "outputs": [],
      "source": [
        "from google.adk.memory import InMemoryMemoryService\n",
        "\n",
        "in_memory_service = InMemoryMemoryService()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mTuaUZ1DWG7L"
      },
      "source": [
        "### `VertexAiMemoryBankService`\n",
        "Memory Bank will only persist information that's meaningful for future interactions. [`VertexAiMemoryBankService`](https://github.com/google/adk-python/blob/main/src/google/adk/memory/vertex_ai_memory_bank_service.py) provides a built-in integration between Memory Bank and ADK.\n",
        "\n",
        "To get started with Memory Bank, you need to first create an Agent Engine. This only takes a few seconds.\n",
        "\n",
        "If you don't have a Google Cloud project, you can use [Vertex in Express mode, which only requires an API key](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/set-up#authentication)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rZP-0IVaQ5_N"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "from google.adk.memory import VertexAiMemoryBankService\n",
        "\n",
        "client = vertexai.Client(project=PROJECT, location=LOCATION)\n",
        "\n",
        "agent_engine = client.agent_engines.create(\n",
        "    config={\n",
        "        \"context_spec\": {\n",
        "            \"memory_bank_config\": {\n",
        "                \"generation_config\": {\n",
        "                    \"model\": f\"projects/{PROJECT}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    }\n",
        ")\n",
        "\n",
        "memory_bank_service = VertexAiMemoryBankService(\n",
        "    agent_engine_id=agent_engine.api_resource.name.split(\"/\")[-1],\n",
        "    project=PROJECT,\n",
        "    location=LOCATION,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-IrevumRepq"
      },
      "source": [
        "## Generating Memories\n",
        "\n",
        "Memories will be generated from your conversation history stored in the `Session` object. The ADK `add_session_to_memory` method is responsible for triggering memory generation for your memory service, but the exact behavior depends on which service you use:\n",
        "\n",
        "* `InMemoryMemoryService` persists all information, including the raw dialogue.\n",
        "\n",
        "* `VertexAiMemoryBankService` only persists *meaningful* information. Not all conversations will result in generated memories. Additionally, the information will be condensed to individual, self-contained memories.\n",
        "\n",
        "In this example, we create two sessions:\n",
        "1. `chit_chat_session`, which is a conversation that contains events that may not be meaningful for future interactions.\n",
        "2. `recommendations_session`, which is a conversation about looking for birthday presents for a family member. It includes implicit information that may be helpful for future conversations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VfTIuJfFAR2z"
      },
      "outputs": [],
      "source": [
        "USER_ID = \"My User\"\n",
        "\n",
        "chit_chat_session = await session_service.create_session(\n",
        "    app_name=APP_NAME, user_id=USER_ID\n",
        ")\n",
        "\n",
        "call_agent(\n",
        "    \"what types of questions do you answer\", chit_chat_session.id, user_id=USER_ID\n",
        ")\n",
        "# Refresh session object.\n",
        "chit_chat_session = await session_service.get_session(\n",
        "    app_name=APP_NAME, user_id=USER_ID, session_id=chit_chat_session.id\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hLvyyXRHEEQH"
      },
      "outputs": [],
      "source": [
        "recommendations_session = await session_service.create_session(\n",
        "    app_name=APP_NAME, user_id=USER_ID\n",
        ")\n",
        "\n",
        "call_agent(\n",
        "    \"I have a three year old niece. What recommendations do you have for birthday presents?\",\n",
        "    recommendations_session.id,\n",
        "    user_id=USER_ID,\n",
        ")\n",
        "call_agent(\n",
        "    \"I like the idea of a bike. Are there any age appropriate options?\",\n",
        "    recommendations_session.id,\n",
        "    user_id=USER_ID,\n",
        ")\n",
        "\n",
        "recommendations_session = await session_service.get_session(\n",
        "    app_name=APP_NAME, user_id=USER_ID, session_id=recommendations_session.id\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0-bWy4LlTjio"
      },
      "source": [
        "### InMemoryMemoryService\n",
        "\n",
        "First, let's see what memories are extracted from the conversation using `InMemoryMemoryService`\n",
        "\n",
        "The memories include the raw dialogue of the conversation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D6PyuoF7_dz2"
      },
      "outputs": [],
      "source": [
        "await in_memory_service.add_session_to_memory(chit_chat_session)\n",
        "await in_memory_service.add_session_to_memory(recommendations_session)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IP1_NjgjTzO0"
      },
      "outputs": [],
      "source": [
        "# Peek at the memories.\n",
        "in_memory_service._session_events"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZY1cRqfi_4Qz"
      },
      "outputs": [],
      "source": [
        "await in_memory_service.search_memory(app_name=APP_NAME, user_id=USER_ID, query=\"test\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wV2OziabBVt5"
      },
      "outputs": [],
      "source": [
        "await in_memory_service.search_memory(\n",
        "    app_name=APP_NAME,\n",
        "    user_id=USER_ID,\n",
        "    query=\"what should I get my mom for mother's day\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j8d9R6NMUgyw"
      },
      "source": [
        "### Memory Bank via `VertexAiMemoryBankService`\n",
        "\n",
        "Now, let's use Memory Bank to generate and store memories. Unlike `InMemoryMemoryService`, only content that's meaningful interactions will be persisted. If there's no meaningful content, the `add_session_to_memory` call will be a no-op; memories are not generated.\n",
        "\n",
        "Memory generation happens in the background. It's a non-blocking operation, so `add_session_to_memory` may complete before memories are actually generated in the background. If you want memory generation to be a blocking operation, use the Agent Engine SDK:\n",
        "\n",
        "```\n",
        "client.agent_engines.memories.generate(\n",
        "    ...,\n",
        "    config={\n",
        "        # Default value is True.\n",
        "        \"wait_for_completion\": True\n",
        "    }\n",
        ")\n",
        "```\n",
        "\n",
        "There are two key operations that Memory Bank performs under-the-hood with memory generation:\n",
        "\n",
        "* **Extraction**: Extracts information about the user from their conversations with the agent. Only information that matches at least one of your instance's memory topics will be persisted.\n",
        "* **Consolidation**: Identifies if existing memories with the same scope should be deleted or updated based on the extracted information. Memory Bank checks that new memories are not duplicative or contradictory before merging them with existing memories. If existing memories don't overlap with the new information, a new memory will be created."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HzYBHoUNSZKn"
      },
      "source": [
        "#### ADK `add_session_to_memory`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ytzu349cBdW6"
      },
      "outputs": [],
      "source": [
        "await memory_bank_service.add_session_to_memory(chit_chat_session)\n",
        "await memory_bank_service.add_session_to_memory(recommendations_session)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3ihzck53quZ"
      },
      "source": [
        "If you run this immediately after `add_session_to_memory`, there may be no memories printed. `add_session_to_memory` is a non-blocking function and memory generation can take a few seconds to run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3emdHmGhBtMd"
      },
      "outputs": [],
      "source": [
        "await memory_bank_service.search_memory(\n",
        "    app_name=APP_NAME,\n",
        "    user_id=USER_ID,\n",
        "    query=\"What should I get my mom for mother's day\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DtdqszuOUyJ0"
      },
      "source": [
        "#### Agent Engine SDK `GenerateMemories`\n",
        "\n",
        "If I want to interact directly with Memory Bank, I can use the Agent Engine SDK. For example, I can use [`RetrieveMemories`](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/fetch-memories#scope-based) to retrieve my memories with or without similarity search.\n",
        "\n",
        "Memories are isolated by their `scope`, which is an arbitrary dictionary. `VertexAiMemoryBankService` uses the scope keys `app_name` and `user_id`. `app_name` is set by your Runner, and `user_id` is set when creating your `Session`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xXxDASPUG3OK"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        ")\n",
        "list(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E-KVnoBHV1G5"
      },
      "source": [
        "You can also use `Retrievememories` to retrieve memories using similarity search."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uwH2UgOZVs36"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        "    similarity_search_params={\n",
        "        \"search_query\": \"what should I get my niece for her birthday?\"\n",
        "    },\n",
        ")\n",
        "list(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QWWmfmn7XHNf"
      },
      "source": [
        "Memory Bank uses a process called \"consolidation\" to ensure that we do not upload duplicative or contradictory information for the same scope."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gyHYYlxIXGKQ"
      },
      "outputs": [],
      "source": [
        "import pprint\n",
        "\n",
        "operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\n",
        "                            \"text\": \"I like the idea of getting my niece a bike for her third birthday\"\n",
        "                        }\n",
        "                    ],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")\n",
        "print(\"***OPERATION RESPONSE***\")\n",
        "pprint.pprint(operation)\n",
        "\n",
        "print(\"\\n***GENERATED MEMORIES***\")\n",
        "for generated_memory in operation.response.generated_memories:\n",
        "    pprint.pprint(client.agent_engines.memories.get(name=generated_memory.memory.name))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FUAJRM8pYiNj"
      },
      "source": [
        "What if your agent extracted memories that you want to combine with your existing memories? You can also use `GenerateMemories` to consolidate pre-extracted facts for the same scope [documenation](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/generate-memories#consolidate-pre-extracted-memories).\n",
        "\n",
        "Memory Bank already contains memories on this topic. So, the relevant existing memories will be updated rather than creating a new, duplicative memory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hVaFGo7RY0zE"
      },
      "outputs": [],
      "source": [
        "operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        "    direct_memories_source={\"direct_memories\": [{\"fact\": \"I got my niece a red bike\"}]},\n",
        ")\n",
        "\n",
        "print(\"***OPERATION RESPONSE***\")\n",
        "pprint.pprint(operation)\n",
        "\n",
        "print(\"\\n***GENERATED MEMORIES***\")\n",
        "for generated_memory in operation.response.generated_memories:\n",
        "    if generated_memory.action != \"DELETED\":\n",
        "        pprint.pprint(\n",
        "            client.agent_engines.memories.get(name=generated_memory.memory.name)\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ql9iDjUO5Mxb"
      },
      "source": [
        "#### Automate with Callbacks\n",
        "\n",
        "To automate the triggering of memory generation, you can implement a callback to generate memories after every turn. See the [`Remember this: Agent state and memory with ADK` blog post](https://cloud.google.com/blog/topics/developers-practitioners/remember-this-agent-state-and-memory-with-adk) for more information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sq4QSggUAuLu"
      },
      "outputs": [],
      "source": [
        "from google import adk\n",
        "\n",
        "\n",
        "async def auto_save_session_to_memory_callback(callback_context):\n",
        "    # Use the invocation context to access the conversation history that should\n",
        "    # be used as the data source for memory generation.\n",
        "    await memory_bank_service.add_session_to_memory(\n",
        "        callback_context._invocation_context.session\n",
        "    )\n",
        "    print(\"\\n****Triggered memory generation****\\n\")\n",
        "\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        "    tools=[adk.tools.preload_memory_tool.PreloadMemoryTool()],\n",
        "    after_agent_callback=auto_save_session_to_memory_callback,\n",
        ")\n",
        "\n",
        "runner = Runner(\n",
        "    agent=agent,\n",
        "    app_name=APP_NAME,\n",
        "    session_service=session_service,\n",
        "    memory_service=memory_bank_service,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o-oU37SyA5J7"
      },
      "outputs": [],
      "source": [
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "\n",
        "call_agent(\n",
        "    \"What christmas present recommendations do you have for my niece?\",\n",
        "    session.id,\n",
        "    user_id=USER_ID,\n",
        ")\n",
        "\n",
        "call_agent(\n",
        "    \"I think I'll get her a doll. What types are good for her age?\",\n",
        "    session.id,\n",
        "    user_id=USER_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E9_pkrRUBBSP"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        "    similarity_search_params={\n",
        "        \"search_query\": \"What was I thinking about getting my niece for Christmas?\"\n",
        "    },\n",
        ")\n",
        "list(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "koK2O-odR36D"
      },
      "source": [
        "The prior example used the entire Session to generate memories. However, this can be costly as the size of the session increases.\n",
        "\n",
        "You can also use callbacks to only save the current turn to Memory Bank. This has the downside, however, where you lose some of the context of the conversation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F5hWwSvEPaaW"
      },
      "outputs": [],
      "source": [
        "async def auto_save_user_turn_to_memory_callback(callback_context):\n",
        "    last_turn = callback_context._invocation_context.user_content\n",
        "    client.agent_engines.memories.generate(\n",
        "        name=agent_engine.api_resource.name,\n",
        "        scope={\"app_name\": APP_NAME, \"user_id\": USER_ID},\n",
        "        direct_contents_source={\"events\": [{\"content\": last_turn}]},\n",
        "        config={\"wait_for_completion\": False},\n",
        "    )\n",
        "    print(\"\\n****Triggered memory generation****\\n\")\n",
        "\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        "    tools=[adk.tools.preload_memory_tool.PreloadMemoryTool()],\n",
        "    before_agent_callback=auto_save_user_turn_to_memory_callback,\n",
        ")\n",
        "\n",
        "runner = Runner(\n",
        "    agent=agent,\n",
        "    app_name=APP_NAME,\n",
        "    session_service=session_service,\n",
        "    memory_service=memory_bank_service,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ziRXQsAJPeYl"
      },
      "outputs": [],
      "source": [
        "call_agent(\n",
        "    \"I think I'll get her a doll. What types are good for her age?\",\n",
        "    session.id,\n",
        "    user_id=USER_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rstDOvEhWogm"
      },
      "source": [
        "## Using Memory in your agent\n",
        "\n",
        "Now, let's link our agent to Memory, so that memories can be fetched by your agent and used for inference. This allows your agent to remember information from prior sessions in a new, empty session.\n",
        "\n",
        "You can use ADK's built-in tools to fetch memories and incorporate them in the prompt. When using the built-in tools, it's important to provide both a memory tool when creating your Agent and a memory service when defining your Runner."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yOiYWk7GZlyN"
      },
      "source": [
        "### `adk.PreloadMemoryTool`\n",
        "The `PreloadMemoryTool` always retrieves memories at the start of each turn and includes the memories in the system instruction.\n",
        "\n",
        "Under the hood, it invokes `memory_service.search_memory` for the given User and App Name. Memory is invoked before the LLM is called, so there will not be an associated tool call logged for memory.\n",
        "\n",
        "We're going to use a [callback](https://google.github.io/adk-docs/callbacks/) to print out the System Instructions. Since we use `before_model_callback`, the callback executes right before the model is called.\n",
        "\n",
        "<img src=\"https://google.github.io/adk-docs/assets/callback_flow.png\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l1ENgxmkWnVl"
      },
      "outputs": [],
      "source": [
        "from google import adk\n",
        "\n",
        "\n",
        "def log_system_instructions(callback_context, llm_request):\n",
        "    \"\"\"A callback to print the LLM request.\"\"\"\n",
        "    print(\n",
        "        f\"\\n*System Instruction*:\\n{llm_request.config.system_instruction}\\n*********\\n\"\n",
        "    )\n",
        "\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        "    tools=[adk.tools.preload_memory_tool.PreloadMemoryTool()],\n",
        "    before_model_callback=log_system_instructions,\n",
        ")\n",
        "\n",
        "runner = 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": "r0g7dxD0aG_H"
      },
      "source": [
        "This conversation uses a new session, so it doesn't have access to the user's conversation history without using memory.\n",
        "\n",
        "Since we're using `PreloadMemoryTool`, the retrieved memories will be appended to the System Instructions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2FD9VV_4Hqja"
      },
      "outputs": [],
      "source": [
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "\n",
        "call_agent(\n",
        "    \"Can you remind me what I got my niece for her christmas?\",\n",
        "    session.id,\n",
        "    user_id=USER_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TSpv-aGuaG4g"
      },
      "source": [
        "### adk.LoadMemoryTool\n",
        "\n",
        "Now, let's use `LoadMemoryTool`. Unlike `PreloadMemoryTool`, this tool acts like a standard tool. Your agent needs to decide whether the tool should be invoked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIbplEW0ZlRR"
      },
      "outputs": [],
      "source": [
        "from google import adk\n",
        "\n",
        "\n",
        "def log_tool_call(tool, args, tool_response, **kwargs):\n",
        "    \"\"\"A callback to print the LLM request.\"\"\"\n",
        "    print(f\"\\n*Tool*:\\n{tool}\")\n",
        "    print(f\"\\n*Tool call*:\\n{args}\")\n",
        "    print(f\"\\n*Tool response*:\\n{tool_response}\\n*********\\n\")\n",
        "\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        "    tools=[adk.tools.load_memory_tool.LoadMemoryTool()],\n",
        "    before_model_callback=log_system_instructions,\n",
        "    after_tool_callback=log_tool_call,\n",
        ")\n",
        "\n",
        "runner = Runner(\n",
        "    agent=agent,\n",
        "    app_name=APP_NAME,\n",
        "    session_service=session_service,\n",
        "    memory_service=memory_bank_service,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u2LLB98aZssb"
      },
      "outputs": [],
      "source": [
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "\n",
        "call_agent(\n",
        "    \"Can you remind me what I got my niece for her christmas?\",\n",
        "    session.id,\n",
        "    user_id=USER_ID,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9uGnTDQO42tf"
      },
      "source": [
        "If the agent decides that Memory is not helpful for a given query, the Memory tool will not be invoked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pmm3SJOsd288"
      },
      "outputs": [],
      "source": [
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "\n",
        "call_agent(\"Hi!\", session.id, user_id=USER_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-M6GeiZzAox6"
      },
      "source": [
        "### Custom callback to retrieve memories\n",
        "\n",
        "If you want more control over how memories are retrieved, you can create your own custom callback.\n",
        "\n",
        "For example, you can use the Agent Engine SDK within the callback to retrieve memories without similarity search. This is lower latency than using similarity search.\n",
        "\n",
        "You can also customize your `scope` rather than using the default \"app_name\" and \"user_id\" scope keys used by ADK's memory service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TUc_jTzr5QSm"
      },
      "outputs": [],
      "source": [
        "def retrieve_memories_callback(callback_context, llm_request):\n",
        "    user_id = callback_context._invocation_context.user_id\n",
        "\n",
        "    response = client.agent_engines.memories.retrieve(\n",
        "        name=agent_engine.api_resource.name,\n",
        "        # Unlike the ADK Memory Service, this does not use App Name in the scope.\n",
        "        scope={\"user_id\": user_id},\n",
        "    )\n",
        "    memories = [f\"* {memory.memory.fact}\" for memory in list(response)]\n",
        "    if not memories:\n",
        "        # No memories to add to System Instructions.\n",
        "        return\n",
        "    # Append formatted memories to the System Instructions\n",
        "    llm_request.config.system_instruction += (\n",
        "        \"\\nHere is information that you have about the user:\\n\"\n",
        "    )\n",
        "    llm_request.config.system_instruction += \"\\n\".join(memories)\n",
        "\n",
        "\n",
        "def log_tool_call(tool, args, tool_response, **kwargs):\n",
        "    \"\"\"A callback to print the LLM request.\"\"\"\n",
        "    print(f\"\\n*Tool*:\\n{tool}\")\n",
        "    print(f\"\\n*Tool call*:\\n{args}\")\n",
        "    print(f\"\\n*Tool response*:\\n{log_tool_call}\\n*********\\n\")\n",
        "\n",
        "\n",
        "agent = LlmAgent(\n",
        "    model=MODEL,\n",
        "    name=\"Generic_QA_Agent\",\n",
        "    instruction=\"Answer the user's questions\",\n",
        "    before_model_callback=retrieve_memories_callback,\n",
        ")\n",
        "\n",
        "runner = 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": "hPHr71m4_hiJ"
      },
      "source": [
        "Since we're using a different scope than the ADK Memory Service, we don't have access to the prior memories. Memories are isolated by their scope dictionary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0zV-bZtz51E6"
      },
      "outputs": [],
      "source": [
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "\n",
        "call_agent(\"Hi!\", session.id, user_id=USER_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F-RoQ0UF_s5v"
      },
      "source": [
        "Let's first generate memories for our new scope."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3PjiIK2L_r43"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\"content\": {\"role\": \"user\", \"parts\": [{\"text\": \"I have four nieces!\"}]}}\n",
        "        ]\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fTzgPdtU_9Qd"
      },
      "outputs": [],
      "source": [
        "call_agent(\"What information do you know about me?\", session.id, user_id=USER_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SAw7Dax-biQW"
      },
      "source": [
        "## Customize your Memory Bank's behavior\n",
        "\n",
        "So far, we've used Memory Bank with it's default settings. Now, let's customize some configurations. The configuration is set when creating or updating your Agent Engine."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6nIKVog-sitb"
      },
      "source": [
        "### Customizing Memory Extraction\n",
        "\n",
        "By default, Memory Bank considers memories that fit the following topics to be meaningful:\n",
        "\n",
        "* **Personal information (`USER_PERSONAL_INFO`)**: Significant personal information about the user, like names, relationships, hobbies, and important dates. For example, \"I work at Google\" or \"My wedding anniversary is on December 31\".\n",
        "* **User preferences (`USER_PREFERENCES`)**: Stated or implied likes, dislikes, preferred styles, or patterns. For example, \"I prefer the middle seat.\"\n",
        "* **Key conversation events and task outcomes (`KEY_CONVERSATION_DETAILS`)**: Important milestones or conclusions within the dialogue. For example, \"I booked plane tickets for a round trip between JFK and SFO. I leave on June 1, 2025 and return on June 7, 2025.\"\n",
        "* **Explicit remember / forget instructions (`EXPLICIT_INSTRUCTIONS`)**: Information that the user explicitly asks the agent to remember or forget. For example, if the user says \"Remember that I primarily use Python,\" Memory Bank generates a memory such as \"I primarily use Python.\"\n",
        "\n",
        "These are managed topics where Memory Bank manages their definition. If you want to use a subset of these topics or use custom topics, you can provide a \"CustomizationConfig\" to customize what information Memory Bank should consider meaningful to persist."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G2pwuUddXf59"
      },
      "source": [
        "#### Managed Topics\n",
        "\n",
        "We're going to update our Memory Bank to only extract user preferences.\n",
        "\n",
        "This will take affect for all requests to Memory Bank, regardless of whether you use ADK or Agent Engine SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6qc7x-1tXWS_"
      },
      "outputs": [],
      "source": [
        "user_preferences_config = {\n",
        "    \"scope_keys\": [\"user_id\"],\n",
        "    \"memory_topics\": [\n",
        "        {\"managed_memory_topic\": {\"managed_topic_enum\": \"USER_PREFERENCES\"}}\n",
        "    ],\n",
        "}\n",
        "\n",
        "\n",
        "client.agent_engines.update(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    config={\n",
        "        \"context_spec\": {\n",
        "            \"memory_bank_config\": {\n",
        "                \"customization_configs\": [user_preferences_config],\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t3RNGsTmYI74"
      },
      "source": [
        "Memories that fit the `USER_PREFERENCES` category will be persisted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g_efAhHIYGEC"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": \"I prefer the aisle seat\"}],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0W7oEz7KYOks"
      },
      "source": [
        "Memories that don't fit `USER_PREFERENCES` won't be persisted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V9sQPIGQYTr0"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": \"I'm going to get my niece a red bike\"}],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SodYXQQrXMqD"
      },
      "source": [
        "#### Custom Topics\n",
        "Alternatively, you can define your own custom topics where you define your own labels, descriptions, and examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o-W0cxdqsjRi"
      },
      "outputs": [],
      "source": [
        "from vertexai.types import MemoryBankCustomizationConfig as CustomizationConfig\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigGenerateMemoriesExample as GenerateMemoriesExample,\n",
        ")\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource as ConversationSource,\n",
        ")\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent as ConversationSourceEvent,\n",
        ")\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory as ExampleGeneratedMemory,\n",
        ")\n",
        "from vertexai.types import MemoryBankCustomizationConfigMemoryTopic as MemoryTopic\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic as CustomMemoryTopic,\n",
        ")\n",
        "\n",
        "memory_topic = MemoryTopic(\n",
        "    custom_memory_topic=CustomMemoryTopic(\n",
        "        label=\"business_feedback\",\n",
        "        description=\"\"\"Specific user feedback about their experience at\n",
        "the coffee shop. This includes opinions on drinks, food, pastries, ambiance,\n",
        "staff friendliness, service speed, cleanliness, and any suggestions for\n",
        "improvement.\"\"\",\n",
        "    )\n",
        ")\n",
        "\n",
        "\n",
        "example = GenerateMemoriesExample(\n",
        "    conversation_source=ConversationSource(\n",
        "        events=[\n",
        "            ConversationSourceEvent(\n",
        "                content=Content(\n",
        "                    role=\"model\",\n",
        "                    parts=[\n",
        "                        Part(\n",
        "                            text=\"Welcome back to The Daily Grind! We'd love to hear your feedback on your visit.\"\n",
        "                        )\n",
        "                    ],\n",
        "                )\n",
        "            ),\n",
        "            ConversationSourceEvent(\n",
        "                content=Content(\n",
        "                    role=\"user\",\n",
        "                    parts=[\n",
        "                        Part(\n",
        "                            text=\"Hey. The drip coffee was a bit lukewarm today, which was a bummer. Also, the music was way too loud, I could barely hear my friend.\"\n",
        "                        )\n",
        "                    ],\n",
        "                )\n",
        "            ),\n",
        "        ]\n",
        "    ),\n",
        "    generated_memories=[\n",
        "        ExampleGeneratedMemory(\n",
        "            fact=\"The user reported that the drip coffee was lukewarm.\"\n",
        "        ),\n",
        "        ExampleGeneratedMemory(\n",
        "            fact=\"The user felt the music in the shop was too loud.\"\n",
        "        ),\n",
        "    ],\n",
        ")\n",
        "\n",
        "noop_example = GenerateMemoriesExample(\n",
        "    conversation_source=ConversationSource(\n",
        "        events=[\n",
        "            ConversationSourceEvent(\n",
        "                content=Content(\n",
        "                    role=\"model\",\n",
        "                    parts=[\n",
        "                        Part(\n",
        "                            text=\"Welcome back to The Daily Grind! We'd love to hear your feedback on your visit.\"\n",
        "                        )\n",
        "                    ],\n",
        "                )\n",
        "            ),\n",
        "            ConversationSourceEvent(\n",
        "                content=Content(\n",
        "                    role=\"user\", parts=[Part(text=\"Thanks for the coffee!\")]\n",
        "                )\n",
        "            ),\n",
        "        ]\n",
        "    ),\n",
        "    generated_memories=[],\n",
        ")\n",
        "\n",
        "client.agent_engines.update(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    config={\n",
        "        \"context_spec\": {\n",
        "            \"memory_bank_config\": {\n",
        "                \"customization_configs\": [\n",
        "                    CustomizationConfig(\n",
        "                        memory_topics=[memory_topic],\n",
        "                        generate_memories_examples=[example, noop_example],\n",
        "                    )\n",
        "                ],\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XE783Z_TDHC_"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": \"123\"},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": \"I prefer the aisle seat\"}],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u1j8iPHBdg6n"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": \"123\"},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": \"You should have almond milk\"}],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DJ7v62dIePo9"
      },
      "outputs": [],
      "source": [
        "list(\n",
        "    client.agent_engines.memories.retrieve(\n",
        "        name=agent_engine.api_resource.name, scope={\"user_id\": \"123\"}\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KVPXrs2GS_ZJ"
      },
      "source": [
        "### Dynamic TTL\n",
        "\n",
        "You can also configure Memory Bank to automatically set the TTL for any generated or manually created (or updated) memories."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dtf5AcltTV2I"
      },
      "source": [
        "#### Default TTL\n",
        "\n",
        "Default TTL applies to all operations that create or update a Memory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vQZKHD-TS-qD"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.update(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    config={\n",
        "        \"context_spec\": {\n",
        "            \"memory_bank_config\": {\n",
        "                \"ttl_config\": {\n",
        "                    # 30 days\n",
        "                    \"default_ttl\": f\"{60 * 60 * 24 * 30}s\"\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vICeAqPyCaT0"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\"content\": {\"role\": \"user\", \"parts\": [{\"text\": \"I have four nieces!\"}]}}\n",
        "        ]\n",
        "    },\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "client.agent_engines.memories.get(\n",
        "    name=response.response.generated_memories[0].memory.name\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oa0A4ogWT8BM"
      },
      "source": [
        "#### Granular TTL\n",
        "\n",
        "The TTL is calculated based on which operation created or updated the Memory. If not set for a given operation, then the operation won't update the Memory's expiration time.\n",
        "\n",
        "Parameters include:\n",
        " * `create_ttl`: Applies to Memories created with CreateMemory operations\n",
        " * `generate_created_ttl`: Applies to Memories created with GenerateMemories\n",
        " * `generate_updated_ttl`: Applies to Memories updated with GenerateMemories"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3g02gtyrT7fG"
      },
      "outputs": [],
      "source": [
        "client.agent_engines.update(\n",
        "    name=agent_engine.api_resource.name,\n",
        "    config={\n",
        "        \"context_spec\": {\n",
        "            \"memory_bank_config\": {\n",
        "                \"ttl_config\": {\n",
        "                    \"granular_ttl_config\": {\n",
        "                        # 365 days\n",
        "                        \"generate_created_ttl\": f\"{24 * 60 * 60 * 365}s\",\n",
        "                    }\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dT3ffaBksw99"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [{\"text\": \"Remember that I work at Google!\"}],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        ")\n",
        "\n",
        "client.agent_engines.memories.get(\n",
        "    name=response.response.generated_memories[0].memory.name\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O6UGzBpCTpHB"
      },
      "outputs": [],
      "source": [
        "response = client.agent_engines.memories.generate(\n",
        "    scope={\"user_id\": USER_ID},\n",
        "    name=agent_engine.api_resource.name,\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\"content\": {\"role\": \"user\", \"parts\": [{\"text\": \"I work at Google!\"}]}}\n",
        "        ]\n",
        "    },\n",
        ")\n",
        "\n",
        "client.agent_engines.memories.get(\n",
        "    name=response.response.generated_memories[0].memory.name\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a678c7a2bf29"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "It's always a best practice in cloud development to clean up resources you no longer need to avoid incurring unexpected costs. This final cell deletes the AgentEngine resources we created throughout this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "10fe83835935"
      },
      "outputs": [],
      "source": [
        "delete_agent_engines = True\n",
        "\n",
        "agent_engine_name = agent_engine.api_resource.name\n",
        "\n",
        "if delete_agent_engines:\n",
        "    # Delete agent engine\n",
        "    client.agent_engines.delete(name=agent_engine_name, force=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_memory_bank_on_adk.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
