{
  "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": [
        "# Building a Multimodal Trip Planner with ADK on Vertex AI Agent Engine 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/tutorial_get_started_with_multimodal_agents_with_memory_bank.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fagents%2Fagent_engine%2Fmemory_bank%2Ftutorial_get_started_with_multimodal_agents_with_memory_bank.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/agents/agent_engine/memory_bank/tutorial_get_started_with_multimodal_agents_with_memory_bank.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/memory_bank/tutorial_get_started_with_multimodal_agents_with_memory_bank.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\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/tutorial_get_started_with_multimodal_agents_with_memory_bank.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/tutorial_get_started_with_multimodal_agents_with_memory_bank.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/tutorial_get_started_with_multimodal_agents_with_memory_bank.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/tutorial_get_started_with_multimodal_agents_with_memory_bank.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/tutorial_get_started_with_multimodal_agents_with_memory_bank.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",
        "| [Yee Sian](https://github.com/yeesian) |\n",
        "| [Shawn Yang](https://github.com/shawn-yang-google) |\n",
        "| [Kimberly Milam](https://github.com/klmilam) |\n",
        "| [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIbgu3eqffcp"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "    <b>Warning:</b> Vertex AI Memory Bank is in Preview. Multimodal memories in this tutorial will depend on the dimensions of your multimodal inputs and the customization features you leverage.\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Imagine a trip planner that remembers your travel photos from Paris, understands your video tour of Tokyo's streets, and recalls your voice notes about that perfect gelato shop in Rome. This tutorial teaches you to build exactly that—a multimodal trip planner powered by Memory Bank, the managed memory service of Vertex AI Agent Engine.\n",
        "\n",
        "### What You'll Learn\n",
        "\n",
        "- **Multimodal Memory Generation**: Extract memories from images, videos, and audio\n",
        "- **Smart Retrieval**: Fetch relevant memories based on context and similarity\n",
        "- **Agent Development**: Build an ADK agent that thinks like a travel companion using your multimodal inputs\n",
        "- **Agent Deployment**: Deploy your multimodal agent to Vertex AI Agent Engine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet \"google-cloud-aiplatform[agent_engines, adk]==1.115.0\" \"google-adk==1.14.1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "from vertexai import Client\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
        "# fmt: on\n",
        "if not BUCKET_NAME or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "# Set variables for ADK\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
        "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"1\"\n",
        "APP_NAME = \"trip_planner\"\n",
        "USER_ID = \"traveler_123\"\n",
        "\n",
        "# Create the staging bucket for Agent Engine\n",
        "!gsutil mb -p $PROJECT_ID -l $LOCATION $BUCKET_URI\n",
        "\n",
        "client = Client(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "We import all the Python libraries we'll need."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "import asyncio\n",
        "import warnings\n",
        "from typing import Any\n",
        "\n",
        "from IPython.display import Audio, Image, Video, display\n",
        "\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "from google.adk import tools\n",
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.memory import VertexAiMemoryBankService\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService, VertexAiSessionService\n",
        "from google.adk.tools import FunctionTool\n",
        "from google.genai import types\n",
        "from vertexai.agent_engines import AgentEngine\n",
        "from vertexai.types import ManagedTopicEnum\n",
        "from vertexai.types import MemoryBankCustomizationConfigMemoryTopic as MemoryTopic\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic as CustomMemoryTopic,\n",
        ")\n",
        "from vertexai.types import (\n",
        "    MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic as ManagedMemoryTopic,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AT7kpogMYzlP"
      },
      "source": [
        "### Helpers\n",
        "\n",
        "We define some helper functions you will use along the tutorial.\n",
        "\n",
        "In particular, the `MultimodalMemoryGenerator` takes different media types (image, video, audio) along with some text context, format them correctly for the Memory Bank API, and trigger the memory generation process.\n",
        "\n",
        "> Note: Memory Bank doesn't just store the raw files. It uses Gemini and its multimodal capabilities to process the file and the context, extracting key facts and concepts. These extracted facts are what become the \"memories.\"\n",
        "\n",
        "The other helpers, `_parse_memories`, `call_agent`, and `call_remote_agent`, are utilities for extracting API responses and running our agent, both locally and when it's deployed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MMlmbBpEY1uS"
      },
      "outputs": [],
      "source": [
        "def display_image(image_uri: str, width: int = None, height: int = None):\n",
        "    \"\"\"Displays an image from a given URI with optional width and height.\"\"\"\n",
        "    display(Image(image_uri, width=width, height=height))\n",
        "\n",
        "\n",
        "def display_video(video_uri: str, width: int = None, height: int = None):\n",
        "    \"\"\"Displays a video from a given URI with optional width and height.\"\"\"\n",
        "    display(Video(video_uri, width=width, height=height))\n",
        "\n",
        "\n",
        "def display_audio(audio_uri: str):\n",
        "    \"\"\"Displays an audio file from a given URI.\"\"\"\n",
        "    display(Audio(audio_uri))\n",
        "\n",
        "\n",
        "class MultimodalMemoryGenerator:\n",
        "    \"\"\"Handles memory generation from various media types.\"\"\"\n",
        "\n",
        "    def __init__(self, client, agent_engine_name):\n",
        "        self.client = client\n",
        "        self.agent_engine_name = agent_engine_name\n",
        "\n",
        "    def generate_from_image(self, image_uri: str, context: str, user_id: str):\n",
        "        \"\"\"Generate memories from travel photos.\"\"\"\n",
        "        mime_type = image_uri.split(\".\")[-1]\n",
        "        if mime_type not in [\"webp\", \"jpeg\", \"png\", \"jpg\"]:\n",
        "            raise ValueError(f\"Unsupported image file type: {mime_type}\")\n",
        "        if mime_type == \"jpg\":\n",
        "            mime_type = \"jpeg\"\n",
        "        mime_type = f\"image/{mime_type}\"\n",
        "\n",
        "        events = [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\"text\": context},\n",
        "                        {\"file_data\": {\"file_uri\": image_uri, \"mime_type\": mime_type}},\n",
        "                    ],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "\n",
        "        response = self.client.agent_engines.memories.generate(\n",
        "            name=self.agent_engine_name,\n",
        "            direct_contents_source={\"events\": events},\n",
        "            scope={\"user_id\": user_id},\n",
        "            config={\"wait_for_completion\": True},\n",
        "        )\n",
        "\n",
        "        return self._parse_memories(response)\n",
        "\n",
        "    def generate_from_video(self, video_uri: str, context: str, user_id: str):\n",
        "        \"\"\"Generate memories from travel videos.\"\"\"\n",
        "        mime_type = video_uri.split(\".\")[-1]\n",
        "        if mime_type not in [\n",
        "            \"mp4\",\n",
        "            \"mpeg\",\n",
        "            \"x-flv\",\n",
        "            \"quicktime\",\n",
        "            \"mpegps\",\n",
        "            \"mpg\",\n",
        "            \"webm\",\n",
        "            \"wmv\",\n",
        "            \"3gpp\",\n",
        "        ]:\n",
        "            raise ValueError(f\"Unsupported video file type: {mime_type}\")\n",
        "        mime_type = f\"video/{mime_type}\"\n",
        "\n",
        "        events = [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\"text\": context},\n",
        "                        {\"file_data\": {\"file_uri\": video_uri, \"mime_type\": mime_type}},\n",
        "                    ],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "\n",
        "        response = self.client.agent_engines.memories.generate(\n",
        "            name=self.agent_engine_name,\n",
        "            direct_contents_source={\"events\": events},\n",
        "            scope={\"user_id\": user_id},\n",
        "            config={\"wait_for_completion\": True},\n",
        "        )\n",
        "\n",
        "        return self._parse_memories(response)\n",
        "\n",
        "    def generate_from_audio(self, audio_uri: str, context: str, user_id: str):\n",
        "        \"\"\"Generate memories from voice notes about travel.\"\"\"\n",
        "        mime_type = audio_uri.split(\".\")[-1]\n",
        "        if mime_type not in [\n",
        "            \"aac\",\n",
        "            \"flac\",\n",
        "            \"mp3\",\n",
        "            \"m4a\",\n",
        "            \"mpeg\",\n",
        "            \"mpga\",\n",
        "            \"mp4\",\n",
        "            \"opus\",\n",
        "            \"pcm\",\n",
        "            \"wav\",\n",
        "            \"webm\",\n",
        "        ]:\n",
        "            raise ValueError(f\"Unsupported audio file type: {mime_type}\")\n",
        "        mime_type = f\"audio/{mime_type}\"\n",
        "\n",
        "        events = [\n",
        "            {\n",
        "                \"content\": {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\"text\": f\"Voice note context: {context}\"},\n",
        "                        {\"file_data\": {\"file_uri\": audio_uri, \"mime_type\": mime_type}},\n",
        "                    ],\n",
        "                }\n",
        "            }\n",
        "        ]\n",
        "\n",
        "        operation = self.client.agent_engines.memories.generate(\n",
        "            name=self.agent_engine_name,\n",
        "            direct_contents_source={\"events\": events},\n",
        "            scope={\"user_id\": user_id},\n",
        "            config={\"wait_for_completion\": True},\n",
        "        )\n",
        "\n",
        "        return self._parse_memories(operation)\n",
        "\n",
        "    def _parse_memories(self, operation):\n",
        "        \"\"\"Parse and format generated memories.\"\"\"\n",
        "        memories = []\n",
        "        if operation.response and operation.response.generated_memories:\n",
        "            for gen_memory in operation.response.generated_memories:\n",
        "                if gen_memory.action != \"DELETED\" and gen_memory.memory:\n",
        "                    try:\n",
        "                        full_memory = client.agent_engines.memories.get(\n",
        "                            name=gen_memory.memory.name\n",
        "                        )\n",
        "                        memories.append({\"fact\": full_memory.fact})\n",
        "                    except Exception as e:\n",
        "                        print(f\"Could not get any memory: {e}\")\n",
        "        else:\n",
        "            print(\"No memories generated\")\n",
        "        return memories\n",
        "\n",
        "\n",
        "def call_agent(runner: Runner, content: str, session_id: str, user_id: str):\n",
        "    \"\"\"Calls the agent and prints the final response.\"\"\"\n",
        "    events = runner.run(user_id=user_id, session_id=session_id, new_message=content)\n",
        "\n",
        "    for event in events:\n",
        "        if event.is_final_response():\n",
        "            final_response = event.content.parts[0].text\n",
        "            print(\"=\" * 50)\n",
        "            print(\"Agent Response:\", final_response)\n",
        "            print(\"=\" * 50)\n",
        "            return final_response\n",
        "\n",
        "\n",
        "async def call_remote_agent(\n",
        "    agent: AgentEngine, session_id: str, message: dict, user_id: str\n",
        "):\n",
        "    \"\"\"Sends a message to the deployed agent, streams the response,\n",
        "    and prints only the final text output.\n",
        "\n",
        "    Args:\n",
        "        session_id: The ID of the current conversation session.\n",
        "        message: The message payload, including text and/or file data.\n",
        "    \"\"\"\n",
        "    final_response = \"\"\n",
        "    async for event in agent.async_stream_query(\n",
        "        user_id=user_id,\n",
        "        session_id=session_id,\n",
        "        message=message,\n",
        "    ):\n",
        "        if event.get(\"content\"):\n",
        "            for part in event[\"content\"].get(\"parts\", []):\n",
        "                if \"text\" in part:\n",
        "                    final_response = part[\"text\"]\n",
        "                    print(\"=\" * 50)\n",
        "                    print(\"Agent Response:\", final_response)\n",
        "                    print(\"=\" * 50)\n",
        "                    return final_response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2MR_3-HKQqr"
      },
      "source": [
        "## Creating multimodal Memory Bank with Travel-Specific Topics\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LvWYLBksYroW"
      },
      "source": [
        "### Create Topics\n",
        "\n",
        "Before we can store memories, we can define their structure. Memory Bank uses \"topics\" to organize information and allows you to customize them in a way to extract memories that are relevant for your use case.\n",
        "\n",
        "We will use a mix of:\n",
        "- Managed Topics: Pre-defined concepts for common concepts like USER_PREFERENCES.\n",
        "- Custom Topics: Concepts we define ourselves, tailored to our specific domain—in this case, travel. This is where we define the schema for our agent's memory. Memory Bank uses 'topics' to categorize information, which helps the underlying model know what's important to extract and save.  \n",
        "\n",
        "In our case, we're defining travel_experiences, travel_preferences, and travel_logistics. The description for each custom topic acts as a prompt, guiding the model on what kind of facts to extract for that category.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kgdhLGfjKTd9"
      },
      "outputs": [],
      "source": [
        "travel_topics = [\n",
        "    MemoryTopic(\n",
        "        managed_memory_topic=ManagedMemoryTopic(\n",
        "            managed_topic_enum=ManagedTopicEnum.USER_PREFERENCES\n",
        "        )\n",
        "    ),\n",
        "    MemoryTopic(\n",
        "        managed_memory_topic=ManagedMemoryTopic(\n",
        "            managed_topic_enum=ManagedTopicEnum.USER_PERSONAL_INFO\n",
        "        )\n",
        "    ),\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"travel_experiences\",\n",
        "            description=\"\"\"Memorable travel experiences including:\n",
        "                - Places visited and impressions\n",
        "                - Favorite restaurants, cafes, and food experiences\n",
        "                - Preferred accommodation types and locations\n",
        "                - Activities enjoyed (museums, hiking, beaches, etc.)\n",
        "                - Travel companions and social preferences\n",
        "                - Photos and videos from trips with location context\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"travel_preferences\",\n",
        "            description=\"\"\"Travel style and preferences:\n",
        "                - Budget preferences (luxury, mid-range, budget)\n",
        "                - Transportation preferences (flying, trains, driving)\n",
        "                - Trip duration preferences\n",
        "                - Season and weather preferences\n",
        "                - Cultural interests and language abilities\n",
        "                - Dietary restrictions and food preferences\"\"\",\n",
        "        )\n",
        "    ),\n",
        "    MemoryTopic(\n",
        "        custom_memory_topic=CustomMemoryTopic(\n",
        "            label=\"travel_logistics\",\n",
        "            description=\"\"\"Practical travel information:\n",
        "                - Passport and visa information\n",
        "                - Frequent flyer numbers and hotel loyalty programs\n",
        "                - Emergency contacts\n",
        "                - Medical considerations and insurance\n",
        "                - Packing preferences and essentials\n",
        "                - Time zone preferences and jet lag strategies\"\"\",\n",
        "        )\n",
        "    ),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9woCwCFsKaNr"
      },
      "source": [
        "### Configure Memory Bank with our travel topics\n",
        "\n",
        "Now we assemble the full memory_bank_config, which will be passed to our Agent Engine when we create it.\n",
        "\n",
        "This config is composed of a few key parts.\n",
        "\n",
        "The customization_configs is where we pass in our travel_topics list from the previous cell, telling Memory Bank to use our custom schema. For similarity_search_config, we're specifying gemini-embedding-001. This model is used for similarity searches when retrieving memories. Finally, the generation_config specifies gemini-2.5-flash as the model. This model will be used for the generation of memories (i.e., processing the inputs and extracting the facts)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "42imKmi7Kaz2"
      },
      "outputs": [],
      "source": [
        "memory_bank_config = {\n",
        "    \"customization_configs\": [\n",
        "        {\n",
        "            \"memory_topics\": travel_topics,\n",
        "        }\n",
        "    ],\n",
        "    \"similarity_search_config\": {\n",
        "        \"embedding_model\": f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-embedding-001\"\n",
        "    },\n",
        "    \"generation_config\": {\n",
        "        \"model\": f\"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash\"\n",
        "    },\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kywl8opiCidn"
      },
      "source": [
        "### Create Agent Engine with Memory Bank\n",
        "\n",
        "This is a key step. We're calling client.agent_engines.create() to create a new Agent Engine instance.\n",
        "\n",
        "An Agent Engine is the managed Vertex AI resource that hosts our agent and its associated services. Crucially, we are not deploying our agent code yet. We are just provisioning the infrastructure that includes the Memory Bank.\n",
        "\n",
        "We pass our memory_bank_config inside the context_spec dictionary. This associates our custom memory schema (topics, models) with this specific Agent Engine. This call will create a persistent resource in your project that you can see in the Google Cloud console."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqfVoZuzCidp"
      },
      "outputs": [],
      "source": [
        "agent_engine = client.agent_engines.create(\n",
        "    config={\n",
        "        \"display_name\": \"Multimodal Memory Bank\",\n",
        "        \"description\": \"Multimodal Memory Bank with Travel-Specific Topics\",\n",
        "        \"context_spec\": {\"memory_bank_config\": memory_bank_config},\n",
        "    }\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J2e0x1vGCidp"
      },
      "source": [
        "### Test the memory\n",
        "\n",
        "Let's instantiate our MultimodalMemoryGenerator helper class. We pass it the client (for making API calls) and the agent_engine.api_resource.name. This resource name is the unique identifier for the Agent Engine we just created, telling the SDK which Memory Bank to send data to."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EQ8OxVQmCidp"
      },
      "outputs": [],
      "source": [
        "memory_gen = MultimodalMemoryGenerator(client, agent_engine.api_resource.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XYfV72K2rXQc"
      },
      "source": [
        "#### Generate memories for audio input\n",
        "\n",
        "Finally, let's ingest an audio file. This .wav file is a voice note about a trip to Gaeta, Italy. We provide the file and the context \"I am travelling in Gaeta.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8F9Z96hK7zed"
      },
      "outputs": [],
      "source": [
        "audio_uri = \"https://storage.googleapis.com/github-repo/audio_ai/gaeta.wav\"\n",
        "display_audio(audio_uri=audio_uri)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZlrXIW5qTsj0"
      },
      "outputs": [],
      "source": [
        "audio_memories = memory_gen.generate_from_audio(\n",
        "    audio_uri=\"gs://github-repo/audio_ai/gaeta.wav\",\n",
        "    context=\"I am travelling in Gaeta\",\n",
        "    user_id=\"traveler_123\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kDhb6VOurceZ"
      },
      "source": [
        "This output is more detailed! Because the audio file itself contained a rich description (castles, beaches, local food), Memory Bank was able to extract multiple facts from that single file, all of which fall under our custom travel_experiences topic.\n",
        "\n",
        "This demonstrates the importance of using multimodal inputs for memory generation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vnBiKEEUWRDL"
      },
      "outputs": [],
      "source": [
        "print(\"Generated memories from audio:\")\n",
        "for memory in audio_memories:\n",
        "    print(f\"  - {memory['fact']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gHZvpsQ-DfEF"
      },
      "source": [
        "#### Generate memories for image\n",
        "\n",
        "Time for our first multimodal memory! We're calling `generate_from_image` passing a photo of Marienplatz, a text context (\"My visit to Germany last summer...\"), and the user_id. This user_id is important as it scopes the memory to this specific user.\n",
        "\n",
        "\n",
        "Behind the scenes, our helper class calls the `agent_engines.memories.generate()` API. Memory Bank receives both the image and the text, uses a multimodal model to understand them together, and then uses our generation_config model (Gemini 2.5 Flash) to extract facts that match our travel_topics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5T-A8tlLDfEH"
      },
      "outputs": [],
      "source": [
        "image_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/landmark1.jpg\"\n",
        "display_image(image_uri, width=500, height=500)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xMAab_EhDfEH"
      },
      "outputs": [],
      "source": [
        "image_memories = memory_gen.generate_from_image(\n",
        "    image_uri=image_uri,\n",
        "    context=\"My visit to Germany last summer, walking around Marienplatz in Munich\",\n",
        "    user_id=\"traveler_123\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B8Rrt0yeDfEH"
      },
      "source": [
        "Let's check the output. The _parse_memories helper in our class retrieves the full memory fact. As you can see, the model correctly extracted the key information and saved it as a simple, factual statement.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M1k4iPkDDfEH"
      },
      "outputs": [],
      "source": [
        "print(\"Generated memories from image:\")\n",
        "for memory in image_memories:\n",
        "    print(f\"  - {memory['fact']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DV3zYy_sDfEI"
      },
      "source": [
        "#### Generate memories for video input\n",
        "\n",
        "Now let's try with a video. The process is identical: we provide an .mp4 file and a simple text context. The MultimodalMemoryGenerator class handles setting the correct video/mp4 MIME type.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x_PLnjsKDfEI"
      },
      "outputs": [],
      "source": [
        "video_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/mediterraneansea.mp4\"\n",
        "display_video(video_uri=video_uri)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bg-t89wUDfEI"
      },
      "outputs": [],
      "source": [
        "video_memories = memory_gen.generate_from_video(\n",
        "    video_uri=\"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/mediterraneansea.mp4\",\n",
        "    context=\"I am travelling in Europe\",\n",
        "    user_id=\"traveler_123\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uzDf-BkjDfEI"
      },
      "source": [
        "Again, we print the result. The model distills the input into a core fact.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ASEPxXYeDfEI"
      },
      "outputs": [],
      "source": [
        "print(\"Generated memories from video:\")\n",
        "for memory in video_memories:\n",
        "    print(f\"  - {memory['fact']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_jt7oGpWY97C"
      },
      "source": [
        "## Building the Trip Planning Agent with multimodal memory\n",
        "\n",
        "Time to build our multimodal agent using ADK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z9NloH-gitk5"
      },
      "source": [
        "### Build custom tools\n",
        "\n",
        "We're defining a simple Python function `calculate_trip_budget`. The key parts are the type hints (destination: str, days: int) and the docstring. The ADK framework and the underlying LLM will inspect this docstring (especially the description \"Calculates estimated budget for a trip.\") to understand when and how to use this tool.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3ULIYqZLjAZv"
      },
      "outputs": [],
      "source": [
        "def calculate_trip_budget(destination: str, days: int, style: str) -> dict[str, Any]:\n",
        "    \"\"\"Calculates estimated budget for a trip.\n",
        "\n",
        "    Use this tool when the user asks about trip costs or budget estimates.\n",
        "\n",
        "    Args:\n",
        "        destination: The destination city or country\n",
        "        days: Number of days for the trip\n",
        "        style: Travel style (budget, mid-range, luxury)\n",
        "\n",
        "    Returns:\n",
        "        A dictionary with budget breakdown.\n",
        "        Example: {'status': 'success', 'total': 3000, 'daily': 500, 'currency': 'USD'}\n",
        "    \"\"\"\n",
        "    daily_rates = {\"budget\": 100, \"mid-range\": 250, \"luxury\": 500}\n",
        "\n",
        "    daily = daily_rates.get(style.lower(), 250)\n",
        "    total = daily * days\n",
        "\n",
        "    return {\n",
        "        \"status\": \"success\",\n",
        "        \"total\": total,\n",
        "        \"daily\": daily,\n",
        "        \"currency\": \"USD\",\n",
        "        \"breakdown\": {\n",
        "            \"accommodation\": daily * 0.4,\n",
        "            \"food\": daily * 0.3,\n",
        "            \"activities\": daily * 0.2,\n",
        "            \"transport\": daily * 0.1,\n",
        "        },\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a0CB2BzFrxyz"
      },
      "source": [
        "Here, we wrap our Python function in the ADK's FunctionTool class. This converts the function into an object that the ADK LlmAgent can understand and execute.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2GAzLNUYjEw3"
      },
      "outputs": [],
      "source": [
        "budget_tool = FunctionTool(func=calculate_trip_budget)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xXF-3noRjJTs"
      },
      "source": [
        "### Create Enhanced Agent with Tools\n",
        "\n",
        "This is where we define the 'brain' of our agent using the ADK's LlmAgent class.\n",
        "\n",
        "We specify gemini-2.5-flash for fast, intelligent responses.\n",
        "\n",
        "The instruction parameter is the system prompt. We're telling it who it is (\"an advanced trip planning assistant\"), what it can do (\"access to user's complete travel history\"), and how to behave (\"Always mention what memories influenced your suggestions\").\n",
        "\n",
        "Finally, the tools list defines what the agent can do. We're giving it our custom budget_tool and, just as importantly, the built-in tools.preload_memory_tool.PreloadMemoryTool(). This built-in tool is what's responsible for automatically fetching relevant memories from Memory Bank at the start of a new conversation turn.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ctp3_ay4jKTv"
      },
      "outputs": [],
      "source": [
        "trip_planner_agent = LlmAgent(\n",
        "    model=\"gemini-2.5-flash\",\n",
        "    name=\"TripPlanner\",\n",
        "    instruction=\"\"\"You are an advanced trip planning assistant with memory and tools.\n",
        "\n",
        "    Your capabilities:\n",
        "    - Access to user's complete travel history through memories\n",
        "    - Use 'calculate_trip_budget' to provide cost estimates\n",
        "    - Remember all conversations and preferences\n",
        "\n",
        "    Tool usage guidelines:\n",
        "    - Always mention what memories influenced your suggestions\n",
        "    - After providing budget, ask if they want adjustments\n",
        "\n",
        "    Be personal and reference specific past experiences when available.\n",
        "    If not available, just keep talking with the user. Don't make up facts.\n",
        "    \"\"\",\n",
        "    tools=[tools.preload_memory_tool.PreloadMemoryTool(), budget_tool],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yULL3hXAkT3b"
      },
      "source": [
        "### Test the agent\n",
        "\n",
        "Let's test our agent. For the sake of a self-contained test, this code creates a new, empty Agent Engine. This means the test that follows will first populate this new memory, and then test retrieval."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sdE5mT7hZPQv"
      },
      "outputs": [],
      "source": [
        "agent_engine = client.agent_engines.create(\n",
        "    config={\n",
        "        \"display_name\": \"Multimodal Memory Bank\",\n",
        "        \"description\": \"Multimodal Memory Bank with Travel-Specific Topics\",\n",
        "        \"context_spec\": {\"memory_bank_config\": memory_bank_config},\n",
        "    }\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gc6LOZt5siw3"
      },
      "source": [
        "Before we can run our agent locally, we need to configure its services. The ADK Runner ties everything together.\n",
        "\n",
        "It needs a session_service to store short-term conversation history; InMemorySessionService is perfect for local testing as it just keeps the session in your notebook's RAM.\n",
        "\n",
        "It also needs a memory_service for long-term memory. We instantiate VertexAiMemoryBankService and point it to our deployed Agent Engine's ID. This tells the ADK LlmAgent (and its PreloadMemoryTool) which cloud-based Memory Bank to connect to.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ctVSLa3fkPMj"
      },
      "outputs": [],
      "source": [
        "session_service = InMemorySessionService()\n",
        "\n",
        "memory_bank_service = VertexAiMemoryBankService(\n",
        "    agent_engine_id=agent_engine.api_resource.name.split(\"/\")[-1],\n",
        "    project=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        ")\n",
        "\n",
        "runner = Runner(\n",
        "    agent=trip_planner_agent,\n",
        "    app_name=APP_NAME,\n",
        "    session_service=session_service,\n",
        "    memory_service=memory_bank_service,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VtuMLOGWtBZ8"
      },
      "source": [
        "This async function defines a complete, multi-turn test conversation with our local agent. It's designed to simulate a user providing multiple pieces of multimodal information across several turns.\n",
        "\n",
        "It starts by creating a session and sends a series of messages (text, image, video, audio) using our call_agent helper.\n",
        "\n",
        "The most critical step is at the end: memory_bank_service.add_session_to_memory(final_session_state). This is the\n",
        "consolidation step. It takes the entire conversation from the\n",
        "InMemorySessionService and sends it to the VertexAiMemoryBankService to be processed and saved permanently.\n",
        "\n",
        "To prove it worked, the function then starts a new session and asks a question (\"suggest a cultural destination...\") to prove that the agent can recall and use the memories from the previous session.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "754qCiWjeH9k"
      },
      "outputs": [],
      "source": [
        "async def test_trip_planner():\n",
        "    session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "    print(f\"🌍 Starting a new trip planning session: {session.id}\")\n",
        "\n",
        "    image_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/landmark1.jpg\"\n",
        "    video_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/mediterraneansea.mp4\"\n",
        "    audio_uri = \"gs://github-repo/audio_ai/gaeta.wav\"\n",
        "\n",
        "    text_message = types.Content(role=\"user\", parts=[{\"text\": \"Hello!\"}])\n",
        "\n",
        "    call_agent(content=text_message, session_id=session.id, user_id=USER_ID)\n",
        "    asyncio.sleep(5)\n",
        "\n",
        "    mime_type = image_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\"webp\", \"jpeg\", \"png\", \"jpg\"]:\n",
        "        raise ValueError(f\"Unsupported image file type: {mime_type}\")\n",
        "    if mime_type == \"jpg\":\n",
        "        mime_type = \"jpeg\"\n",
        "    mime_type = f\"image/{mime_type}\"\n",
        "\n",
        "    image_message = types.Content(\n",
        "        role=\"user\",\n",
        "        parts=[\n",
        "            {\n",
        "                \"text\": \"I'm planning a trip. First, here is a picture that shows you the kind of place I like.\"\n",
        "            },\n",
        "            {\"file_data\": {\"file_uri\": image_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    )\n",
        "    call_agent(content=image_message, session_id=session.id, user_id=USER_ID)\n",
        "    asyncio.sleep(5)\n",
        "\n",
        "    mime_type = video_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\n",
        "        \"mp4\",\n",
        "        \"mpeg\",\n",
        "        \"x-flv\",\n",
        "        \"quicktime\",\n",
        "        \"mpegps\",\n",
        "        \"mpg\",\n",
        "        \"webm\",\n",
        "        \"wmv\",\n",
        "        \"3gpp\",\n",
        "    ]:\n",
        "        raise ValueError(f\"Unsupported video file type: {mime_type}\")\n",
        "    mime_type = f\"video/{mime_type}\"\n",
        "    video_message = types.Content(\n",
        "        role=\"user\",\n",
        "        parts=[\n",
        "            {\n",
        "                \"text\": \"Next, here's a video. I also enjoy cities close to Mediterranean sea.\"\n",
        "            },\n",
        "            {\"file_data\": {\"file_uri\": video_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    )\n",
        "    call_agent(content=video_message, session_id=session.id, user_id=USER_ID)\n",
        "    asyncio.sleep(5)\n",
        "\n",
        "    mime_type = audio_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\n",
        "        \"aac\",\n",
        "        \"flac\",\n",
        "        \"mp3\",\n",
        "        \"m4a\",\n",
        "        \"mpeg\",\n",
        "        \"mpga\",\n",
        "        \"mp4\",\n",
        "        \"opus\",\n",
        "        \"pcm\",\n",
        "        \"wav\",\n",
        "        \"webm\",\n",
        "    ]:\n",
        "        raise ValueError(f\"Unsupported audio file type: {mime_type}\")\n",
        "    mime_type = f\"audio/{mime_type}\"\n",
        "    audio_message = types.Content(\n",
        "        role=\"user\",\n",
        "        parts=[\n",
        "            {\"text\": \"Finally, I loved Gaeta. To give you an idea, Here is an audio\"},\n",
        "            {\"file_data\": {\"file_uri\": audio_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    )\n",
        "    call_agent(content=audio_message, session_id=session.id, user_id=USER_ID)\n",
        "    asyncio.sleep(5)\n",
        "\n",
        "    print(\"\\n---------------------------------------------------\")\n",
        "    print(\"Conversation finished. Consolidating all memories at once...\")\n",
        "    final_session_state = await session_service.get_session(\n",
        "        app_name=APP_NAME, user_id=USER_ID, session_id=session.id\n",
        "    )\n",
        "    await memory_bank_service.add_session_to_memory(final_session_state)\n",
        "    print(\"✅ Full conversation context (Image, Video, Audio) saved to Memory Bank.\")\n",
        "    print(\"---------------------------------------------------\")\n",
        "\n",
        "    new_session = await session_service.create_session(\n",
        "        app_name=APP_NAME, user_id=USER_ID\n",
        "    )\n",
        "    print(\n",
        "        f\"\\n🌅 Starting a NEW session ({new_session.id}) to test cumulative memory...\"\n",
        "    )\n",
        "\n",
        "    text_message = types.Content(role=\"user\", parts=[{\"text\": \"Hello!\"}])\n",
        "\n",
        "    call_agent(text_message, new_session.id, USER_ID)\n",
        "\n",
        "    verification_message = types.Content(\n",
        "        role=\"user\",\n",
        "        parts=[\n",
        "            {\n",
        "                \"text\": \"Based on the picture, video, AND audio I shared with you before, suggest a cultural destination for me.\"\n",
        "            }\n",
        "        ],\n",
        "    )\n",
        "\n",
        "    call_agent(verification_message, new_session.id, USER_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEXgxWVntVjm"
      },
      "source": [
        "Let's run the test! Watch the output. The first few turns are the agent acknowledging the new information (image, video, audio). After the 'Consolidating all memories' message, a new session starts.\n",
        "\n",
        "Notice how the agent's response to \"Hello!\" is now different. It says \"Hello again!\" and immediately references past memories. This is the PreloadMemoryTool in action!\n",
        "\n",
        "The final response is a detailed recommendation for Naples, which explicitly references the image, the video, and the audio note about Gaeta, demonstrating a sophisticated, context-aware agent using its multimodal memory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WZGHMP1aNVvw"
      },
      "outputs": [],
      "source": [
        "await test_trip_planner()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5XEzUS61EMl7"
      },
      "source": [
        "## Deploy the agent on Vertex AI Agent Engine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUlyRQJNqiXc"
      },
      "source": [
        "### Deploy the agent\n",
        "\n",
        "It's time to deploy your agent! So far, our LlmAgent has only been running locally in this notebook. We will now deploy our agent to the cloud as a scalable, serverless endpoint.\n",
        "\n",
        "We call `client.agent_engines.create` again, but this time we pass our local agent=trip_planner_agent object. The `config` dictionary now defines the deployment itself. We provide a `display_name` and, critically, pass our `memory_bank_config` again in the context_spec. This links our deployed agent code to the same Memory Bank configuration. We also provide\n",
        "requirements and the `staging_bucket` URI we defined at the start.\n",
        "\n",
        "The create call will automatically \"pickle\" our trip_planner_agent object, bundle it with our requirements, and upload it to this bucket. This process takes a few minutes as it builds a container image and provisions the serving infrastructure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jeG75V8RZcAE"
      },
      "outputs": [],
      "source": [
        "remote_mm_adk_agent = client.agent_engines.create(\n",
        "    agent=trip_planner_agent,\n",
        "    config={\n",
        "        \"display_name\": \"Multimodal Memory Bank\",\n",
        "        \"description\": \"Multimodal Memory Bank with Travel-Specific Topics\",\n",
        "        \"context_spec\": {\"memory_bank_config\": memory_bank_config},\n",
        "        \"requirements\": [\n",
        "            \"google-cloud-aiplatform[agent_engines,adk]==1.115.0\",\n",
        "            \"google-adk==1.14.1\",\n",
        "        ],\n",
        "        \"staging_bucket\": BUCKET_URI,\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blb6BDRFRwcV"
      },
      "source": [
        "### Submit a multimodal query\n",
        "\n",
        "The `create` command returns a client object representing the deployed agent.  We'll grab its api_resource.name, which is the full, unique identifier for our deployed agent. We need this name to send it queries. And we initialize the new agent engine endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jlbMi7u5GkMC"
      },
      "outputs": [],
      "source": [
        "remote_mm_adk_agent_name = remote_mm_adk_agent.api_resource.name\n",
        "remote_mm_adk_agent = agent_engines.get(remote_mm_adk_agent_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xx0ZGJGXCJAh"
      },
      "source": [
        "Now you are ready to hit the endpoint. Here you have a test function which is similar to our local test, but it's designed to talk to the deployed agent.  The function uses `remote_mm_adk_agent.async_create_session()` to create a new session ID and formats the messages (e.g., image_message) as simple dictionaries for the async_stream_query API. It also explicitly performing the session retrieval and memory consolidation from the client side using `VertexAiSessionService` and `VertexAiMemoryBankService`. This demonstrates how you can have fine-grained control over memory persistence, even with a deployed agent, just as we did in the local test."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HwwWjC2bqsJq"
      },
      "outputs": [],
      "source": [
        "async def test_remote_trip_planner():\n",
        "    \"\"\"Simulates a multi-turn, multimodal conversation with the deployed agent,\n",
        "    mirroring the exact structure of the original notebook's test function.\n",
        "    \"\"\"\n",
        "    conversation_session = await remote_mm_adk_agent.async_create_session(\n",
        "        user_id=USER_ID\n",
        "    )\n",
        "    session_id = conversation_session[\"id\"]\n",
        "    print(f\"🌍 Starting a new trip planning session: {session_id}\\n\")\n",
        "\n",
        "    image_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/landmark1.jpg\"\n",
        "    video_uri = \"https://storage.googleapis.com/github-repo/img/gemini/multimodality_usecases_overview/mediterraneansea.mp4\"\n",
        "    audio_uri = \"gs://github-repo/audio_ai/gaeta.wav\"\n",
        "\n",
        "    text_message = {\"role\": \"user\", \"parts\": [{\"text\": \"Hello!\"}]}\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=session_id,\n",
        "        message=text_message,\n",
        "    )\n",
        "    await asyncio.sleep(5)\n",
        "\n",
        "    mime_type = image_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\"webp\", \"jpeg\", \"png\", \"jpg\"]:\n",
        "        raise ValueError(f\"Unsupported image file type: {mime_type}\")\n",
        "    if mime_type == \"jpg\":\n",
        "        mime_type = \"jpeg\"\n",
        "    mime_type = f\"image/{mime_type}\"\n",
        "    image_message = {\n",
        "        \"role\": \"user\",\n",
        "        \"parts\": [\n",
        "            {\n",
        "                \"text\": \"I'm planning a trip. First, here is a picture that shows you the kind of place I like.\"\n",
        "            },\n",
        "            {\"file_data\": {\"file_uri\": image_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    }\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=session_id,\n",
        "        message=image_message,\n",
        "    )\n",
        "    await asyncio.sleep(5)\n",
        "\n",
        "    mime_type = video_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\n",
        "        \"mp4\",\n",
        "        \"mpeg\",\n",
        "        \"x-flv\",\n",
        "        \"quicktime\",\n",
        "        \"mpegps\",\n",
        "        \"mpg\",\n",
        "        \"webm\",\n",
        "        \"wmv\",\n",
        "        \"3gpp\",\n",
        "    ]:\n",
        "        raise ValueError(f\"Unsupported video file type: {mime_type}\")\n",
        "    mime_type = f\"video/{mime_type}\"\n",
        "    video_message = {\n",
        "        \"role\": \"user\",\n",
        "        \"parts\": [\n",
        "            {\n",
        "                \"text\": \"Next, here's a video. I also enjoy cities close to the Mediterranean sea.\"\n",
        "            },\n",
        "            {\"file_data\": {\"file_uri\": video_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    }\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=session_id,\n",
        "        message=video_message,\n",
        "    )\n",
        "    await asyncio.sleep(5)\n",
        "\n",
        "    mime_type = audio_uri.split(\".\")[-1]\n",
        "    if mime_type not in [\n",
        "        \"aac\",\n",
        "        \"flac\",\n",
        "        \"mp3\",\n",
        "        \"m4a\",\n",
        "        \"mpeg\",\n",
        "        \"mpga\",\n",
        "        \"mp4\",\n",
        "        \"opus\",\n",
        "        \"pcm\",\n",
        "        \"wav\",\n",
        "        \"webm\",\n",
        "    ]:\n",
        "        raise ValueError(f\"Unsupported audio file type: {mime_type}\")\n",
        "    mime_type = f\"audio/{mime_type}\"\n",
        "    audio_message = {\n",
        "        \"role\": \"user\",\n",
        "        \"parts\": [\n",
        "            {\n",
        "                \"text\": \"Finally, I loved Gaeta. To give you an idea, here is an audio note about it.\"\n",
        "            },\n",
        "            {\"file_data\": {\"file_uri\": audio_uri, \"mime_type\": mime_type}},\n",
        "        ],\n",
        "    }\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=session_id,\n",
        "        message=audio_message,\n",
        "    )\n",
        "\n",
        "    print(\"\\n---------------------------------------------------\")\n",
        "    print(\n",
        "        \"Conversation finished. For a deployed agent, context is automatically saved to Memory Bank.\"\n",
        "    )\n",
        "\n",
        "    vertex_session_service = VertexAiSessionService(PROJECT_ID, LOCATION)\n",
        "    final_session_state = await vertex_session_service.get_session(\n",
        "        app_name=remote_mm_adk_agent_name.split(\"/\")[-1],\n",
        "        user_id=USER_ID,\n",
        "        session_id=session_id,\n",
        "    )\n",
        "\n",
        "    vertex_memory_bank_service = VertexAiMemoryBankService(\n",
        "        agent_engine_id=remote_mm_adk_agent_name.split(\"/\")[-1],\n",
        "        project=PROJECT_ID,\n",
        "        location=LOCATION,\n",
        "    )\n",
        "    await vertex_memory_bank_service.add_session_to_memory(final_session_state)\n",
        "    print(\"✅ Full conversation context is now persistent.\")\n",
        "    print(\"---------------------------------------------------\\n\")\n",
        "    await asyncio.sleep(10)\n",
        "\n",
        "    verification_session = await remote_mm_adk_agent.async_create_session(\n",
        "        user_id=USER_ID\n",
        "    )\n",
        "    new_session_id = verification_session[\"id\"]\n",
        "    print(\n",
        "        f\"🌅 Starting a NEW session ({new_session_id}) to test cumulative memory...\\n\"\n",
        "    )\n",
        "\n",
        "    hello_again_message = {\"role\": \"user\", \"parts\": [{\"text\": \"Hello again!\"}]}\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=new_session_id,\n",
        "        message=hello_again_message,\n",
        "    )\n",
        "    await asyncio.sleep(5)\n",
        "\n",
        "    verification_message = {\n",
        "        \"role\": \"user\",\n",
        "        \"parts\": [\n",
        "            {\n",
        "                \"text\": \"Based on the picture, video, AND audio I shared with you before, suggest a cultural destination for me.\"\n",
        "            }\n",
        "        ],\n",
        "    }\n",
        "    await call_remote_agent(\n",
        "        agent=remote_mm_adk_agent,\n",
        "        user_id=USER_ID,\n",
        "        session_id=new_session_id,\n",
        "        message=verification_message,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pZDL7Ohacex2"
      },
      "outputs": [],
      "source": [
        "# Run the test\n",
        "await test_remote_trip_planner()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Finally, let's clean up our resources to avoid incurring costs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IMKGJ56Im9JF"
      },
      "outputs": [],
      "source": [
        "delete_agent_engine = True\n",
        "\n",
        "if delete_agent_engine:\n",
        "    agent_engines = client.agent_engines.list()\n",
        "    for agent_engine in agent_engines:\n",
        "        agent_engine_name = agent_engine.api_resource.name\n",
        "        if agent_engine.api_resource.display_name == \"Multimodal Memory Bank\":\n",
        "            agent_engine.delete(force=True)\n",
        "            print(f\"Deleted agent engine: {agent_engine_name}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_get_started_with_multimodal_agents_with_memory_bank.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
