{
  "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": [
        "# Customizing Memory Topics\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_custom_topics.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_custom_topics.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_custom_topics.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_custom_topics.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/memory_bank/get_started_with_memory_bank_custom_topics.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_custom_topics.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_custom_topics.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_custom_topics.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_custom_topics.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 tutorial demonstrates how to customize **Vertex AI Memory Bank** for specialized domains by defining custom memory topics and providing few-shot examples. You will build an AI-powered financial advisor assistant that remembers client investment goals, risk preferences, and financial situations across multiple conversations.\n",
        "\n",
        "While Memory Bank's default managed topics (like `USER_PERSONAL_INFO` and `USER_PREFERENCES`) work well for general-purpose agents, specialized domains like financial services require more granular, domain-specific memory extraction. This tutorial shows you how to teach Memory Bank to recognize and extract financial-specific information with precision.\n",
        "\n",
        "By the end of this tutorial, you'll have an assistant that extracts and remembers these specialized financial details with high accuracy, enabling truly personalized financial advice across multiple client interactions.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n",
        "\n",
        "First, we'll install the Vertex AI SDK. We need version 1.111.0 or higher to access all 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 and ensures the new packages are properly loaded.\n"
      ]
    },
    {
      "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)."
      ]
    },
    {
      "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)\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 essential Python libraries and 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": "7WdrJibDyrGt"
      },
      "source": [
        "### Define helper function for displaying memories\n",
        "\n",
        "This helper function provides consistent formatting when displaying generated memories. It shows the action type (CREATED or UPDATED) and retrieves the full memory details for each generated memory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EU8ek3jjyuuZ"
      },
      "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",
        "    Returns:\n",
        "        List of memory facts for further processing\n",
        "    \"\"\"\n",
        "    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",
        "                    # Retrieve full memory details including topics\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",
        "                    memories.append(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",
        "    return memories\n",
        "\n",
        "\n",
        "print(\"✅ Helper function defined successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Generate memories with default managed topics\n",
        "\n",
        "Before defining custom topics, let's see what memories are automatically extracted using Memory Bank's default configuration. This baseline will help us understand the limitations and motivate the need for domain-specific customization.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "onQ1_f7-y8ST"
      },
      "source": [
        "### Understanding default managed topics\n",
        "\n",
        "Memory Bank comes with four pre-defined managed topics that work well for general conversational agents:\n",
        "\n",
        "| Managed Topic | Description | Example |\n",
        "|---------------|-------------|---------|\n",
        "| `USER_PERSONAL_INFO` | Basic personal details about the user | \"Client name is Michael Chen, age 42\" |\n",
        "| `USER_PREFERENCES` | General preferences and likes/dislikes | \"Client prefers email communication\" |\n",
        "| `KEY_CONVERSATION_DETAILS` | Important outcomes or milestones | \"Client scheduled follow-up for next month\" |\n",
        "| `EXPLICIT_INSTRUCTIONS` | Direct remember/forget requests | \"Client asked to remember account number\" |\n",
        "\n",
        "These topics work well for general conversational agents, but they lack the granularity needed for specialized domains like financial services.\n",
        "\n",
        "**The Problem:**\n",
        "When a client says \"I want to retire at 65 with 80,000 dollars annual income, and I have a conservative risk tolerance,\" the default topics might capture \"client wants to retire\" as a preference, but they won't extract:\n",
        "- Specific retirement age (65)\n",
        "- Target retirement income (80,000 dollars /year)\n",
        "- Risk tolerance category (conservative)\n",
        "- Investment timeline (23 years if client is 42)\n",
        "\n",
        "This is why we need custom topics."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fD1DOTGHzOBn"
      },
      "source": [
        "### Create Agent Engine with default configuration\n",
        "\n",
        "Let's create a Memory Bank instance using only the default managed topics. This serves as our baseline for comparison.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7uKlj0WNzYzR"
      },
      "outputs": [],
      "source": [
        "print(\"🛠️ Creating Agent Engine with default managed topics...\\n\")\n",
        "print(\"📋 This configuration uses ONLY the default topics:\")\n",
        "print(\"   - USER_PERSONAL_INFO\")\n",
        "print(\"   - USER_PREFERENCES\")\n",
        "print(\"   - KEY_CONVERSATION_DETAILS\")\n",
        "print(\"   - EXPLICIT_INSTRUCTIONS\\n\")\n",
        "\n",
        "# Configure Memory Bank with default settings (no customization)\n",
        "default_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 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",
        "    # Note: No customization_configs provided, so default managed topics are used\n",
        ")\n",
        "\n",
        "# Create the Agent Engine resource\n",
        "# This provisions the backend infrastructure for Memory Bank\n",
        "default_agent_engine = client.agent_engines.create(\n",
        "    config={\"context_spec\": {\"memory_bank_config\": default_memory_config}}\n",
        ")\n",
        "\n",
        "default_engine_name = default_agent_engine.api_resource.name\n",
        "\n",
        "print(\"✅ Agent Engine created with default configuration!\")\n",
        "print(f\"   Resource Name: {default_engine_name}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tfHL3gdxzmjK"
      },
      "source": [
        "### Create a session for a financial advisor client\n",
        "\n",
        "A **Session** represents a single conversation between a client and the financial advisor. Each session is linked to a `user_id`, which allows Memory Bank to associate memories with specific clients across multiple conversations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UoPa4T8ozquk"
      },
      "outputs": [],
      "source": [
        "print(\"💬 Creating session for financial advisor client...\\n\")\n",
        "\n",
        "# Generate unique client ID\n",
        "client_id = \"client_michael_\" + str(uuid.uuid4())[:4]\n",
        "\n",
        "# Create a session for this client\n",
        "session = client.agent_engines.sessions.create(\n",
        "    name=default_engine_name,\n",
        "    user_id=client_id,\n",
        "    config={\"display_name\": f\"Financial consultation for {client_id}\"},\n",
        ")\n",
        "\n",
        "session_name = session.response.name\n",
        "\n",
        "print(\"✅ Session created successfully!\")\n",
        "print(f\"   Session Name: {session_name}\")\n",
        "print(f\"   Client ID: {client_id}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P6G_R8liztI3"
      },
      "source": [
        "### Add initial financial advisor conversation\n",
        "\n",
        "Now we'll add a realistic conversation between a financial advisor and a client. This conversation contains rich financial information that should ideally be extracted into structured memories.\n",
        "\n",
        "As you read this conversation, notice the specific financial details:\n",
        "- Investment goals (retirement at 65, college fund for daughter)\n",
        "- Risk tolerance (conservative, doesn't want to lose savings)\n",
        "- Financial capacity (150k income, 200k in 401k, 50k savings)\n",
        "- Asset details (house worth 600k, 300k mortgage)\n",
        "- Liquidity timeline (100k needed in 8 years for college)\n",
        "\n",
        "Let's see how well the default topics capture these details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2uMDVzin0Okb"
      },
      "outputs": [],
      "source": [
        "print(\"Adding initial financial advisor conversation to session...\\n\")\n",
        "\n",
        "# Realistic conversation between financial advisor and client\n",
        "# This contains the structured financial data we want to extract\n",
        "financial_conversation = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Hi, I'm Michael Chen. I'm 42 years old and want to start planning for retirement.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Hello Michael! I'm glad you're thinking about retirement planning. Let's discuss your financial goals and current situation.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"My main goal is to retire at 65 with a comfortable lifestyle. I also want to save for my daughter's college education - she's 10 now.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Excellent. Two important goals: retirement in 23 years and college savings in about 8 years. What's your current income and savings situation?\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"I earn $150,000 annually. I have about $200,000 in my 401k and $50,000 in savings. My house is worth $600,000 with a $300,000 mortgage.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Thank you for sharing. That gives us a good starting point. What's your comfort level with investment risk?\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"I'm pretty conservative. I don't want to risk losing my savings, but I understand I need some growth for retirement.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"That's a balanced approach. We can structure a portfolio that's conservative but still provides growth potential. Let me note your risk tolerance as moderate-conservative.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"message\": \"Also, I'll need to access about $100,000 in 8 years for my daughter's college tuition. The rest can stay invested for retirement.\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"model\",\n",
        "        \"message\": \"Perfect. We'll structure your investments with different time horizons - more liquid assets for the college fund and longer-term growth investments for retirement.\",\n",
        "    },\n",
        "]\n",
        "\n",
        "# Add each conversation turn to the session\n",
        "# We track invocation_id to maintain conversation order\n",
        "invocation_id = 0\n",
        "\n",
        "for turn in financial_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=session_name,\n",
        "        author=client_id,  # Required: who is speaking\n",
        "        invocation_id=str(invocation_id),  # Required: conversation sequence\n",
        "        timestamp=datetime.datetime.now(tz=datetime.timezone.utc),  # Required: when\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 added to session successfully!\")\n",
        "print(\"💡 Now let's see what memories are extracted with default topics...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qD_hg0Wp0YK_"
      },
      "source": [
        "### Generate memories with default topics\n",
        "\n",
        "Now let's trigger memory generation using the default managed topics. The `generate` method will:\n",
        "- Read the entire conversation from the session\n",
        "- Extract facts that match the default managed topics\n",
        "- Store them as memories in the Memory Bank\n",
        "\n",
        "In particular, behind the scenes:\n",
        "\n",
        "- Gemini 2.5 Flash reads the conversation\n",
        "- It identifies information matching default topics (USER_PERSONAL_INFO, USER_PREFERENCES, etc.)\n",
        "- It consolidates similar facts to avoid duplicates\n",
        "- It stores the extracted memories with the user's scope\n",
        "\n",
        "**Note:** `wait_for_completion=True` makes this a blocking call, which is useful for this tutorial. In production, you might set it to `False` to run asynchronously.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lq0I1cMh0j5J"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"GENERATING MEMORIES WITH DEFAULT MANAGED TOPICS\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"📋 This will extract information matching default topics:\")\n",
        "print(\"   - USER_PERSONAL_INFO (name, age, basic details)\")\n",
        "print(\"   - USER_PREFERENCES (general likes/dislikes)\")\n",
        "print(\"   - KEY_CONVERSATION_DETAILS (important outcomes)\")\n",
        "print(\"   - EXPLICIT_INSTRUCTIONS (direct remember/forget requests)\\n\")\n",
        "\n",
        "print(\"⏳ Processing conversation...\")\n",
        "\n",
        "# Generate memories from the financial conversation\n",
        "# This is a long-running operation that processes the entire session\n",
        "default_operation = client.agent_engines.memories.generate(\n",
        "    name=default_engine_name,\n",
        "    vertex_session_source={\"session\": session_name},\n",
        "    config={\"wait_for_completion\": True},  # Wait for completion (blocking call)\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Memory generation complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18iaptPW2VFq"
      },
      "source": [
        "### Display memories extracted with default topics\n",
        "\n",
        "Let's examine what memories were extracted using the default configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KYMC391R2UeD"
      },
      "outputs": [],
      "source": [
        "# Display the generated memories using our helper function\n",
        "default_memories = display_generated_memories(\n",
        "    default_operation, client, \"Memories Extracted with DEFAULT Managed Topics\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qkV5PJ9N22cg"
      },
      "source": [
        "### Retrieve and analyze baseline memories\n",
        "\n",
        "Let's retrieve all memories and get a concrete count of what was captured.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BfkJrSXl23PV"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"BASELINE RESULTS: Complete Memory List\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "# Retrieve all memories for this client\n",
        "results = client.agent_engines.memories.retrieve(\n",
        "    name=default_engine_name, scope={\"user_id\": client_id}\n",
        ")\n",
        "\n",
        "all_default_memories = list(results)\n",
        "\n",
        "print(f\"✅ Total memories captured with default topics: {len(all_default_memories)}\\n\")\n",
        "\n",
        "if all_default_memories:\n",
        "    print(\"📋 Complete list of extracted memories:\")\n",
        "    for i, memory in enumerate(all_default_memories, 1):\n",
        "        print(f\"   {i}. {memory.memory.fact}\")\n",
        "else:\n",
        "    print(\"⚠️ No memories were extracted (this would indicate a problem)\")\n",
        "\n",
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"KEY OBSERVATION:\")\n",
        "print(\"=\" * 80)\n",
        "print(\"\\nDefault topics provide a basic understanding but:\")\n",
        "print(\"  - Merge distinct financial concepts into generic categories\")\n",
        "print(\"  - Miss specific numerical details crucial for financial planning\")\n",
        "print(\"  - Don't distinguish between different investment timelines\")\n",
        "print(\"  - Lack structured categorization needed for domain expertise\")\n",
        "print(\"\\nThis is exactly why financial services need CUSTOM TOPICS.\")\n",
        "print(\"Let's build them now...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L3rY2yyA3Dfk"
      },
      "source": [
        "## Define custom financial advisor topics\n",
        "\n",
        "Now let's define specialized topics that capture the specific information a financial advisor needs to remember about each client.\n",
        "\n",
        "### Understanding custom topics\n",
        "\n",
        "**Custom Topics** allow you to teach Memory Bank exactly what information matters in your domain. For financial services, we need topics that capture:\n",
        "\n",
        "- **Investment Goals**: What the client wants to achieve (retirement, education, wealth accumulation)\n",
        "- **Risk Tolerance**: How much volatility they can accept (conservative, moderate, aggressive)\n",
        "- **Financial Capacity**: Income, assets, liabilities, net worth\n",
        "- **Account Preferences**: Account types, trading preferences, investment vehicles\n",
        "- **Liquidity Timeline**: When funds will be needed (short-term vs. long-term)\n",
        "\n",
        "Each custom topic has:\n",
        "- **label**: A short identifier (e.g., \"risk_tolerance\")\n",
        "- **description**: Detailed instructions telling the extraction model what to look for. It is like a job description for the extraction model. The more specific and detailed you are, the better the extraction quality.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXs-y5CK3X7L"
      },
      "source": [
        "### Define the five custom financial topics\n",
        "\n",
        "Let's create our five custom topics with detailed, instructive descriptions.\n",
        "\n",
        "Each topic focuses on one dimension of financial planning. This creates clear boundaries and prevents overlapping extraction, which improves both precision and consistency."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QDfcq-HH3WE4"
      },
      "outputs": [],
      "source": [
        "print(\"🎨 Defining custom financial advisor topics...\\n\")\n",
        "\n",
        "custom_financial_topics = [\n",
        "    # Topic 1: Investment Goals\n",
        "    # What the client wants to achieve financially\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"investment_goals\",\n",
        "            description=\"\"\"Extract the client's specific financial objectives and goals. Include:\n",
        "\n",
        "            - Primary goals (retirement planning, education funding, wealth accumulation, major purchases)\n",
        "            - Specific targets with timelines (e.g., \"retire at 65\", \"college fund needed in 8 years\")\n",
        "            - Target amounts or income requirements (e.g., \"$80,000 annual retirement income\")\n",
        "            - Life milestones tied to goals (e.g., \"daughter turns 18\", \"house purchase at 45\")\n",
        "            - Multiple goals should be captured separately (retirement vs. education vs. home purchase)\n",
        "\n",
        "            Format: Extract as specific, actionable goals with quantitative details when available.\n",
        "            Example: \"Client's primary goal is retirement at age 65 with $80,000 annual income\"\n",
        "            Example: \"Client needs $200,000 for daughter's college education in 8 years (daughter currently age 10)\"\n",
        "\n",
        "            Do NOT include risk tolerance or current financial situation here - those belong in other topics.\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    # Topic 2: Risk Tolerance\n",
        "    # How much investment volatility the client can accept\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"risk_tolerance\",\n",
        "            description=\"\"\"Extract the client's comfort level with investment risk and volatility. Include:\n",
        "\n",
        "            - Risk profile category (conservative, moderate-conservative, moderate, moderate-aggressive, aggressive)\n",
        "            - Specific risk preferences and attitudes (e.g., \"can't afford to lose principal\", \"wants growth\", \"balanced approach\")\n",
        "            - Emotional comfort with market fluctuations\n",
        "            - Past experiences that influence risk appetite (e.g., \"lost money in 2008\", \"comfortable with volatility\")\n",
        "            - Willingness to accept short-term losses for long-term gains\n",
        "            - Specific risk constraints (e.g., \"needs guaranteed income\", \"cannot risk college fund\")\n",
        "\n",
        "            Format: Capture both stated risk tolerance and behavioral/emotional indicators.\n",
        "            Example: \"Client has conservative risk tolerance - doesn't want to risk losing principal but understands need for some growth\"\n",
        "            Example: \"Client is aggressive investor comfortable with high volatility given 20-year timeline\"\n",
        "\n",
        "            Do NOT include investment goals or financial capacity here - focus only on risk attitudes.\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    # Topic 3: Financial Capacity\n",
        "    # Current financial situation and resources\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"financial_capacity\",\n",
        "            description=\"\"\"Extract the client's complete current financial situation and resources. Include:\n",
        "\n",
        "            - Annual income (salary, bonuses, other income sources with specific amounts)\n",
        "            - Current savings and investments (401k balance, IRA, brokerage accounts, emergency fund)\n",
        "            - Regular contribution ability (monthly/annual amounts client can invest)\n",
        "            - Real estate holdings (primary residence value, investment properties, mortgage balances)\n",
        "            - Other assets (business ownership, vehicle value, inheritances expected)\n",
        "            - Liabilities (mortgages, loans, credit card debt with amounts)\n",
        "            - Net worth or equity positions when mentioned\n",
        "            - Cash flow and saving capacity\n",
        "\n",
        "            Format: Extract specific dollar amounts and account types. Calculate derived values when helpful.\n",
        "            Example: \"Client earns $150,000 annually\"\n",
        "            Example: \"Client has $200,000 in 401k, $50,000 in savings (total liquid assets: $250,000)\"\n",
        "            Example: \"Client owns home worth $600,000 with $300,000 mortgage (equity: $300,000)\"\n",
        "\n",
        "            Focus on current financial state, not goals or preferences.\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    # Topic 4: Account Preferences\n",
        "    # Preferred account types and investment vehicles\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"account_preferences\",\n",
        "            description=\"\"\"Extract the client's preferences for account types and investment vehicles. Include:\n",
        "\n",
        "            - Preferred account types (401k, traditional IRA, Roth IRA, 529 education plans, taxable brokerage)\n",
        "            - Investment vehicle preferences (mutual funds, ETFs, individual stocks, bonds, alternatives)\n",
        "            - Trading frequency and style (buy-and-hold, active trading, rebalancing preferences)\n",
        "            - Tax considerations (preference for tax-advantaged accounts, tax-loss harvesting interest)\n",
        "            - Employer benefits (401k match utilization, ESPP participation)\n",
        "            - Diversification preferences (sector preferences, international exposure, concentration limits)\n",
        "            - ESG or ethical investing preferences if mentioned\n",
        "\n",
        "            Format: Capture both existing account usage and stated preferences for future investments.\n",
        "            Example: \"Client prefers Roth IRA for tax-free growth in retirement\"\n",
        "            Example: \"Client wants diversified portfolio across asset classes\"\n",
        "            Example: \"Client prefers index funds and ETFs over individual stock picking\"\n",
        "\n",
        "            Do NOT include dollar amounts or balances here - those belong in financial_capacity.\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    # Topic 5: Liquidity Timeline\n",
        "    # When the client needs access to funds\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"liquidity_timeline\",\n",
        "            description=\"\"\"Extract when the client will need access to invested funds and liquidity requirements. Include:\n",
        "\n",
        "            - Short-term needs (within 3 years): emergency fund, planned purchases, near-term expenses\n",
        "            - Medium-term needs (3-10 years): home down payment, education funding, business investment\n",
        "            - Long-term needs (10+ years): retirement, legacy planning, long-term wealth building\n",
        "            - Specific amounts needed at specific times (e.g., \"$100,000 in 8 years for college\")\n",
        "            - Flexibility in timing (hard deadlines vs. flexible goals)\n",
        "            - Partial vs. full liquidity requirements\n",
        "            - Rolling or phased access needs (e.g., \"annual withdrawals starting at 65\")\n",
        "\n",
        "            Format: Be specific about amounts, timelines, and flexibility. Separate different time horizons clearly.\n",
        "            Example: \"Client needs $100,000 in 8 years for daughter's college tuition (hard deadline)\"\n",
        "            Example: \"Client has 23-year timeline until retirement at 65 (flexible, long-term)\"\n",
        "            Example: \"Client wants to maintain $50,000 emergency fund for immediate access\"\n",
        "\n",
        "            This helps structure portfolios with appropriate investment horizons and liquidity.\"\"\",\n",
        "        )\n",
        "    ),\n",
        "]\n",
        "\n",
        "print(\"✅ Custom financial topics defined successfully!\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MK8RyzKi3p5Z"
      },
      "source": [
        "### Create customization configuration\n",
        "\n",
        "We package our custom topics into a `CustomizationConfig` object. This configuration will guide the memory extraction process. It tells Memory Bank to use our 5 custom topics instead of (or in addition to) the default managed topics. During memory generation, Gemini will be instructed to extract information matching our detailed topic descriptions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U1zVH_Xe317y"
      },
      "outputs": [],
      "source": [
        "print(\"⚙️ Creating customization configuration...\\n\")\n",
        "\n",
        "# Package custom topics into a CustomizationConfig\n",
        "# This object will be passed to the MemoryBankConfig\n",
        "financial_customization = CustomizationConfig(memory_topics=custom_financial_topics)\n",
        "\n",
        "print(\"✅ Customization configuration created!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVWkXJXz37bJ"
      },
      "source": [
        "### Create Memory Bank config with custom topics\n",
        "\n",
        "Now we create a new `MemoryBankConfig` that includes our custom topic configuration. This config will use our specialized financial topics instead of the default generic ones."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YpG28XhX37wP"
      },
      "outputs": [],
      "source": [
        "print(\"⚙️ Creating Memory Bank config with custom topics...\\n\")\n",
        "\n",
        "# Create Memory Bank config with our custom topics\n",
        "custom_memory_config = MemoryBankConfig(\n",
        "    # Same embedding model for similarity search (text-embedding-005)\n",
        "    similarity_search_config=SimilaritySearchConfig(\n",
        "        embedding_model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005\"\n",
        "    ),\n",
        "    # Same LLM for memory extraction (gemini-2.5-flash)\n",
        "    generation_config=GenerationConfig(\n",
        "        model=f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "    ),\n",
        "    # NEW: Add our custom financial topics\n",
        "    # This is what makes the difference!\n",
        "    customization_configs=[financial_customization],\n",
        ")\n",
        "\n",
        "print(\"✅ Memory Bank configuration created with custom topics!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LA9gZmb_4idB"
      },
      "source": [
        "## Create new Agent Engine with custom topics\n",
        "\n",
        "To see the impact of custom topics clearly, we'll create a new Agent Engine. This ensures a clean comparison between default and custom topic extraction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n56SmN3j4t3W"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"CREATING AGENT ENGINE WITH CUSTOM TOPICS\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"🛠️ Creating new Agent Engine with custom financial topics...\")\n",
        "\n",
        "# Create new Agent Engine with custom financial topics\n",
        "custom_agent_engine = client.agent_engines.create(\n",
        "    config={\"context_spec\": {\"memory_bank_config\": custom_memory_config}}\n",
        ")\n",
        "\n",
        "custom_engine_name = custom_agent_engine.api_resource.name\n",
        "\n",
        "print(\"\\n✅ Agent Engine created with custom topics!\")\n",
        "print(f\"   Resource Name: {custom_engine_name}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xPSfCiXu4yNQ"
      },
      "source": [
        "### Create session for custom engine\n",
        "\n",
        "We'll create a new session associated with our custom-configured Agent Engine. This allows us to run the same conversation through a different memory extraction pipeline.\n",
        "\n",
        "**Important:**\n",
        "We're using the same `client_id` as before. This means memories from both engines are associated with the same user, which would allow consolidation if we were using a single engine in production.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hrvJRm7F4lea"
      },
      "outputs": [],
      "source": [
        "print(\"💬 Creating session for custom topics engine...\\n\")\n",
        "\n",
        "# Create new session for the custom engine\n",
        "# Using the same client_id to maintain user identity\n",
        "custom_session = client.agent_engines.sessions.create(\n",
        "    name=custom_engine_name,\n",
        "    user_id=client_id,  # Same client ID as baseline\n",
        "    config={\"display_name\": f\"Custom topics session for {client_id}\"},\n",
        ")\n",
        "\n",
        "custom_session_name = custom_session.response.name\n",
        "\n",
        "print(\"✅ Session created successfully!\")\n",
        "print(f\"   Session Name: {custom_session_name}\")\n",
        "print(f\"   Client ID: {client_id} (same as baseline)\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpDH0Iie49N9"
      },
      "source": [
        "### Add the same conversation to custom engine\n",
        "\n",
        "We'll add the exact same financial advisor conversation to this new session. This enables an \"apples-to-apples\" comparison of default vs. custom topic extraction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zs21aC7f4fLO"
      },
      "outputs": [],
      "source": [
        "print(\"⬆️ Adding the same financial conversation to custom session...\\n\")\n",
        "\n",
        "# Reset invocation counter for new session\n",
        "invocation_id = 0\n",
        "\n",
        "# Add the same conversation to the custom engine's session\n",
        "for turn in financial_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=custom_session_name,\n",
        "        author=client_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",
        "\n",
        "print(\"✅ Conversation added to custom session successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "imSd5YsA5JiX"
      },
      "source": [
        "### Generate memories with custom topics\n",
        "\n",
        "Now we trigger memory generation on the custom-configured engine. The extraction model will now be guided by our detailed financial topic descriptions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9BQA1gPH5Ef8"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"GENERATING MEMORIES WITH CUSTOM FINANCIAL TOPICS\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"⏳ Processing conversation...\")\n",
        "\n",
        "# Generate memories using custom topics\n",
        "# This is where the magic happens!\n",
        "custom_operation = client.agent_engines.memories.generate(\n",
        "    name=custom_engine_name,\n",
        "    vertex_session_source={\"session\": custom_session_name},\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Memory generation complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0-ZO8ZSY5RC7"
      },
      "source": [
        "### Display memories extracted with custom topics\n",
        "\n",
        "Let's examine the memories extracted using our custom financial topics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5SMsVEw5UCw"
      },
      "outputs": [],
      "source": [
        "# Display memories extracted with custom topics\n",
        "custom_memories = display_generated_memories(\n",
        "    custom_operation, client, \"Memories Extracted with CUSTOM Financial Topics\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_UeV9KQl5jNy"
      },
      "source": [
        "## Compare default vs. custom topic extraction\n",
        "\n",
        "Now let's do a side-by-side comparison to see the impact of custom topics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7PV7VcQ35sqL"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n",
        "print(\"  SIDE-BY-SIDE COMPARISON: Default vs. Custom Topics\")\n",
        "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n",
        "\n",
        "# Retrieve all memories from default engine\n",
        "print(\"📊 BASELINE - Memories with DEFAULT Managed Topics:\")\n",
        "default_results = client.agent_engines.memories.retrieve(\n",
        "    name=default_engine_name, scope={\"user_id\": client_id}\n",
        ")\n",
        "default_memories_list = list(default_results)\n",
        "\n",
        "if default_memories_list:\n",
        "    for i, memory in enumerate(default_memories_list, 1):\n",
        "        print(f\"   {i}. {memory.memory.fact}\")\n",
        "else:\n",
        "    print(\"   (No memories extracted)\")\n",
        "\n",
        "print(f\"\\n✅ Total: {len(default_memories_list)} memories\")\n",
        "\n",
        "print(\"\\n\" + \"-\" * 80 + \"\\n\")\n",
        "\n",
        "# Retrieve all memories from custom engine\n",
        "print(\"📊 IMPROVED - Memories with CUSTOM Financial Topics:\")\n",
        "custom_results = client.agent_engines.memories.retrieve(\n",
        "    name=custom_engine_name, scope={\"user_id\": client_id}\n",
        ")\n",
        "custom_memories_list = list(custom_results)\n",
        "\n",
        "if custom_memories_list:\n",
        "    for i, memory in enumerate(custom_memories_list, 1):\n",
        "        # Get full memory to access topics\n",
        "        full_memory = client.agent_engines.memories.get(name=memory.memory.name)\n",
        "        print(f\"   {i}. {memory.memory.fact}\")\n",
        "else:\n",
        "    print(\"   (No memories extracted)\")\n",
        "\n",
        "print(f\"\\n✅ Total: {len(custom_memories_list)} memories\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e-Hn6REq6Nwq"
      },
      "source": [
        "## Add few-shot examples for even better extraction\n",
        "\n",
        "Custom topics tell the model *what* to extract. Few-shot examples show the model *how* to extract it. By providing example conversations and the exact memory facts you expect, you teach the model your domain's nuances and preferred extraction style.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6srFKW-B6Q_i"
      },
      "source": [
        "### Understanding few-shot examples\n",
        "\n",
        "**Few-shot examples** are training samples that demonstrate the desired extraction behavior. Each example contains:\n",
        "1. **conversation_source**: A sample conversation snippet\n",
        "2. **generated_memories**: The exact memory facts you want to be extracted from that conversation\n",
        "\n",
        "**Best Practices:**\n",
        "- Provide 2-5 examples per domain (more is not always better)\n",
        "- Use realistic conversations from your domain\n",
        "- Show different scenarios (conservative client, aggressive client, complex situation)\n",
        "- Demonstrate the desired level of granularity\n",
        "- Include calculated or derived values if you want the model to infer them"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sOdj3b5p6YSm"
      },
      "source": [
        "### Create few-shot examples for financial topics\n",
        "\n",
        "We'll create 3 few-shot examples covering different aspects of financial advising:\n",
        "1. **Conservative retiree** - near-term needs, risk-averse\n",
        "2. **Aggressive young investor** - long-term horizon, high risk tolerance\n",
        "3. **Balanced investor** - specific account preferences, structured portfolio\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tcJUbTaL6hQy"
      },
      "outputs": [],
      "source": [
        "print(\"🎨 Creating few-shot examples for financial memory extraction...\\n\")\n",
        "\n",
        "financial_few_shot_examples = [\n",
        "    # Example 1: Conservative Retiree - Near-Term Needs\n",
        "    # Demonstrates: near retirement, conservative risk, specific timeline\n",
        "    GenerateMemoriesExample(\n",
        "        conversation_source=ConversationSource(\n",
        "            events=[\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"user\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"I'm retiring in 3 years and need to move my 401k into something safer. I have about $400,000 saved. I can't take any major risks at this point.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"model\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"I understand. With 3 years until retirement, capital preservation is key. Let's discuss safe investment options.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "        generated_memories=[\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client's primary investment goal is retirement in 3 years; needs to transition 401k to safer investments\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has conservative risk tolerance - cannot accept major risks due to near retirement timeline\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(fact=\"Client has $400,000 in current 401k savings\"),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has 3-year timeline until retirement; needs fund access at that point (short-term liquidity requirement)\"\n",
        "            ),\n",
        "        ],\n",
        "    ),\n",
        "    # Example 2: Aggressive Young Investor - Long-Term Growth\n",
        "    # Demonstrates: long horizon, high risk tolerance, specific investment preferences\n",
        "    GenerateMemoriesExample(\n",
        "        conversation_source=ConversationSource(\n",
        "            events=[\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"user\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"I just got a $10,000 bonus and want to invest it aggressively. I'm 28 and won't need this money for at least 20 years. I'm thinking tech stocks or crypto.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"model\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"Great! With a 20-year horizon at your age, you can handle volatility. Let's discuss growth-oriented options.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "        generated_memories=[\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client's investment goal is long-term wealth building with 20+ year investment horizon\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has aggressive risk tolerance - comfortable with high volatility given long timeline and young age (28)\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has $10,000 from bonus available for investment; currently age 28\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client prefers growth-oriented investments: tech stocks and cryptocurrency\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has 20+ year timeline with no liquidity needs until then (long-term investment horizon)\"\n",
        "            ),\n",
        "        ],\n",
        "    ),\n",
        "    # Example 3: Balanced Investor - Structured Portfolio\n",
        "    # Demonstrates: specific account type, asset allocation, regular contributions\n",
        "    GenerateMemoriesExample(\n",
        "        conversation_source=ConversationSource(\n",
        "            events=[\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"user\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"I want to open a Roth IRA and contribute the maximum each year. I'm comfortable with 70% stocks and 30% bonds. I can invest $6,500 annually.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "                ConversationSourceEvent(\n",
        "                    content=Content(\n",
        "                        role=\"model\",\n",
        "                        parts=[\n",
        "                            Part(\n",
        "                                text=\"Excellent choice! A Roth IRA with that 70/30 allocation is a solid retirement strategy.\"\n",
        "                            )\n",
        "                        ],\n",
        "                    )\n",
        "                ),\n",
        "            ]\n",
        "        ),\n",
        "        generated_memories=[\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client's investment goal is to maximize Roth IRA contributions for tax-advantaged retirement savings\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client has moderate to moderate-aggressive risk tolerance with 70% stocks / 30% bonds asset allocation preference\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client can contribute $6,500 annually (maximum Roth IRA contribution limit)\"\n",
        "            ),\n",
        "            ExampleGeneratedMemory(\n",
        "                fact=\"Client prefers Roth IRA account type with 70% stock / 30% bond allocation\"\n",
        "            ),\n",
        "        ],\n",
        "    ),\n",
        "]\n",
        "\n",
        "print(\"✅ Few-shot examples created successfully!\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cd60cIHc6ofZ"
      },
      "source": [
        "### Update customization config with few-shot examples\n",
        "\n",
        "Now we add our few-shot examples to the customization configuration, creating our most advanced extraction setup combining:\n",
        "\n",
        "1. **Custom Topics** (what to extract) - 5 financial topics\n",
        "2. **Few-Shot Examples** (how to extract it) - 3 demonstration examples\n",
        "\n",
        "This is the gold standard for domain-specific memory extraction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jp3msgjE6vwG"
      },
      "outputs": [],
      "source": [
        "print(\"\\n⚙️ Updating customization configuration with few-shot examples...\\n\")\n",
        "\n",
        "# Create advanced customization with both custom topics AND few-shot examples\n",
        "advanced_financial_customization = CustomizationConfig(\n",
        "    memory_topics=custom_financial_topics,  # Our 5 custom topics (WHAT to extract)\n",
        "    generate_memories_examples=financial_few_shot_examples,  # Our 3 few-shot examples (HOW to extract)\n",
        ")\n",
        "\n",
        "print(\"✅ Advanced customization created!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Mvwf7nS60VU"
      },
      "source": [
        "### Update Agent Engine with few-shot configuration\n",
        "\n",
        "Instead of creating a new engine, we'll update our existing custom engine with the enhanced configuration. This demonstrates how to evolve an agent's capabilities in production without losing existing memories."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "btYyb-x_66-I"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"UPDATING AGENT ENGINE WITH FEW-SHOT EXAMPLES\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "# Create Memory Bank config with custom topics AND few-shot examples\n",
        "advanced_memory_config = MemoryBankConfig(\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",
        "    customization_configs=[\n",
        "        advanced_financial_customization\n",
        "    ],  # Updated with few-shot examples\n",
        ")\n",
        "\n",
        "# Update the existing custom engine\n",
        "print(\"⚙️ Updating Agent Engine configuration...\")\n",
        "updated_engine = client.agent_engines.update(\n",
        "    name=custom_engine_name,\n",
        "    config={\"context_spec\": {\"memory_bank_config\": advanced_memory_config}},\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Agent Engine updated with few-shot examples!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NlnlCp0b7C5X"
      },
      "source": [
        "### Create new session to test few-shot impact\n",
        "\n",
        "We'll create another session to test the impact of few-shot examples on memory extraction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lg42vZCq7DnO"
      },
      "outputs": [],
      "source": [
        "print(\"\\n💬 Creating new session to test few-shot examples...\\n\")\n",
        "\n",
        "# Create new session for few-shot testing\n",
        "# Still using the same client_id for continuity\n",
        "fewshot_session = client.agent_engines.sessions.create(\n",
        "    name=custom_engine_name,  # Using the updated engine with few-shot\n",
        "    user_id=client_id,\n",
        "    config={\"display_name\": f\"Few-shot test session for {client_id}\"},\n",
        ")\n",
        "\n",
        "fewshot_session_name = fewshot_session.response.name\n",
        "\n",
        "print(\"✅ Session created successfully!\")\n",
        "print(f\"   Session Name: {fewshot_session_name}\")\n",
        "print(f\"   Engine: {custom_engine_name} (now with few-shot examples)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gw_h8LcD7KYY"
      },
      "source": [
        "## Add same conversation to few-shot session\n",
        "\n",
        "Once again, we add the same conversation to ensure a fair comparison.\n",
        "\n",
        "In fact, we now have three memory extractions from the same conversation:\n",
        "1. Default managed topics (baseline)\n",
        "2. Custom topics only\n",
        "3. Custom topics + few-shot examples (best)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T70W4aku7IN_"
      },
      "outputs": [],
      "source": [
        "print(\"\\n⬆️ Adding same financial conversation to few-shot session...\\n\")\n",
        "\n",
        "# Reset invocation counter\n",
        "invocation_id = 0\n",
        "\n",
        "# Add conversation to few-shot session\n",
        "for turn in financial_conversation:\n",
        "    client.agent_engines.sessions.events.append(\n",
        "        name=fewshot_session_name,\n",
        "        author=client_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",
        "\n",
        "print(\"✅ Conversation added to few-shot session!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RoJRPtXt7VKT"
      },
      "source": [
        "### Generate memories with few-shot examples\n",
        "\n",
        "Now we generate memories with the fully configured system: custom topics + few-shot examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aheEiAQ-7Y0V"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"GENERATING MEMORIES WITH CUSTOM TOPICS + FEW-SHOT EXAMPLES\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"⏳ Processing conversation...\")\n",
        "\n",
        "# Generate memories with few-shot configuration\n",
        "fewshot_operation = client.agent_engines.memories.generate(\n",
        "    name=custom_engine_name,\n",
        "    vertex_session_source={\"session\": fewshot_session_name},\n",
        "    config={\"wait_for_completion\": True},\n",
        ")\n",
        "\n",
        "print(\"\\n✅ Memory generation complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g6SxaFCa7c0R"
      },
      "source": [
        "### Display memories with few-shot examples\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S7u2Ijrw7a86"
      },
      "outputs": [],
      "source": [
        "# Display memories generated with few-shot examples\n",
        "fewshot_memories = display_generated_memories(\n",
        "    fewshot_operation,\n",
        "    client,\n",
        "    \"Memories Extracted with CUSTOM TOPICS + FEW-SHOT EXAMPLES\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TzvefJ0b7xvJ"
      },
      "source": [
        "## Three-way comparison: Default vs. Custom vs. Custom + Few-Shot\n",
        "\n",
        "Now let's see all three approaches side-by-side to demonstrate the progressive improvement."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vbdmbm1_74p8"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n",
        "print(\"  COMPREHENSIVE COMPARISON: Three Approaches to Memory Extraction\")\n",
        "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"📊 APPROACH 1: Default Managed Topics (Baseline)\")\n",
        "print(\"\\n\")\n",
        "print(\"⚙️ Configuration:\")\n",
        "print(\"   - 4 generic managed topics (USER_PERSONAL_INFO, USER_PREFERENCES, etc.)\")\n",
        "print(\"   - No domain customization\")\n",
        "print(\"   - No few-shot examples\\n\")\n",
        "\n",
        "print(\"📋 Extracted Memories:\")\n",
        "default_results = client.agent_engines.memories.retrieve(\n",
        "    name=default_engine_name, scope={\"user_id\": client_id}\n",
        ")\n",
        "default_list = list(default_results)\n",
        "\n",
        "if default_list:\n",
        "    for i, memory in enumerate(default_list, 1):\n",
        "        print(f\"   {i}. {memory.memory.fact}\")\n",
        "else:\n",
        "    print(\"   (No memories extracted)\")\n",
        "\n",
        "print(\"\\n\" + \"-\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"📊 APPROACH 2: Custom Financial Topics\")\n",
        "print(\"\\n\")\n",
        "print(\"⚙️ Configuration:\")\n",
        "print(\"   - 5 domain-specific topics (investment_goals, risk_tolerance, etc.)\")\n",
        "print(\"   - Detailed topic descriptions\")\n",
        "print(\"   - No few-shot examples\\n\")\n",
        "\n",
        "print(\"📋 Extracted Memories:\")\n",
        "custom_results = client.agent_engines.memories.retrieve(\n",
        "    name=custom_engine_name, scope={\"user_id\": client_id}\n",
        ")\n",
        "custom_list = list(custom_results)\n",
        "\n",
        "if custom_list:\n",
        "    for i, memory in enumerate(custom_list, 1):\n",
        "        full_memory = client.agent_engines.memories.get(name=memory.memory.name)\n",
        "        print(f\"   {i}. {memory.memory.fact}\")\n",
        "else:\n",
        "    print(\"   (No memories extracted)\")\n",
        "\n",
        "print(\"\\n\" + \"-\" * 80 + \"\\n\")\n",
        "\n",
        "print(\"📊 APPROACH 3: Custom Topics + Few-Shot Examples (Optimized)\")\n",
        "print(\"\\n\")\n",
        "print(\"⚙️ Configuration:\")\n",
        "print(\"   - 5 domain-specific topics (WHAT to extract)\")\n",
        "print(\"   - 3 few-shot examples (HOW to extract it)\")\n",
        "print(\"   - Detailed topic descriptions + extraction patterns\\n\")\n",
        "\n",
        "print(\"📋 Extracted Memories:\")\n",
        "# The fewshot_memories were already extracted, display them\n",
        "for i, fact in enumerate(fewshot_memories, 1):\n",
        "    print(f\"   {i}. {fact}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8V0Sn6gC-_17"
      },
      "source": [
        "## Putting everything together\n",
        "\n",
        "Now let's demonstrate how these improved memories enable truly personalized financial advice."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0qqlmwqO_O7U"
      },
      "source": [
        "### Create function to generate personalized advice\n",
        "\n",
        "This function fetches relevant memories using similarity search and uses Gemini to create personalized financial advice.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aQCK1sBY_JLT"
      },
      "outputs": [],
      "source": [
        "def generate_personalized_financial_advice(client_id, query, engine_name):\n",
        "    \"\"\"Generate personalized financial advice using memories and Gemini.\"\"\"\n",
        "    from google import genai\n",
        "\n",
        "    genai_client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "    print(\"\\n🔍 Generating personalized advice for query:\")\n",
        "    print(f'   \"{query}\"\\n')\n",
        "\n",
        "    # Retrieve relevant memories using similarity search\n",
        "    results = client.agent_engines.memories.retrieve(\n",
        "        name=engine_name,\n",
        "        scope={\"user_id\": client_id},\n",
        "        similarity_search_params={\n",
        "            \"search_query\": query,\n",
        "            \"top_k\": 5,  # Get top 5 most relevant memories\n",
        "        },\n",
        "    )\n",
        "\n",
        "    memories = list(results)\n",
        "\n",
        "    print(\"📋 Client Context (from Memory Bank):\\n\")\n",
        "    memory_context = []\n",
        "\n",
        "    if memories:\n",
        "        for i, mem in enumerate(memories, 1):\n",
        "            print(f\"   {i}. {mem.memory.fact}\")\n",
        "            memory_context.append(mem.memory.fact)\n",
        "    else:\n",
        "        print(\"⚠️ No relevant memories found.\")\n",
        "        print(\"\\n💡 Without memory context, Gemini can only provide generic advice.\")\n",
        "        return None\n",
        "\n",
        "    # Create prompt augmented with memory context\n",
        "    prompt = f\"\"\"You are an expert financial advisor. Based on the following client information\n",
        "from their previous conversations, provide personalized financial advice.\n",
        "\n",
        "Client Information from Memory Bank:\n",
        "{chr(10).join(f\"- {fact}\" for fact in memory_context)}\n",
        "\n",
        "Client Question: {query}\n",
        "\n",
        "Provide specific, actionable financial advice that:\n",
        "1. Addresses their question directly\n",
        "2. Considers their stated risk tolerance\n",
        "3. Aligns with their investment goals and timeline\n",
        "4. Takes into account their financial capacity\n",
        "5. Respects their liquidity needs\n",
        "\n",
        "Format your response with:\n",
        "1. Direct Answer (personalized to their situation)\n",
        "2. Rationale (based on their specific profile)\n",
        "3. Specific Recommendations (with concrete numbers when appropriate)\n",
        "4. Important Considerations\n",
        "5. Next Steps\"\"\"\n",
        "\n",
        "    # Generate personalized response with Gemini\n",
        "    print(\"\\n\")\n",
        "    print(\"💬 Personalized Financial Advice from Gemini:\\n\")\n",
        "\n",
        "    response = genai_client.models.generate_content(\n",
        "        model=\"gemini-2.5-flash\",\n",
        "        contents=prompt,\n",
        "    )\n",
        "\n",
        "    print(response.text)\n",
        "\n",
        "    return response.text\n",
        "\n",
        "\n",
        "print(\"✅ Function defined successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JS1JN3fc_pZM"
      },
      "source": [
        "### Test with default topics engine\n",
        "\n",
        "Let's first see what advice is generated using memories from the default topics engine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4JMAfdCY_kKh"
      },
      "outputs": [],
      "source": [
        "print(\"📊 TEST 1: Advice using DEFAULT MANAGED TOPICS\")\n",
        "print(\"-\" * 80)\n",
        "print(\n",
        "    \"❓ Query: 'Should I invest more aggressively to maximize returns for retirement?'\\n\"\n",
        ")\n",
        "print(\"📋 Memories available: From default managed topics (generic)\\n\")\n",
        "\n",
        "advice_default = generate_personalized_financial_advice(\n",
        "    client_id=client_id,\n",
        "    query=\"Should I invest more aggressively to maximize returns for retirement?\",\n",
        "    engine_name=default_engine_name,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PAin11HXAxqj"
      },
      "source": [
        "### Test with custom topics + few-shot engine\n",
        "\n",
        "Now let's see the improvement when using memories from our optimized engine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R2XUjuVCAvq0"
      },
      "outputs": [],
      "source": [
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"📊 TEST 2: Advice using CUSTOM TOPICS + FEW-SHOT EXAMPLES\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "print(\n",
        "    \"❓ Query: 'Should I invest more aggressively to maximize returns for retirement?'\\n\"\n",
        ")\n",
        "print(\n",
        "    \"📋 Memories available: From custom topics + few-shot (structured, comprehensive)\\n\"\n",
        ")\n",
        "\n",
        "advice_custom = generate_personalized_financial_advice(\n",
        "    client_id=client_id,\n",
        "    query=\"Should I invest more aggressively to maximize returns for retirement?\",\n",
        "    engine_name=custom_engine_name,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nFjTnsVgA8ez"
      },
      "source": [
        "## Cleanup\n",
        "\n",
        "To avoid incurring unexpected costs, let's clean up the resources we created.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cvkDIf0eBBac"
      },
      "outputs": [],
      "source": [
        "print(\"🧹 Cleaning up resources...\\n\")\n",
        "\n",
        "delete_agent_engines = True  # Set to False if you want to keep the resources\n",
        "\n",
        "if delete_agent_engines:\n",
        "    try:\n",
        "        # Delete default topics engine\n",
        "        client.agent_engines.delete(name=default_engine_name, force=True)\n",
        "\n",
        "        # Delete custom topics engine (includes few-shot configuration)\n",
        "        client.agent_engines.delete(name=custom_engine_name, force=True)\n",
        "\n",
        "        print(\"✅ All resources cleaned up successfully!\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\n⚠️ Error during cleanup: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i0zuDHCCBikA"
      },
      "source": [
        "**Congratulations!**\n",
        "\n",
        "Congratulations! You've completed the tutorial on customizing Memory Bank for specialized domains.\n",
        "\n",
        "**What's Next?**\n",
        "\n",
        "Now that you understand custom topics and few-shot learning, you can:\n",
        "\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_custom_topics.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
