{
  "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\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.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.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.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.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| [Kimberly Milam](https://github.com/klmilam) |\n",
        "| [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook is a hands-on guide to mastering **Vertex AI Memory Bank**, a service for building stateful, context-aware conversational AI agents. You will learn how to give your agent a persistent, long-term memory, allowing it to recall guest preferences and past interactions across multiple sessions to provide truly personalized hospitality experiences. We will apply these concepts to a practical, real-world hotel scenario: building a sophisticated hotel concierge assistant.\n",
        "\n",
        "By the end of this tutorial, you will not only understand the core concepts of Memory Bank but also know how to apply them to build an assistant that remembers guest preferences, dietary restrictions, room preferences, and maintains context across conversations to deliver exceptional, personalized service.\n",
        "\n",
        "Here's a high-level overview of the steps we'll take:\n",
        "\n",
        "* **Initial Setup**: We will begin with the fundamentals, configuring a new Memory Bank instance and learning how to create guest sessions to store and retrieve conversation history.\n",
        "* **Basic Memory Operations**: We will explore how to generate memories from conversations and retrieve them to understand what the system remembers about each guest.\n",
        "* **Real-World Application**: We will see how to use these memories to personalize guest interactions when they return to the hotel.\n",
        "* **Resource Management**: Finally, we will address essential operational aspects by properly cleaning up resources."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and required packages\n",
        "\n",
        "First, let's install the Vertex AI SDK.\n",
        "\n",
        "**Note**: This will install the SDK. Colab may prompt you to restart the runtime after installation. This is expected behavior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "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).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "# Initialize the Vertex AI client\n",
        "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "print(\"✅ Vertex AI client initialized!\")\n",
        "print(f\"   Project: {PROJECT_ID}\")\n",
        "print(f\"   Location: {LOCATION}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "We're importing standard Python libraries and, importantly, several class-based types from the Vertex AI SDK.\n",
        "\n",
        "To make the code more readable, we're creating shorter aliases for these long class names. This is a common Python practice that helps keep our code clean and concise without sacrificing the benefits of using the typed classes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import os\n",
        "import uuid\n",
        "import warnings\n",
        "\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "# Import class-based types for Memory Bank\n",
        "from vertexai import types\n",
        "\n",
        "# Basic configuration types\n",
        "MemoryBankConfig = types.ReasoningEngineContextSpecMemoryBankConfig\n",
        "SimilaritySearchConfig = (\n",
        "    types.ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig\n",
        ")\n",
        "GenerationConfig = types.ReasoningEngineContextSpecMemoryBankConfigGenerationConfig\n",
        "\n",
        "# Advanced configuration types\n",
        "TtlConfig = types.ReasoningEngineContextSpecMemoryBankConfigTtlConfig\n",
        "GranularTtlConfig = (\n",
        "    types.ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig\n",
        ")\n",
        "CustomizationConfig = types.MemoryBankCustomizationConfig\n",
        "MemoryTopic = types.MemoryBankCustomizationConfigMemoryTopic\n",
        "ManagedMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic\n",
        "CustomMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic\n",
        "GenerateMemoriesExample = types.MemoryBankCustomizationConfigGenerateMemoriesExample\n",
        "ConversationSource = (\n",
        "    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource\n",
        ")\n",
        "ConversationSourceEvent = (\n",
        "    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent\n",
        ")\n",
        "ExampleGeneratedMemory = (\n",
        "    types.MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory\n",
        ")\n",
        "ManagedTopicEnum = types.ManagedTopicEnum\n",
        "\n",
        "print(\"✅ Libraries imported successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ivZ7M1tFoxJs"
      },
      "source": [
        "### Define helper function for displaying memories\n",
        "\n",
        "This helper function provides consistent formatting when displaying generated memories throughout the tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7tkFZu8RowqZ"
      },
      "outputs": [],
      "source": [
        "def display_generated_memories(operation, client, title=\"Generated Memories\"):\n",
        "    \"\"\"Display memories from a generation operation with consistent formatting.\n",
        "\n",
        "    Args:\n",
        "        operation: The result from client.agent_engines.memories.generate()\n",
        "        client: The Vertex AI client instance\n",
        "        title: Title to display above the memories\n",
        "    \"\"\"\n",
        "    if operation.response and operation.response.generated_memories:\n",
        "        print(f\"\\n✅ {title}: {len(operation.response.generated_memories)}\\n\")\n",
        "\n",
        "        for i, gen_memory in enumerate(operation.response.generated_memories, 1):\n",
        "            if gen_memory.action != \"DELETED\" and gen_memory.memory:\n",
        "                try:\n",
        "                    full_memory = client.agent_engines.memories.get(\n",
        "                        name=gen_memory.memory.name\n",
        "                    )\n",
        "                    action_icon = \"🆕\" if gen_memory.action == \"CREATED\" else \"🔄\"\n",
        "                    print(f\"   {action_icon} {i}. {full_memory.fact}\")\n",
        "                except Exception as e:\n",
        "                    print(f\"   ⚠️ Could not retrieve memory: {e}\")\n",
        "    else:\n",
        "        print(f\"\\n📭 No {title.lower()} found\")\n",
        "\n",
        "\n",
        "print(\"✅ Helper function defined!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Create your Hotel Concierge Memory Bank\n",
        "\n",
        "Now let's set up the foundation for our hotel concierge assistant by creating an Agent Engine with Memory Bank capabilities."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QGMXOQF4pFvt"
      },
      "source": [
        "### Create Agent Engine with Memory Bank configuration\n",
        "\n",
        "The AgentEngine resource acts as the top-level container for your Memory Bank instance. To create one, we need to provide a configuration.\n",
        "\n",
        "Here, MemoryBankConfig has two key parts:\n",
        "\n",
        "1. **`similarity_search_config`**: This specifies the **embedding model** used for similarity searches. When a guest asks \"What did I request last time?\", this model helps find the most relevant memories. We're using `text-embedding-005`, which is excellent for English conversations. If you expect multilingual guests, consider `text-multilingual-embedding-002`.\n",
        "\n",
        "2. **`generation_config`**: This defines the **LLM** that will extract and consolidate memories from conversations. The default, `gemini-2.5-flash`, is a fast and capable model perfect for this task. It reads conversations and intelligently extracts key facts like dietary preferences or room temperature preferences.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bR3ru_61pHMg"
      },
      "outputs": [],
      "source": [
        "print(\"🧠 Creating Memory Bank configuration for hotel concierge...\\n\")\n",
        "\n",
        "basic_memory_config = MemoryBankConfig(\n",
        "    # Which embedding model to use for similarity search\n",
        "    similarity_search_config=SimilaritySearchConfig(\n",
        "        embedding_model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005\"\n",
        "    ),\n",
        "    # Which LLM to use for extracting memories from conversations\n",
        "    generation_config=GenerationConfig(\n",
        "        model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(\"✅ Memory Bank configuration created!\")\n",
        "print(\"   Embedding model: text-embedding-005\")\n",
        "print(\"   Generation model: gemini-2.5-flash\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YD6vm2kRpTdw"
      },
      "source": [
        "Now, we create the AgentEngine resource. By default, Memory Bank is enabled when you create an Agent Engine. This call provisions the necessary backend infrastructure to store and retrieve guest memories.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T53ZZDbQpKIY"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🛠️ Creating Agent Engine with Memory Bank...\\n\")\n",
        "print(\"⏳ This provisions the backend infrastructure for guest memory storage...\")\n",
        "\n",
        "agent_engine = client.agent_engines.create(\n",
        "    config={\"context_spec\": {\"memory_bank_config\": basic_memory_config}}\n",
        ")\n",
        "\n",
        "agent_engine_name = agent_engine.api_resource.name\n",
        "\n",
        "print(\"\\n✅ Agent Engine created successfully!\")\n",
        "print(f\"   Resource Name: {agent_engine_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RJDfHDmtp6nQ"
      },
      "source": [
        "## Store a guest conversation\n",
        "\n",
        "Now let's simulate a real hotel check-in conversation and see how Memory Bank extracts and stores important guest information.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GlxRLL6Lp9uF"
      },
      "source": [
        "### Create a Session for a guest\n",
        "\n",
        "A **Session** is a chronological log of a single interaction between a guest and your concierge agent. It's the raw material from which memories are made.\n",
        "\n",
        "Each session is linked to a `user_id` (in this case, the guest's identifier). This allows the agent to recall information for a specific guest across different stays and interactions. Think of it as the guest's profile ID in your hotel system.\n",
        "\n",
        "**Important**: Using Vertex AI Agent Engine Session is not the only option supported. You can also [provide the source conversation directly in JSON format](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/generate-memories#json-format) if you're using a different session storage system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M7LCXzRMp89n"
      },
      "outputs": [],
      "source": [
        "print(\"💬 Creating a session for guest check-in...\\n\")\n",
        "\n",
        "# Generate a unique guest identifier\n",
        "guest_id = \"guest_emma_\" + str(uuid.uuid4())[:4]\n",
        "\n",
        "# Create a session for this guest\n",
        "session = client.agent_engines.sessions.create(\n",
        "    name=agent_engine_name,\n",
        "    user_id=guest_id,\n",
        "    config={\"display_name\": f\"Check-in conversation for {guest_id}\"},\n",
        ")\n",
        "\n",
        "session_name = session.response.name\n",
        "\n",
        "print(\"✅ Session created successfully!\")\n",
        "print(f\"   Session: {session_name}\")\n",
        "print(f\"   Guest ID: {guest_id}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rIBwUZt2qRHx"
      },
      "source": [
        "### Define the guest conversation\n",
        "\n",
        "This is the raw conversational data we'll use. It's a simple list of dictionaries, each representing a turn in the dialogue between the guest and the hotel concierge.\n",
        "\n",
        "This conversation contains valuable information about the guest's preferences that we want to remember for future stays. Notice how the guest mentions dietary restrictions, room preferences, and special requests.\n",
        "\n",
        "As you read this conversation, think about what a good concierge would remember:\n",
        "- The guest's name (Emma)\n",
        "- Dietary restrictions (vegetarian, lactose intolerant)\n",
        "- Room preferences (high floor, away from elevator, quiet)\n",
        "- Special requests (extra pillows, white noise machine)\n",
        "- Temperature preferences (cooler room, 68°F)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wh0SNBUrOKo"
      },
      "outputs": [],
      "source": [
        "# A realistic hotel check-in conversation\n",
        "guest_conversation = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Hi! I'm Emma Chen. I have a reservation for a deluxe room for three nights.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Welcome to our hotel, Ms. Chen! I have your reservation right here. Let me get you checked in.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Great! Just so you know, I'm vegetarian and lactose intolerant, so please note that for room service and the breakfast buffet.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Absolutely, I've noted your dietary preferences. Our kitchen will be informed, and we have excellent vegetarian options.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Thank you! Also, I prefer a high floor room, away from the elevator. I'm a light sleeper, so somewhere quiet would be ideal.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Of course! I'll assign you a corner room on the 12th floor. It's one of our quietest locations.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Perfect! One more thing - can I get extra pillows and a white noise machine? And I prefer the room temperature around 68°F.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Certainly! I'll have housekeeping bring extra pillows and a white noise machine to your room, and I'll set the thermostat to 68°F before your arrival.\",\n",
        "    },\n",
        "]\n",
        "\n",
        "print(\"✅ Guest conversation defined!\")\n",
        "print(f\"   {len(guest_conversation)} conversation turns\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nn0qg_yhrg25"
      },
      "source": [
        "### Add conversation to the session\n",
        "\n",
        "Here, we loop through our conversation and append each turn as an event to the session we created earlier. This persists the conversation history, making it available for memory generation.\n",
        "\n",
        "Memory Bank needs the complete conversation stored in a session before it can analyze and extract memories. Each turn is timestamped and attributed to either the guest (user) or the concierge (model)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0c7MD8pdriEh"
      },
      "outputs": [],
      "source": [
        "print(\"⬆️ Adding conversation to session...\\n\")\n",
        "\n",
        "invocation_id = 0\n",
        "\n",
        "for turn in guest_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=session_name,\n",
        "        author=guest_id,  # Required by Sessions\n",
        "        invocation_id=str(invocation_id),  # Required by Sessions\n",
        "        timestamp=datetime.datetime.now(\n",
        "            tz=datetime.timezone.utc\n",
        "        ),  # Required by Sessions\n",
        "        config={\n",
        "            \"content\": {\"role\": turn[\"role\"], \"parts\": [{\"text\": turn[\"message\"]}]}\n",
        "        },\n",
        "    )\n",
        "\n",
        "    invocation_id += 1\n",
        "    icon = \"👤\" if turn[\"role\"] == \"user\" else \"🤖\"\n",
        "    print(f\"{icon} {turn['message']}\")\n",
        "\n",
        "print(\"\\n✅ Conversation successfully added to session!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aptoNJ1-rlNv"
      },
      "source": [
        "## Generate memories from the conversation\n",
        "\n",
        "Now let's see what the AI automatically remembers from this conversation.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p7i9wnSZrpe3"
      },
      "source": [
        "### Extract and consolidate guest memories\n",
        "\n",
        "This is the core of memory generation. The generate method kicks off an operation that performs two main steps.\n",
        "\n",
        "This is how Memory Bank works:\n",
        "\n",
        "1. **Extraction**: The generation model (gemini-2.5-flash) reads the conversation and extracts key facts. With the default configuration, it looks for information that matches pre-defined **Managed Topics** like:\n",
        "   - `USER_PERSONAL_INFO`: Personal details like the guest's name\n",
        "   - `USER_PREFERENCES`: Stated preferences about rooms, amenities, temperature, etc.\n",
        "   - `EXPLICIT_INSTRUCTIONS`: Specific requests the guest made\n",
        "\n",
        "2. **Consolidation**: Memory Bank intelligently merges new facts with existing memories (if the guest has stayed before), avoiding duplicates and resolving contradictions. For example, if a guest previously preferred 70°F but now says 68°F, the system updates the memory.\n",
        "\n",
        "\n",
        "**Note**: The `wait_for_completion=True` flag makes this a blocking call, which is useful for this tutorial. In production, you would typically set it to `False` to run in the background and poll for results later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b6k_TEP5r781"
      },
      "outputs": [],
      "source": [
        "print(\"🧠 Analyzing conversation and extracting memories...\\n\")\n",
        "print(\n",
        "    \"⏳ The AI is reading the conversation and identifying important facts to remember...\"\n",
        ")\n",
        "\n",
        "# Generate memories from the session\n",
        "operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine_name,\n",
        "    vertex_session_source={\"session\": session_name},\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Memory extraction complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3PphqV9nr-3P"
      },
      "source": [
        "### Display extracted memories\n",
        "\n",
        "Let's see what the AI remembered from Emma's check-in conversation.\n",
        "\n",
        "Memory Bank will have extracted structured facts about:\n",
        "- Guest identity and personal information\n",
        "- Dietary preferences and restrictions\n",
        "- Room location preferences\n",
        "- Temperature preferences\n",
        "- Special amenity requests\n",
        "\n",
        "Each memory is a concise, factual statement that can be used to personalize future interactions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P3DNtaGHsBNk"
      },
      "outputs": [],
      "source": [
        "# Display the generated memories using our helper function\n",
        "display_generated_memories(operation, client, \"Guest Preferences Extracted\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lz3fvHFusGDh"
      },
      "source": [
        "## Retrieve and use guest memories\n",
        "\n",
        "Now let's see how to retrieve these memories and use them to personalize the guest experience."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hsOCq4X4sILh"
      },
      "source": [
        "### Retrieve all memories for a guest\n",
        "\n",
        "Let's retrieve all the memories we just created for our guest. The simplest method is scope-based retrieval.\n",
        "\n",
        "A \"scope\" is a set of key-value pairs that defines a collection of memories. By providing `{\"user_id\": guest_id}`, we're asking for all memories that belong to this specific guest. This is like querying a database: \"Give me all records where user_id = 'guest_emma_1234'\"\n",
        "\n",
        "When a guest returns or calls the front desk, you want to quickly access everything you know about their preferences."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dmUM2EsKsf95"
      },
      "outputs": [],
      "source": [
        "print(f\"📚 Retrieving all memories for guest: {guest_id}...\\n\")\n",
        "\n",
        "# Simple retrieval - get all memories for this guest\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": guest_id}\n",
        ")\n",
        "\n",
        "all_memories = list(results)\n",
        "\n",
        "print(f\"✅ Found {len(all_memories)} memories for this guest:\\n\")\n",
        "\n",
        "for i, retrieved_memory in enumerate(all_memories, 1):\n",
        "    print(f\"   {i}. {retrieved_memory.memory.fact}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZerLO7Vfskq6"
      },
      "source": [
        "## Putting everything together\n",
        "\n",
        "Using Vertex AI Memory Bank, when Emma returns for another stay, we can instantly recall her preferences and provide exceptional, personalized service So, instead of asking Emma to repeat all her preferences, we proactively set up her room exactly how she likes it.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mJAnxc4HuESm"
      },
      "outputs": [],
      "source": [
        "print(\"=\" * 80)\n",
        "print(\"SCENARIO: Guest Returns for Second Stay\")\n",
        "print(\"=\" * 80)\n",
        "print(\"\\n📅 Two weeks later, Emma returns to the hotel...\\n\")\n",
        "\n",
        "# Simulate a return visit conversation\n",
        "print(\"👤 Hi! I'm Emma Chen, checking in again.\")\n",
        "print(\"🤖 Welcome back, Ms. Chen! Let me pull up your preferences...\\n\")\n",
        "\n",
        "# Retrieve Emma's memories\n",
        "print(\"🔍 Retrieving guest preferences from previous stay...\")\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": guest_id}\n",
        ")\n",
        "\n",
        "memories = list(results)\n",
        "\n",
        "print(f\"✅ Found {len(memories)} preferences on file\\n\")\n",
        "\n",
        "# Display how we'll use these memories\n",
        "print(\"🤖 Perfect! Based on your previous stay, I have:\")\n",
        "print()\n",
        "\n",
        "for memory in memories:\n",
        "    fact = memory.memory.fact\n",
        "\n",
        "    # Provide contextual responses based on memory content\n",
        "    if \"vegetarian\" in fact.lower() or \"lactose\" in fact.lower():\n",
        "        print(f\"   🥗 Dietary: {fact}\")\n",
        "        print(\"      → I've already notified the kitchen and breakfast staff\")\n",
        "    elif (\n",
        "        \"floor\" in fact.lower() or \"quiet\" in fact.lower() or \"elevator\" in fact.lower()\n",
        "    ):\n",
        "        print(f\"   🏨 Room: {fact}\")\n",
        "        print(\"      → I've assigned you the same corner room on floor 12\")\n",
        "    elif \"temperature\" in fact.lower() or \"68\" in fact:\n",
        "        print(f\"   🌡️  Climate: {fact}\")\n",
        "        print(\"      → Thermostat pre-set before your arrival\")\n",
        "    elif \"pillow\" in fact.lower() or \"white noise\" in fact.lower():\n",
        "        print(f\"   🛏️  Amenities: {fact}\")\n",
        "        print(\"      → Already placed in your room\")\n",
        "    else:\n",
        "        print(f\"   ℹ️  {fact}\")\n",
        "\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKmPdaPauKpb"
      },
      "source": [
        "## Advanced retrieval: Similarity Search\n",
        "\n",
        "In the previous examples, we retrieved **all** memories for a guest using scope-based retrieval. But what if you only want the most **relevant** memories for a specific question? This is where **similarity search** shines.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "21MWjWI3uMum"
      },
      "source": [
        "### Understanding the difference: Scope vs. Similarity\n",
        "\n",
        "Imagine Emma has stayed at your hotel 10 times and has 50+ memories stored. When the front desk needs to answer \"What are Emma's dietary restrictions?\", they don't need all 50 memories—just the 2-3 most relevant ones about food preferences.\n",
        "\n",
        "This is the idea that justifies similarity search in memory. To recap:\n",
        "\n",
        "- **Scope-based retrieval**: Returns ALL memories matching a scope (e.g., all memories for user_id = \"guest_emma_1234\")\n",
        "- **Similarity search**: Returns only the TOP K most relevant memories for a specific query using semantic similarity\n",
        "\n",
        "Let's see this in action by adding more conversation history for Emma."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p6YSXCT-uelO"
      },
      "source": [
        "#### Add more guest interactions\n",
        "\n",
        "Emma calls the front desk with questions about spa services and breakfast options. Let's record this conversation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Z6Xs85_uhjN"
      },
      "outputs": [],
      "source": [
        "print(\"📞 Emma calls the front desk with some questions...\\n\")\n",
        "\n",
        "# Additional conversation turns\n",
        "additional_conversation = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Hi, I'd like to book a spa treatment tomorrow. Do you have deep tissue massage? I prefer morning appointments, around 9 AM.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Absolutely! I can book you a deep tissue massage at 9 AM tomorrow. I've noted your preference for morning spa appointments.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Great! Also, what vegetarian options do you have for breakfast? I need dairy-free choices because of my lactose intolerance.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"We have several vegetarian and dairy-free options: oat milk smoothie bowls, avocado toast, fresh fruit platters, and coconut yogurt parfaits.\",\n",
        "    },\n",
        "]\n",
        "\n",
        "# Add these turns to the same session\n",
        "invocation_id_counter = len(guest_conversation)\n",
        "\n",
        "for turn in additional_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=session_name,\n",
        "        author=guest_id,\n",
        "        invocation_id=str(invocation_id_counter),\n",
        "        timestamp=datetime.datetime.now(tz=datetime.timezone.utc),\n",
        "        config={\n",
        "            \"content\": {\"role\": turn[\"role\"], \"parts\": [{\"text\": turn[\"message\"]}]}\n",
        "        },\n",
        "    )\n",
        "\n",
        "    invocation_id_counter += 1\n",
        "    icon = \"👤\" if turn[\"role\"] == \"user\" else \"🤖\"\n",
        "    print(f\"{icon} {turn['message']}\")\n",
        "\n",
        "print(\"\\n✅ Additional conversation added to session!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HC3ECbQLulZx"
      },
      "source": [
        "#### Generate new memories from extended conversation\n",
        "\n",
        "Now let's generate memories from this extended conversation. Memory Bank will analyze the entire session and extract new facts while consolidating them with existing memories.\n",
        "\n",
        "This demonstrates Memory Bank's consolidation capability—it won't create duplicate memories. Instead, it merges new information with what it already knows about Emma.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iDoHG2mVurR3"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🧠 Generating memories from extended conversation...\\n\")\n",
        "print(\"⏳ Analyzing the full conversation history and consolidating memories...\")\n",
        "\n",
        "# Generate additional memories\n",
        "operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine_name,\n",
        "    vertex_session_source={\"session\": session_name},\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Memory generation complete!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LGZviss5u51I"
      },
      "outputs": [],
      "source": [
        "# Display newly generated memories\n",
        "display_generated_memories(operation, client, \"New Memories Generated\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xcz-5Fd1vNAs"
      },
      "source": [
        "#### Compare scope-based vs. similarity-based retrieval\n",
        "\n",
        "Now let's compare both retrieval methods side-by-side to see the difference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v3EVb70WvOCa"
      },
      "outputs": [],
      "source": [
        "print(\"COMPARISON: Scope-Based vs. Similarity-Based Retrieval\")\n",
        "\n",
        "# Method 1: Scope-based retrieval (returns ALL memories)\n",
        "print(\"\\n📊 Method 1: Scope-Based Retrieval\")\n",
        "print(\"Query: Get all memories for Emma\")\n",
        "print(\"-\" * 80)\n",
        "\n",
        "all_results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": guest_id}\n",
        ")\n",
        "\n",
        "all_memories = list(all_results)\n",
        "print(f\"\\n✅ Found {len(all_memories)} total memories:\\n\")\n",
        "\n",
        "for i, mem in enumerate(all_memories, 1):\n",
        "    print(f\"   {i}. {mem.memory.fact}\")\n",
        "\n",
        "# Method 2: Similarity search (returns only RELEVANT memories)\n",
        "print(\"\\n🔍 Method 2: Similarity Search\")\n",
        "print(\"Query: What are Emma's dietary restrictions?\")\n",
        "print(\"-\" * 80)\n",
        "\n",
        "search_results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name,\n",
        "    scope={\"user_id\": guest_id},\n",
        "    similarity_search_params={\n",
        "        \"search_query\": \"What are Emma's dietary restrictions?\",\n",
        "        \"top_k\": 3,  # Get top 3 most relevant memories\n",
        "    },\n",
        ")\n",
        "\n",
        "relevant_memories = list(search_results)\n",
        "print(f\"\\n✅ Found {len(relevant_memories)} relevant memories:\\n\")\n",
        "\n",
        "for i, mem in enumerate(relevant_memories, 1):\n",
        "    distance = mem.distance if hasattr(mem, \"distance\") else \"N/A\"\n",
        "    print(f\"   {i}. {mem.memory.fact}\")\n",
        "    print(f\"      → Relevance distance: {distance} (lower = more relevant)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mT5nT1anvaog"
      },
      "source": [
        "### Use similarity search for specific questions\n",
        "\n",
        "Let's test similarity search with several specific front desk questions.\n",
        "\n",
        "In real-world scenarios, hotel staff need quick, targeted answers. Similarity search makes this efficient by retrieving only what matters for each specific question.\n",
        "\n",
        "**Understanding Similarity Search Parameters**\n",
        "\n",
        "| Parameter | Description | Example Value |\n",
        "|-----------|-------------|---------------|\n",
        "| `search_query` | The natural language question or topic | \"What are the guest's room preferences?\" |\n",
        "| `top_k` | Maximum number of results to return | 3 (returns top 3 most relevant) |\n",
        "| `distance` | Distance between query and memory embeddings | Automatically calculated (lower = more similar) |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UtmK2kQjvZHz"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🏨 Testing Similarity Search with Common Front Desk Questions\\n\")\n",
        "\n",
        "# Test with various specific queries\n",
        "search_queries = [\n",
        "    \"What are Emma's room location preferences?\",\n",
        "    \"What are Emma's dietary restrictions?\",\n",
        "    \"What spa services does Emma prefer?\",\n",
        "    \"What temperature does Emma like in her room?\",\n",
        "]\n",
        "\n",
        "for query in search_queries:\n",
        "    print(f'\\n❓ Query: \"{query}\"')\n",
        "    print(\"-\" * 80)\n",
        "\n",
        "    # Perform similarity search\n",
        "    results = client.agent_engines.memories.retrieve(\n",
        "        name=agent_engine_name,\n",
        "        scope={\"user_id\": guest_id},\n",
        "        similarity_search_params={\n",
        "            \"search_query\": query,\n",
        "            \"top_k\": 2,  # Get top 2 most relevant memories\n",
        "        },\n",
        "    )\n",
        "\n",
        "    memories = list(results)\n",
        "\n",
        "    if memories:\n",
        "        print(\"\\n🎯 Most Relevant Memories:\")\n",
        "        for i, mem in enumerate(memories, 1):\n",
        "            distance = mem.distance if hasattr(mem, \"distance\") else \"N/A\"\n",
        "            print(f\"   {i}. {mem.memory.fact}\")\n",
        "            print(\n",
        "                f\"      (Distance: {distance:.4f})\"\n",
        "                if isinstance(distance, (int, float))\n",
        "                else f\"      (Distance: {distance})\"\n",
        "            )\n",
        "    else:\n",
        "        print(\"   No relevant memories found\")\n",
        "\n",
        "    print()\n",
        "\n",
        "print(\"✅ Similarity search enables fast, targeted information retrieval!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aHwtSN7Nv0Zi"
      },
      "source": [
        "### Best practices: When to use each retrieval method\n",
        "\n",
        "**Use Scope-Based Retrieval When**:\n",
        "- You need a complete profile of the guest\n",
        "- You're displaying all preferences in a dashboard\n",
        "- The number of memories is small\n",
        "- You want to ensure nothing is missed\n",
        "\n",
        "**Use Similarity Search When**:\n",
        "- You're answering a specific question\n",
        "- The guest has many memories\n",
        "- You need fast, targeted responses\n",
        "- You're building a conversational agent that needs context-aware responses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring charges to your Google Cloud account for the resources used in this tutorial, delete the Agent Engine we created.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YgRHdO2XwfUX"
      },
      "outputs": [],
      "source": [
        "print(\"🧹 Cleaning up resources...\\n\")\n",
        "\n",
        "delete_agent_engine = True\n",
        "\n",
        "if delete_agent_engine:\n",
        "    # Delete the agent engine and all its memories\n",
        "    client.agent_engines.delete(name=agent_engine_name, force=True)\n",
        "\n",
        "    print(\"✅ Agent Engine deleted successfully!\")\n",
        "    print(\"   All guest memories have been removed\")\n",
        "else:\n",
        "    print(\"⏭️  Skipping cleanup - Agent Engine will remain active\")\n",
        "    print(f\"   Resource: {agent_engine_name}\")\n",
        "    print(\"\\n⚠️  Remember to delete this resource later to avoid charges\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rx2JHxpkwJTW"
      },
      "source": [
        "## Congratulations!\n",
        "\n",
        "You've completed the \"Get Started with Vertex AI Memory Bank\" tutorial!\n",
        "\n",
        "You now have the foundational knowledge to build context-aware, personalized AI agents that remember user preferences across sessions. The hotel concierge scenario you built here can be adapted to countless use cases where personalization and memory matter.\n",
        "\n",
        "**What's Next?**\n",
        "- Explore advanced Memory Bank features in our intermediate tutorials\n",
        "- Check out the [Memory Bank documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank)\n",
        "- Join the [Google Cloud AI community](https://discuss.google.dev/c/google-cloud/14) to share your projects\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_memory_bank.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
