{
  "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": [
        "# Governance 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_governance.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_governance.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_governance.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_governance.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",
        "**Share to:**\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_governance.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_governance.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_governance.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_governance.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_governance.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| 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 demonstrates how to build **customer support systems** using **Vertex AI Memory Bank's governance features**. You will learn to implement data retention policies, maintain audit trails, and handle regulatory compliance requirements while building an intelligent customer support agent.\n",
        "\n",
        "By the end of this tutorial, you will understand how to:\n",
        "\n",
        "* **Manage Basic Memory Operations**: Create, update, delete, and retrieve memories manually\n",
        "* **Generate Memories from Conversations**: Extract structured facts from customer interactions\n",
        "* **Configure Granular TTL (Time-To-Live)**: Set different retention periods (30/90/365 days) for different data types to comply with data retention regulations\n",
        "* **Filter Memories by Topic**: Retrieve specific categories of data\n",
        "* **Track Memory Revisions**: Maintain complete audit trails showing how customer data evolved over time\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "### Install Vertex AI SDK\n",
        "\n",
        "First, install the Vertex AI SDK. We're specifying version 1.111.0 or higher to ensure we have all the latest Memory Bank features.\n",
        "\n",
        "**Note**: This will install the SDK. Colab may prompt you to restart the runtime after installation. This is expected behavior - simply click \"Restart Runtime\" when prompted."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform>=1.123.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment\n",
        "\n",
        "If you are running this notebook in **Google Colab**, run the cell below to authenticate your account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "LOCATION = \"us-central1\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "if not LOCATION:\n",
        "    LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\")\n",
        "\n",
        "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "We're importing standard Python libraries and several specialized classes from the Vertex AI SDK.\n"
      ]
    },
    {
      "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 for governance\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": "DQFPukme7p07"
      },
      "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": "zMEtiWs27s48"
      },
      "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",
        "            # Skip deleted memories - we only want to show active ones\n",
        "            if gen_memory.action != \"DELETED\" and gen_memory.memory:\n",
        "                try:\n",
        "                    # Fetch full memory details including topics and metadata\n",
        "                    full_memory = client.agent_engines.memories.get(\n",
        "                        name=gen_memory.memory.name\n",
        "                    )\n",
        "                    # Visual indicators: new vs updated memories\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": "at-z4pQE7xuS"
      },
      "source": [
        "## Create a basic Memory Bank on Vertex AI Agent Engine\n",
        "\n",
        "Let's start with the fundamentals by creating a simple Memory Bank with basic configuration. We'll add advanced governance features later. The AgentEngine resource acts as the container for your Memory Bank instance. Let's create one with a minimal configuration to get started.\n",
        "\n",
        "**Basic Configuration Components**:\n",
        "\n",
        "1. **Similarity Search Config**: Specifies the embedding model for semantic memory retrieval\n",
        "   - We're using `text-embedding-005`, which is excellent for English conversations\n",
        "\n",
        "2. **Generation Config**: Defines the LLM that extracts memories from conversations\n",
        "   - We're using `gemini-2.5-flash`, a fast and capable model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "te9P-JUa756L"
      },
      "outputs": [],
      "source": [
        "print(\"🧠 Creating basic Memory Bank configuration...\\n\")\n",
        "\n",
        "# Create minimal Memory Bank configuration\n",
        "basic_memory_config = MemoryBankConfig(\n",
        "    # Embedding model 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",
        "    # LLM for memory extraction\n",
        "    generation_config=GenerationConfig(\n",
        "        model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(\"✅ Basic Memory Bank configuration created!\")\n",
        "print(\"   Embedding model: text-embedding-005\")\n",
        "print(\"   Generation model: gemini-2.5-flash\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FGOu_Na47-Ky"
      },
      "source": [
        "Now let's create the Agent Engine with this basic configuration.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0nijEqyq8Ata"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🛠️ Creating Agent Engine with basic Memory Bank...\\n\")\n",
        "print(\"⏳ Provisioning backend infrastructure...\")\n",
        "\n",
        "# Create the Agent Engine\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": "9DxjCyCq8Le4"
      },
      "source": [
        "## Basic Memory Operations\n",
        "\n",
        "Now that we have a Memory Bank, let's learn the fundamental CRUD (Create, Read, Update, Delete) operations. These are the building blocks for all memory management."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1GVr94288Ns6"
      },
      "source": [
        "### Manually create a memory\n",
        "\n",
        "Let's create our first memory manually. This is useful when you need to record information that wasn't captured in a conversation.\n",
        "\n",
        "**Note:** Create method is quite limited. You don't get the benefit of consolidation. If you want consolidation, use GenerateMemories with `direct_memories_source`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pp6_0Exu8Qka"
      },
      "outputs": [],
      "source": [
        "print(\"🆕 Creating your first memory manually\\n\")\n",
        "\n",
        "# Generate a unique customer ID for this tutorial\n",
        "customer_id = \"customer_alex_\" + str(uuid.uuid4())[:4]\n",
        "\n",
        "# Create a memory manually\n",
        "first_memory = client.agent_engines.memories.create(\n",
        "    name=agent_engine_name,\n",
        "    fact=\"Customer prefers to be contacted via email for all communications\",\n",
        "    scope={\"user_id\": customer_id},\n",
        ")\n",
        "\n",
        "print(\"✅ Memory created successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-gvhEdhV9NCQ"
      },
      "source": [
        "### Retrieve memories for a customer\n",
        "\n",
        "Now let's retrieve the memory we just created. This demonstrates how to fetch all memories associated with a specific customer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xvuQaFGL85t0"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🔍 Retrieving all memories for the customer\\n\")\n",
        "\n",
        "# Retrieve all memories for this customer\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": customer_id}\n",
        ")\n",
        "\n",
        "all_memories = list(results)\n",
        "\n",
        "print(f\"✅ Found {len(all_memories)} memory for customer {customer_id}:\\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": "8n1jiBJj9Z4y"
      },
      "source": [
        "### Update an existing memory\n",
        "\n",
        "Customer preferences change over time. Let's update our memory to reflect a new communication preference.\n",
        "\n",
        "To update memories, you can generate new memories from updated information. Memory Bank intelligently recognizes existing memories and updates them, creating revision history automatically."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8QOHNSCX9aH8"
      },
      "outputs": [],
      "source": [
        "print(\"\\n✏️ Updating the customer's communication preference\\n\")\n",
        "\n",
        "first_memory = all_memories[0]\n",
        "\n",
        "print(f\"📋 Current preference: '{first_memory.memory.fact}'\")\n",
        "print(\n",
        "    \"\\n💡 Customer calls to update: 'I prefer SMS for urgent issues, email for everything else'\\n\"\n",
        ")\n",
        "\n",
        "# Update the memory by generating from new information\n",
        "# Memory Bank will recognize this updates the existing memory\n",
        "update_operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine_name,\n",
        "    scope={\"user_id\": customer_id},\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"parts\": [\n",
        "                        {\n",
        "                            \"text\": \"Customer prefers SMS for urgent issues and email for all other communications\"\n",
        "                        }\n",
        "                    ]\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "# Get the updated memory details\n",
        "if update_operation.response and update_operation.response.generated_memories:\n",
        "    updated_memory_ref = update_operation.response.generated_memories[0]\n",
        "    updated_memory = client.agent_engines.memories.get(\n",
        "        name=updated_memory_ref.memory.name\n",
        "    )\n",
        "\n",
        "    print(\"✅ Memory updated successfully!\")\n",
        "    print(f\"   📝 New fact: {updated_memory.fact}\")\n",
        "    print(f\"   🔄 Updated: {updated_memory.update_time}\")\n",
        "    print(f\"   🔄 Action: {updated_memory_ref.action}\")\n",
        "\n",
        "    if (\n",
        "        hasattr(updated_memory_ref, \"previous_revision\")\n",
        "        and updated_memory_ref.previous_revision\n",
        "    ):\n",
        "        print(\n",
        "            f\"   📜 Previous revision preserved: {updated_memory_ref.previous_revision.split('/')[-1]}\"\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7I7b5xzyL6YL"
      },
      "source": [
        "### Create additional memories\n",
        "\n",
        "Let's create a few more memories to make our examples more realistic. We'll add different types of customer information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pTj2fDp1Pgn_"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📝 Adding more customer information\\n\")\n",
        "\n",
        "# Create several more memories\n",
        "additional_memories = [\n",
        "    \"Customer's account type is Enterprise Pro with priority support\",\n",
        "    \"Customer reported issue with API authentication returning 401 errors\",\n",
        "    \"Customer prefers morning appointments between 9 AM and 11 AM EST\",\n",
        "]\n",
        "\n",
        "created_memories = []\n",
        "\n",
        "for fact in additional_memories:\n",
        "    memory = client.agent_engines.memories.create(\n",
        "        name=agent_engine_name, fact=fact, scope={\"user_id\": customer_id}\n",
        "    )\n",
        "    created_memories.append(memory)\n",
        "    print(f\"   ✅ Created: {fact}\")\n",
        "\n",
        "print(f\"\\n✅ Created {len(created_memories)} additional memories\")\n",
        "print(f\"   Total memories for {customer_id}: {len(created_memories) + 1}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V1kQJe97P0tB"
      },
      "source": [
        "### Retrieve and display all memories\n",
        "\n",
        "Let's see everything we've stored for this customer so far."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lWcapirHP7Mf"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📚 Complete customer profile\\n\")\n",
        "\n",
        "# Retrieve all memories\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": customer_id}\n",
        ")\n",
        "\n",
        "all_customer_memories = list(results)\n",
        "\n",
        "print(f\"✅ Found {len(all_customer_memories)} total memories for {customer_id}:\\n\")\n",
        "\n",
        "for i, mem in enumerate(all_customer_memories, 1):\n",
        "    print(f\"   {i}. {mem.memory.fact}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2LBQ-LBrQHNS"
      },
      "source": [
        "### Delete a specific memory\n",
        "\n",
        "Let's demonstrate how to delete a specific memory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XCXMO_IPQHf3"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🗑️ Deleting a specific memory\\n\")\n",
        "\n",
        "# Let's delete the API issue memory\n",
        "memory_to_delete = all_customer_memories[2]\n",
        "\n",
        "print(\"🔍 Memory to delete:\")\n",
        "print(f\"   📝 {memory_to_delete.memory.fact}\")\n",
        "print(f\"   🆔 ID: {memory_to_delete.memory.name.split('/')[-1]}\")\n",
        "\n",
        "print(\"\\n⚠️ Confirming deletion...\")\n",
        "\n",
        "# Delete the memory\n",
        "client.agent_engines.memories.delete(name=memory_to_delete.memory.name)\n",
        "\n",
        "print(\"\\n✅ Memory deleted successfully!\")\n",
        "\n",
        "# Verify deletion\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": customer_id}\n",
        ")\n",
        "\n",
        "remaining_memories = list(results)\n",
        "\n",
        "print(\n",
        "    f\"\\n📊 Remaining memories: {len(remaining_memories)} (was {len(all_customer_memories)})\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZkPhsIISt7J"
      },
      "source": [
        "## Generate memories from conversations\n",
        "\n",
        "So far we've created memories manually. Now let's see how Memory Bank can automatically extract structured facts from natural conversations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFpWWwucS0aj"
      },
      "source": [
        "### Create a session for customer support\n",
        "\n",
        "A **Session** represents a single conversation between a customer and support. Let's create one and add a realistic support conversation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kxu3oJhtSxXj"
      },
      "outputs": [],
      "source": [
        "print(\"\\n💬 Creating customer support conversation\\n\")\n",
        "\n",
        "# Create session for this customer\n",
        "session = client.agent_engines.sessions.create(\n",
        "    name=agent_engine_name,\n",
        "    user_id=customer_id,\n",
        "    config={\"display_name\": f\"Support session for {customer_id}\"},\n",
        ")\n",
        "\n",
        "session_name = session.response.name\n",
        "\n",
        "print(\"✅ Session created successfully!\")\n",
        "print(f\"   Session: {session_name}\")\n",
        "print(f\"   Customer ID: {customer_id}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaZQe1l7S6mb"
      },
      "source": [
        "### Add support conversation to the session\n",
        "\n",
        "Here's a realistic customer support conversation containing important information that should be remembered."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eIIEZ1dUS4Ao"
      },
      "outputs": [],
      "source": [
        "# Define realistic customer support conversation\n",
        "support_conversation = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Hi, I'm having issues with my billing - I was charged twice for my enterprise subscription last month.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Hello! I apologize for the billing issue. Let me look into your enterprise subscription charges right away.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"I'm on the Enterprise Pro plan at $299/month. The duplicate charge was on March 15th. My company is DataTech Solutions.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Thank you for those details. I see your Enterprise Pro subscription. I'll investigate the duplicate charge from March 15th.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Also, I prefer email communication over phone calls. I'm in EST timezone, so please don't call before 9 AM my time.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Noted! I've recorded your preference for email communication and your EST timezone availability after 9 AM.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"I also opened ticket #SUP-2024-0847 last week about API authentication failures. Still not resolved.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"I see ticket SUP-2024-0847 regarding API authentication. Let me prioritize getting that resolved for you.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"The API returns 401 errors when using our production API key. We're using the REST API v2.1 endpoint.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Thank you for the technical details. I'll escalate the API authentication issue with our engineering team and get back to you within 24 hours via email.\",\n",
        "    },\n",
        "]\n",
        "\n",
        "print(\"\\n⬆️ Adding conversation to session...\\n\")\n",
        "\n",
        "invocation_id = 0\n",
        "\n",
        "# Add each conversation turn to the session with timestamps\n",
        "for turn in support_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=session_name,\n",
        "        author=customer_id,\n",
        "        invocation_id=str(invocation_id),\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 += 1\n",
        "    # Display conversation with visual indicators\n",
        "    icon = \"👤\" if turn[\"role\"] == \"user\" else \"🤖\"\n",
        "    print(f\"{icon} {turn['message']}\")\n",
        "\n",
        "print(\"\\n✅ Conversation added to session!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0r6rGzx5TZDl"
      },
      "source": [
        "### Generate memories from the conversation\n",
        "\n",
        "Memory Bank analyzes the entire conversation and automatically extracts structured facts.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f8LRkKsnTU6u"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🧠 Generating memories from conversation\\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(\"✅ Memory generation complete!\")\n",
        "\n",
        "# Display the generated memories\n",
        "display_generated_memories(operation, client, \"Automatically Extracted Memories\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E-sIlPIpTnv6"
      },
      "source": [
        "### View all memories (manual + generated)\n",
        "\n",
        "Let's see the complete customer profile now - combining our manually created memories with the auto-generated ones."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gz6X9fZVTgZQ"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📊 Complete Customer Profile (Manual + Auto-Generated)\\n\")\n",
        "\n",
        "# Retrieve all memories for this customer\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=agent_engine_name, scope={\"user_id\": customer_id}\n",
        ")\n",
        "\n",
        "complete_profile = list(results)\n",
        "\n",
        "print(f\"✅ Total memories for {customer_id}: {len(complete_profile)}\\n\")\n",
        "\n",
        "for i, mem in enumerate(complete_profile, 1):\n",
        "    print(f\"   {i}. {mem.memory.fact}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iCVdOPxGYW0A"
      },
      "source": [
        "## Granular TTL (Time-To-Live)\n",
        "\n",
        "Now let's add automatic data retention policies using **Granular TTL**. This ensures data is automatically deleted after specified retention periods.\n",
        "\n",
        "Granular TTL lets you specify three different retention periods:\n",
        "\n",
        "1. **create_ttl**: How long manually created memories persist (30 days)\n",
        "2. **generate_created_ttl**: How long newly generated memories persist (90 days)\n",
        "3. **generate_updated_ttl**: How long updated memories persist (365 days)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UdLmueLic06J"
      },
      "source": [
        "### Define granular TTL configuration\n",
        "\n",
        "Let's configure our TTL policy for automatic data lifecycle management:\n",
        "\n",
        "**TTL Configuration Parameters**:\n",
        "\n",
        "| Parameter | Duration | Seconds | Use Case |\n",
        "|-----------|----------|---------|----------|\n",
        "| `create_ttl` | 30 days | 2592000s | Manually created operational data |\n",
        "| `generate_created_ttl` | 90 days | 7776000s | Newly extracted support memories |\n",
        "| `generate_updated_ttl` | 365 days | 31536000s | Consolidated, validated account data |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OhCiMGHAZ74m"
      },
      "outputs": [],
      "source": [
        "print(\"\\n⏱️ Configuring granular TTL for automatic data retention\\n\")\n",
        "\n",
        "# Define granular TTL configuration with different retention periods\n",
        "ttl_config = TtlConfig(\n",
        "    granular_ttl_config=GranularTtlConfig(\n",
        "        create_ttl=\"2592000s\",  # 30 days  - Manual creates (temporary notes)\n",
        "        generate_created_ttl=\"7776000s\",  # 90 days  - New generated memories\n",
        "        generate_updated_ttl=\"31536000s\",  # 365 days - Updated/consolidated memories\n",
        "    )\n",
        ")\n",
        "\n",
        "print(\"✅ Granular TTL configuration created!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lF4-sMEjc-sx"
      },
      "source": [
        "### Update Agent Engine with TTL configuration\n",
        "\n",
        "Now let's update our existing Agent Engine with the TTL governance configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3hBQDkO4c5Rk"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🔧 Updating Agent Engine with TTL governance\\n\")\n",
        "\n",
        "# Create enhanced configuration with TTL\n",
        "governance_memory_config = MemoryBankConfig(\n",
        "    # Keep existing models\n",
        "    similarity_search_config=SimilaritySearchConfig(\n",
        "        embedding_model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005\"\n",
        "    ),\n",
        "    generation_config=GenerationConfig(\n",
        "        model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "    ),\n",
        "    # Add TTL governance\n",
        "    ttl_config=ttl_config,\n",
        ")\n",
        "\n",
        "# Update the existing engine\n",
        "updated_engine = client.agent_engines.update(\n",
        "    name=agent_engine_name,\n",
        "    config={\"context_spec\": {\"memory_bank_config\": governance_memory_config}},\n",
        ")\n",
        "\n",
        "print(\"✅ Agent Engine updated with TTL!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_sTi6f-BdKgN"
      },
      "source": [
        "### Test TTL with new memories\n",
        "\n",
        "Let's create a new memory and see the expiration time based on our TTL configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QSm0g7undF2h"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🧪 Testing TTL with new memory creation\\n\")\n",
        "\n",
        "# Create a new memory with manual creation (30-day TTL)\n",
        "test_memory = client.agent_engines.memories.create(\n",
        "    name=agent_engine_name,\n",
        "    fact=\"Customer requested demo of new analytics dashboard feature\",\n",
        "    scope={\"user_id\": customer_id},\n",
        ")\n",
        "\n",
        "print(\"✅ Memory created with TTL!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DqWvGt-EfXdB"
      },
      "source": [
        "## Topic-Based Filtering with Default Managed Topics\n",
        "\n",
        "One of the most powerful features for organizing memories is topic-based filtering. Even without custom topics, Memory Bank automatically categorizes memories using **default managed topics**.\n",
        "\n",
        "**Default Managed Topics**:\n",
        "\n",
        "| Topic | Description | Example |\n",
        "|-------|-------------|---------|\n",
        "| `USER_PERSONAL_INFO` | Personal details about the user | \"Customer name is Alex Chen, works at DataTech Solutions\" |\n",
        "| `USER_PREFERENCES` | Stated preferences and likes/dislikes | \"Customer prefers email communication over phone\" |\n",
        "| `KEY_CONVERSATION_DETAILS` | Important outcomes or milestones | \"Ticket #SUP-2024-0847 created for API authentication issue\" |\n",
        "| `EXPLICIT_INSTRUCTIONS` | Direct remember/forget requests | \"Customer asked to remember account number\" |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCb0p1j-fcin"
      },
      "source": [
        "### Inspect memories with topic labels\n",
        "\n",
        "Let's inspect memories in detail to see their automatically attached topic labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jN6vhfWQfX1n"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🔍 Inspecting memories with automatic topic labels\\n\")\n",
        "\n",
        "# Get all memories for detailed inspection\n",
        "pager = client.agent_engines.memories.list(name=agent_engine_name)\n",
        "sample_memories = list(pager)[:5]  # First 5 memories\n",
        "\n",
        "if sample_memories:\n",
        "    print(f\"Showing {len(sample_memories)} memories with full details:\\n\")\n",
        "\n",
        "    for i, memory in enumerate(sample_memories, 1):\n",
        "        # Get complete memory details\n",
        "        full_memory = client.agent_engines.memories.get(name=memory.name)\n",
        "\n",
        "        print(f\"\\n📌 Memory {i}:\")\n",
        "        print(f\"   Fact: {full_memory.fact}\")\n",
        "        print(f\"   Scope: {full_memory.scope}\")\n",
        "        print(f\"   Created: {full_memory.create_time}\")\n",
        "        print(f\"   Expires: {full_memory.expire_time}\")\n",
        "\n",
        "        # Show automatically attached topics\n",
        "        if hasattr(full_memory, \"topics\") and full_memory.topics:\n",
        "            topics = [\n",
        "                topic.managed_memory_topic\n",
        "                if hasattr(topic, \"managed_memory_topic\")\n",
        "                else str(topic)\n",
        "                for topic in full_memory.topics\n",
        "            ]\n",
        "            print(f\"   Topics: {', '.join(str(t) for t in topics)}\")\n",
        "        else:\n",
        "            print(\"   Topics: None\")\n",
        "else:\n",
        "    print(\"No memories found\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AQb4o4FBglrh"
      },
      "source": [
        "### Filter memories by specific topic\n",
        "\n",
        "Now let's demonstrate topic-based filtering. This is useful for retrieving specific categories of information.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ah1E1YuqgpRN"
      },
      "outputs": [],
      "source": [
        "def get_memories_by_topic(engine_name, topic_enum_value):\n",
        "    \"\"\"Retrieve memories filtered by a specific managed topic.\n",
        "\n",
        "    Args:\n",
        "        engine_name: The Agent Engine resource name\n",
        "        topic_enum_value: The topic to filter by (e.g., 'USER_PERSONAL_INFO')\n",
        "\n",
        "    Returns:\n",
        "        List of memories matching the topic\n",
        "    \"\"\"\n",
        "    # Get all memories\n",
        "    all_memories_pager = client.agent_engines.memories.list(name=engine_name)\n",
        "    filtered_memories = []\n",
        "\n",
        "    # Filter by topic\n",
        "    for memory in all_memories_pager:\n",
        "        full_memory = client.agent_engines.memories.get(name=memory.name)\n",
        "        topics = (\n",
        "            full_memory.topics\n",
        "            if hasattr(full_memory, \"topics\") and full_memory.topics\n",
        "            else []\n",
        "        )\n",
        "\n",
        "        for topic in topics:\n",
        "            if hasattr(topic, \"managed_memory_topic\"):\n",
        "                topic_value = str(topic.managed_memory_topic)\n",
        "                if topic_enum_value in topic_value:\n",
        "                    filtered_memories.append(full_memory)\n",
        "                    break\n",
        "\n",
        "    return filtered_memories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aVtUbFYBgrMP"
      },
      "source": [
        "Now let's demonstrate filtering for each default managed topic.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VnZCu5fafiv2"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📊 Demonstrating topic-based filtering\\n\")\n",
        "\n",
        "# List of default managed topics to filter by\n",
        "topics_to_filter = [\n",
        "    (\"USER_PERSONAL_INFO\", \"Personal Information: Name, company, account details\"),\n",
        "    (\"USER_PREFERENCES\", \"Preferences: Communication channels, support preferences\"),\n",
        "    (\"KEY_CONVERSATION_DETAILS\", \"Key Details: Ticket numbers, important outcomes\"),\n",
        "]\n",
        "\n",
        "for topic_label, description in topics_to_filter:\n",
        "    print(f\"\\n🔍 {description}\")\n",
        "    print(f\"   Filtering by topic: {topic_label}\\n\")\n",
        "\n",
        "    # Get memories for this topic\n",
        "    topic_memories = get_memories_by_topic(agent_engine_name, topic_label)\n",
        "\n",
        "    if topic_memories:\n",
        "        print(f\"   ✅ Found {len(topic_memories)} memory/memories:\\n\")\n",
        "        for i, memory in enumerate(topic_memories, 1):\n",
        "            print(f\"      {i}. {memory.fact}\")\n",
        "            if memory.expire_time:\n",
        "                print(f\"         ⏱️ Expires: {memory.expire_time}\")\n",
        "    else:\n",
        "        print(f\"   📭 No memories found with topic '{topic_label}'\")\n",
        "\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h6Fg-1O2hAlv"
      },
      "source": [
        "## Memory Revisions: Complete Audit Trails\n",
        "\n",
        "Memory Bank's **revision history** provides a complete, immutable record of how data evolved over time. This is essential for compliance, debugging, and understanding data changes.\n",
        "\n",
        "Each time a memory is created, updated, or deleted, Memory Bank creates a **revision snapshot**. Think of it as version control for customer data.\n",
        "\n",
        "**Revision Actions**:\n",
        "- **CREATED**: First time the memory was generated\n",
        "- **UPDATED**: Memory was modified (e.g., through consolidation or manual update)\n",
        "- **DELETED**: Memory was removed (manually or via TTL expiration)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n-jDy1QShN1W"
      },
      "source": [
        "### List revision history for a memory\n",
        "\n",
        "Let's select one of our memories and examine its complete revision history."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ct5gQEeWhTYO"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📋 Viewing Complete Revision History\\n\")\n",
        "\n",
        "# Get a memory with revision history\n",
        "all_memories_pager = client.agent_engines.memories.list(name=agent_engine_name)\n",
        "memory_with_revisions = None\n",
        "\n",
        "# Find first memory for demonstration\n",
        "for memory in all_memories_pager:\n",
        "    full_memory = client.agent_engines.memories.get(name=memory.name)\n",
        "    memory_with_revisions = full_memory\n",
        "    break\n",
        "\n",
        "if memory_with_revisions:\n",
        "    print(f\"📌 Memory Selected: {memory_with_revisions.fact}\\n\")\n",
        "    print(\"📜 Revision History:\")\n",
        "\n",
        "    try:\n",
        "        # List all revisions for this memory\n",
        "        revisions_pager = client.agent_engines.memories.revisions.list(\n",
        "            name=memory_with_revisions.name\n",
        "        )\n",
        "\n",
        "        revisions = list(revisions_pager)\n",
        "\n",
        "        if revisions:\n",
        "            for i, revision in enumerate(revisions, 1):\n",
        "                print(f\"\\n📜 Revision {i}:\")\n",
        "                print(f\"   🆔 Revision ID: {revision.name.split('/')[-1]}\")\n",
        "                print(f\"   ⏱️ Timestamp: {revision.create_time}\")\n",
        "                print(f\"   📝 Fact: {revision.fact}\")\n",
        "                if hasattr(revision, \"expire_time\") and revision.expire_time:\n",
        "                    print(f\"   🔚 Expires: {revision.expire_time}\")\n",
        "        else:\n",
        "            print(\"\\n   📭 No revision history available yet\")\n",
        "            print(\"   💡 Revisions accumulate as memories are updated or consolidated\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\n   ⚠️ Note: {e}\")\n",
        "else:\n",
        "    print(\"📭 No memories available for revision inspection\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EzZrECEHhhuF"
      },
      "source": [
        "### Retrieve a specific revision (Time-Travel Query)\n",
        "\n",
        "One of the most powerful compliance features: retrieving the exact state of a memory at a specific point in time. Think about a data access request asks \"What data did you have about me on March 20th?\" With revisions, you can answer this precisely."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BnHTL3iehbF4"
      },
      "outputs": [],
      "source": [
        "print(\"\\n⏰ Time-Travel Query: Retrieving Specific Revision\\n\")\n",
        "\n",
        "if memory_with_revisions and revisions:\n",
        "    try:\n",
        "        if revisions:\n",
        "            # Get the first revision (earliest version)\n",
        "            target_revision = revisions[0]\n",
        "            revision_id = target_revision.name.split(\"/\")[-1]\n",
        "\n",
        "            print(f\"🎯 Retrieving revision: {revision_id}\")\n",
        "            print(f\"   ⏱️ Timestamp: {target_revision.create_time}\\n\")\n",
        "\n",
        "            # Fetch the specific revision\n",
        "            specific_revision = client.agent_engines.memories.revisions.get(\n",
        "                name=target_revision.name\n",
        "            )\n",
        "\n",
        "            print(\"📋 Revision Details:\")\n",
        "            print(f\"   📝 Memory Fact: {specific_revision.fact}\")\n",
        "            print(f\"   🛠️ Created: {specific_revision.create_time}\")\n",
        "            if (\n",
        "                hasattr(specific_revision, \"expire_time\")\n",
        "                and specific_revision.expire_time\n",
        "            ):\n",
        "                print(f\"   🔚 Expires: {specific_revision.expire_time}\")\n",
        "        else:\n",
        "            print(\"📭 No revisions available for time-travel query\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"⚠️ Could not retrieve revision: {e}\")\n",
        "else:\n",
        "    print(\"📭 No memory available for time-travel demonstration\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0KSjGjCrlQZ3"
      },
      "source": [
        "### Rollback a Memory to a Previous Revision\n",
        "\n",
        "Now that we've seen how to view and compare revisions, let's explore one of the most powerful governance features: the ability to **rollback** a memory to a\n",
        "previous state.\n",
        "\n",
        "**When to use rollback**:\n",
        "- A memory was incorrectly updated during consolidation\n",
        "- You need to revert to a known-good state for compliance\n",
        "- An incorrect fact needs to be corrected quickly\n",
        "\n",
        "Let's demonstrate this with a new memory."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-IQN39KSm314"
      },
      "source": [
        "#### Create the initial memory with revision labels\n",
        "\n",
        "We'll create a memory and tag it with revision labels to track its source and verification status.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OZx6OT-HlYrg"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🎯 Create initial memory with revision tracking\\n\")\n",
        "\n",
        "# Create a new customer memory with revision labels for tracking\n",
        "rollback_customer_id = \"customer_sarah_\" + str(uuid.uuid4())[:4]\n",
        "\n",
        "initial_operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine_name,\n",
        "    scope={\"user_id\": rollback_customer_id},\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"parts\": [\n",
        "                        {\n",
        "                            \"text\": \"Customer prefers email communication only, no phone calls\"\n",
        "                        }\n",
        "                    ]\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        "    config={\n",
        "        \"wait_for_completion\": True,\n",
        "        \"revision_labels\": {\"data_source\": \"initial_preference\", \"verified\": \"true\"},\n",
        "    },\n",
        ")\n",
        "\n",
        "if initial_operation.response and initial_operation.response.generated_memories:\n",
        "    rollback_memory = initial_operation.response.generated_memories[0].memory\n",
        "    rollback_memory_name = rollback_memory.name\n",
        "\n",
        "    # Get full memory details\n",
        "    full_memory = client.agent_engines.memories.get(name=rollback_memory_name)\n",
        "\n",
        "    print(\"✅ Initial memory created!\")\n",
        "    print(f\"   📝 Fact: {full_memory.fact}\")\n",
        "    print(f\"   🆔 Memory ID: {rollback_memory_name.split('/')[-1]}\")\n",
        "else:\n",
        "    print(\"⚠️ Failed to create initial memory\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SpEPUPT5nK7N"
      },
      "source": [
        "#### Simulate an incorrect update\n",
        "\n",
        "Let's update the memory with incorrect information (simulating a misinterpretation during consolidation). Notice the `verified: \"false\"` label.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P88MwBf_lld3"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🎯 Update the memory (simulating an incorrect update)\\n\")\n",
        "\n",
        "# Simulate an incorrect update - customer preference misinterpreted\n",
        "update_operation = client.agent_engines.memories.generate(\n",
        "    name=agent_engine_name,\n",
        "    scope={\"user_id\": rollback_customer_id},\n",
        "    direct_contents_source={\n",
        "        \"events\": [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"parts\": [\n",
        "                        {\"text\": \"Customer prefers phone calls for all communications\"}\n",
        "                    ]\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "    },\n",
        "    config={\n",
        "        \"wait_for_completion\": True,\n",
        "        \"revision_labels\": {\n",
        "            \"data_source\": \"updated_preference\",\n",
        "            \"verified\": \"false\",  # This update was not verified!\n",
        "        },\n",
        "    },\n",
        ")\n",
        "\n",
        "if update_operation.response and update_operation.response.generated_memories:\n",
        "    updated_ref = update_operation.response.generated_memories[0]\n",
        "\n",
        "    # Get the updated memory\n",
        "    updated_memory = client.agent_engines.memories.get(name=updated_ref.memory.name)\n",
        "\n",
        "    print(\"✅ Memory updated (incorrectly)!\")\n",
        "    print(f\"   📝 New fact: {updated_memory.fact}\")\n",
        "    print(f\"   🔄 Action: {updated_ref.action}\")\n",
        "\n",
        "    # Save the previous revision for rollback\n",
        "    if hasattr(updated_ref, \"previous_revision\") and updated_ref.previous_revision:\n",
        "        previous_revision_name = updated_ref.previous_revision\n",
        "        previous_revision_id = previous_revision_name.split(\"/\")[-1]\n",
        "        print(f\"   📜 Previous revision ID: {previous_revision_id}\")\n",
        "else:\n",
        "    print(\"⚠️ Failed to update memory\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4DMqbrexnTPp"
      },
      "source": [
        "#### View the revision history before rollback\n",
        "\n",
        "Let's examine both revisions to see the incorrect update in the history.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99Rid5rbmP3Y"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📜 Viewing revision history before rollback\\n\")\n",
        "\n",
        "# List all revisions to see the history\n",
        "revisions_before = list(\n",
        "    client.agent_engines.memories.revisions.list(name=rollback_memory_name)\n",
        ")\n",
        "\n",
        "print(f\"✅ Found {len(revisions_before)} revisions:\\n\")\n",
        "\n",
        "for i, rev in enumerate(revisions_before, 1):\n",
        "    revision_id = rev.name.split(\"/\")[-1]\n",
        "    labels = rev.labels if hasattr(rev, \"labels\") and rev.labels else {}\n",
        "\n",
        "    print(f\"📌 Revision {i}:\")\n",
        "    print(f\"   🆔 ID: {revision_id}\")\n",
        "    print(f\"   📝 Fact: {rev.fact}\")\n",
        "    print(f\"   ⏱️ Created: {rev.create_time}\")\n",
        "    print(f\"   🏷️ Labels: {labels}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ixzu0s5toGS_"
      },
      "source": [
        "#### Perform the rollback\n",
        "\n",
        "Now we'll rollback to the previous (correct) revision using the `rollback()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6_FU1mG7oG0u"
      },
      "outputs": [],
      "source": [
        "print(\"\\nRollback to the previous (correct) revision\\n\")\n",
        "\n",
        "# Rollback to the previous revision using the revision ID we saved\n",
        "rollback_operation = client.agent_engines.memories.rollback(\n",
        "    name=rollback_memory_name, target_revision_id=previous_revision_id\n",
        ")\n",
        "\n",
        "print(\"✅ Rollback operation completed!\")\n",
        "print(f\"   🎯 Target revision: {previous_revision_id}\")\n",
        "\n",
        "# Verify the rollback by getting the current memory state\n",
        "restored_memory = client.agent_engines.memories.get(name=rollback_memory_name)\n",
        "\n",
        "print(\"\\n📋 Memory after rollback:\")\n",
        "print(f\"   📝 Current fact: {restored_memory.fact}\")\n",
        "print(f\"   ⏱️ Last updated: {restored_memory.update_time}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IOKq7YP8h2tV"
      },
      "source": [
        "#### Compare revision history after rollback\n",
        "\n",
        "Notice how rollback removes the newer (incorrect) revision, maintaining data integrity.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IgZXjmYYhsYt"
      },
      "outputs": [],
      "source": [
        "print(\"\\n📊 Comparing revision history after rollback\\n\")\n",
        "\n",
        "# List revisions after rollback\n",
        "revisions_after = list(\n",
        "    client.agent_engines.memories.revisions.list(name=rollback_memory_name)\n",
        ")\n",
        "\n",
        "print(f\"✅ Revision count after rollback: {len(revisions_after)}\\n\")\n",
        "\n",
        "for i, rev in enumerate(revisions_after, 1):\n",
        "    revision_id = rev.name.split(\"/\")[-1]\n",
        "    labels = rev.labels if hasattr(rev, \"labels\") and rev.labels else {}\n",
        "\n",
        "    print(f\"📌 Revision {i} (after rollback):\")\n",
        "    print(f\"   🆔 ID: {revision_id}\")\n",
        "    print(f\"   📝 Fact: {rev.fact}\")\n",
        "    print(f\"   🏷️ Labels: {labels}\")\n",
        "\n",
        "    # Highlight if this is the verified original\n",
        "    if labels.get(\"verified\") == \"true\":\n",
        "        print(\"   ✅ This is the verified original revision\")\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-hGvdGDFok4h"
      },
      "source": [
        "### Filtering Revisions by Labels\n",
        "\n",
        "Revision labels allow you to track metadata about when and how memories were created or updated. You can filter revisions by these labels for advanced governance workflows."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "McuwANaKo1tw"
      },
      "source": [
        "#### Filter revisions using label queries\n",
        "\n",
        "Let's demonstrate two common filtering scenarios: finding verified revisions and filtering by data source.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x2tVcf1fozjq"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🔍 Filtering revisions by labels\\n\")\n",
        "\n",
        "# Example 1: Get all verified revisions\n",
        "print(\"Example 1: Filter for verified revisions only\\n\")\n",
        "\n",
        "try:\n",
        "    verified_revisions = list(\n",
        "        client.agent_engines.memories.revisions.list(\n",
        "            name=rollback_memory_name, config={\"filter\": 'labels.verified=\"true\"'}\n",
        "        )\n",
        "    )\n",
        "\n",
        "    print(f\"✅ Found {len(verified_revisions)} verified revision(s):\\n\")\n",
        "\n",
        "    for rev in verified_revisions:\n",
        "        print(f\"   📝 {rev.fact}\")\n",
        "        print(f\"   ⏱️ Created: {rev.create_time}\")\n",
        "        print(f\"   🏷️ Labels: {rev.labels}\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"⚠️ Filter operation: {e}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GttiyPalo_Ip"
      },
      "outputs": [],
      "source": [
        "# Example 2: Get revisions from a specific data source\n",
        "print(\"\\nExample 2: Filter by data source\\n\")\n",
        "\n",
        "try:\n",
        "    source_revisions = list(\n",
        "        client.agent_engines.memories.revisions.list(\n",
        "            name=rollback_memory_name,\n",
        "            config={\"filter\": 'labels.data_source=\"initial_preference\"'},\n",
        "        )\n",
        "    )\n",
        "\n",
        "    print(\n",
        "        f\"✅ Found {len(source_revisions)} revision(s) from 'initial_preference' source:\\n\"\n",
        "    )\n",
        "\n",
        "    for rev in source_revisions:\n",
        "        print(f\"   📝 {rev.fact}\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"⚠️ Filter operation: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BjY4p9vCitMN"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring unnecessary costs, it's best practice to delete resources you no longer need."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kdHamkNxh77Q"
      },
      "outputs": [],
      "source": [
        "print(\"\\n🧹 Cleaning up resources...\\n\")\n",
        "\n",
        "delete_agent_engine = True\n",
        "\n",
        "if delete_agent_engine:\n",
        "    print(\"⏳ Deleting Agent Engine (this will remove all memories and sessions)...\\n\")\n",
        "\n",
        "    # Delete the agent engine\n",
        "    # The force=True parameter ensures all contained resources are also deleted\n",
        "    client.agent_engines.delete(name=agent_engine_name, force=True)\n",
        "\n",
        "    print(\"✅ Agent Engine deleted successfully!\")\n",
        "else:\n",
        "    print(\"⚠️ Cleanup skipped - Agent Engine and all data retained\")\n",
        "    print(f\"   📋 Resource Name: {agent_engine_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-47e4jPoi7n-"
      },
      "source": [
        "## Congratulations!\n",
        "\n",
        "You've completed the \"Get started with Memory Bank - Governance\" tutorial!\n",
        "\n",
        "You now have a better understanding of how better manage memories, configure Granular TTL (Time-To-Live), filter Memories by topic and\n",
        "track memory revisions.  \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"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_memory_bank_governance.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
