chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,603 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ur8xi4C7S06n"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright 2025 Google LLC\n",
|
||||
"#\n",
|
||||
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||
"# you may not use this file except in compliance with the License.\n",
|
||||
"# You may obtain a copy of the License at\n",
|
||||
"#\n",
|
||||
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||
"#\n",
|
||||
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||
"# See the License for the specific language governing permissions and\n",
|
||||
"# limitations under the License."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JAPoU8Sm5E6e"
|
||||
},
|
||||
"source": [
|
||||
"# Get started with Vertex AI Memory Bank - ADK\n",
|
||||
"\n",
|
||||
"<table align=\"left\">\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
|
||||
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagent-engine%2Fmemory%2Fget_started_with_memory_bank_adk.ipynb\">\n",
|
||||
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
|
||||
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\">\n",
|
||||
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
"</table>\n",
|
||||
"\n",
|
||||
"<div style=\"clear: both;\"></div>\n",
|
||||
"\n",
|
||||
"<b>Share to:</b>\n",
|
||||
"\n",
|
||||
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
|
||||
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
||||
"</a>\n",
|
||||
"\n",
|
||||
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
|
||||
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
||||
"</a>\n",
|
||||
"\n",
|
||||
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
|
||||
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
||||
"</a>\n",
|
||||
"\n",
|
||||
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
|
||||
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
||||
"</a>\n",
|
||||
"\n",
|
||||
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/memory/get_started_with_memory_bank_adk.ipynb\" target=\"_blank\">\n",
|
||||
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
||||
"</a>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "84f0f73a0f76"
|
||||
},
|
||||
"source": [
|
||||
"| Author(s) |\n",
|
||||
"| --- |\n",
|
||||
"| [ Kimberly Milam, Ivan Nardini](https://github.com/inardini)|"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvgnzT1CKxrO"
|
||||
},
|
||||
"source": [
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"This tutorial demonstrates how to build an agent with long-term memory using the Google Agent Development Kit (ADK) and Vertex AI Memory Bank.\n",
|
||||
"\n",
|
||||
"Vertex AI Memory Bank enables you to dynamically generate and store long-term memories from user conversations. This allows an agent to access personalized information across multiple sessions for a particular user, leading to more contextual and continuous interactions.\n",
|
||||
"\n",
|
||||
"### What you'll learn\n",
|
||||
"\n",
|
||||
"By the end of this tutorial, you will be able to:\n",
|
||||
"* Understand how long-term memory enhances agent capabilities.\n",
|
||||
"* Create a Vertex AI Agent Engine instance to use Memory Bank.\n",
|
||||
"* Build an ADK agent that uses the `PreloadMemoryTool` tool to retrieve information.\n",
|
||||
"* See how the agent stores conversation history and recalls relevant facts from past sessions.\n",
|
||||
"* Build more sophisticated, context-aware Agents.\n",
|
||||
"\n",
|
||||
"### Why this is important\n",
|
||||
"\n",
|
||||
"Traditional LLM-based agents often lack the ability to recall information from previous interactions, treating each conversation as a new one. This \"amnesia\" prevents them from maintaining context over time. Long-term memory, like that provided by Vertex AI Memory Bank, allows agents to:\n",
|
||||
"\n",
|
||||
"* Maintain context over extended periods.\n",
|
||||
"* Personalize interactions based on user history.\n",
|
||||
"* Build a deeper understanding of user preferences and needs.\n",
|
||||
"\n",
|
||||
"Google's ADK provides a robust framework for building agents, and by extending it with services like Memory Bank, you can unlock a new level of sophistication.\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"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "tFy3H3aPgx12"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet \"google-cloud-aiplatform>=1.100.0\" \"google-adk>=1.5.0\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"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",
|
||||
"import uuid\n",
|
||||
"\n",
|
||||
"import vertexai\n",
|
||||
"\n",
|
||||
"# Project configuration\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",
|
||||
" if not PROJECT_ID:\n",
|
||||
" raise ValueError(\"Project ID not found. Please set it in the form above or as the GOOGLE_CLOUD_PROJECT environment variable.\")\n",
|
||||
"\n",
|
||||
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
|
||||
"\n",
|
||||
"# Set environment variables required for ADK\n",
|
||||
"os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"TRUE\"\n",
|
||||
"os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
|
||||
"os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
|
||||
"\n",
|
||||
"# Agent configuration\n",
|
||||
"MODEL_NAME = \"gemini-2.5-flash\"\n",
|
||||
"USER_ID = f\"user_{uuid.uuid4()}\"\n",
|
||||
"\n",
|
||||
"print(f\"Project: {PROJECT_ID}\")\n",
|
||||
"print(f\"Location: {LOCATION}\")\n",
|
||||
"print(f\"Session User ID: {USER_ID}\")\n",
|
||||
"\n",
|
||||
"# Initialize Vertex AI client\n",
|
||||
"_ = vertexai.Client(\n",
|
||||
" project=PROJECT_ID,\n",
|
||||
" location=LOCATION,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5303c05f7aa6"
|
||||
},
|
||||
"source": [
|
||||
"### Import libraries\n",
|
||||
"\n",
|
||||
"Import the Python libraries required for this tutorial, including components from the ADK, Vertex AI SDK, and standard libraries.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "6fc324893334"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import uuid\n",
|
||||
"import vertexai\n",
|
||||
"from vertexai import agent_engines\n",
|
||||
"from google import adk\n",
|
||||
"from google.adk.memory import VertexAiMemoryBankService\n",
|
||||
"from google.adk.sessions import VertexAiSessionService\n",
|
||||
"from google.genai import types"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qwEDAkpMqyR4"
|
||||
},
|
||||
"source": [
|
||||
"### Helper functions\n",
|
||||
"\n",
|
||||
"We will define a helper function to simplify running the agent and capturing its response. This will make our interaction loop cleaner and easier to read.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "PlnHiogXq0JP"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def run_single_turn(query, session, user_id):\n",
|
||||
" \"\"\"Run a single conversation turn.\"\"\"\n",
|
||||
" content = types.Content(role=\"user\", parts=[types.Part(text=query)])\n",
|
||||
" events = runner.run(user_id=user_id, session_id=session, new_message=content)\n",
|
||||
"\n",
|
||||
" response_content = None\n",
|
||||
" for event in events:\n",
|
||||
" if event.is_final_response():\n",
|
||||
" response_content = event.content.parts[0].text\n",
|
||||
" return response_content\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def chat_loop(session, user_id) -> None:\n",
|
||||
" \"\"\"Main chat interface loop.\"\"\"\n",
|
||||
" print(\"\\nStarting chat. Type 'exit' or 'quit' to end.\")\n",
|
||||
" print(\"Every message will be automatically stored in memory.\\n\")\n",
|
||||
"\n",
|
||||
" while True:\n",
|
||||
" user_input = input(\"\\nYou: \")\n",
|
||||
" if user_input.lower() in [\"quit\", \"exit\", \"bye\"]:\n",
|
||||
" print(\"\\nAssistant: Thank you for chatting. Have a great day!\")\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
" response = run_single_turn(user_input, session, user_id)\n",
|
||||
" if response:\n",
|
||||
" print(f\"\\nAssistant: {response}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EdvJRUWRNGHE"
|
||||
},
|
||||
"source": [
|
||||
"## Creating the Agent Engine with Memory Bank\n",
|
||||
"\n",
|
||||
"Vertex AI Memory Bank is a component of the Vertex AI Agent Engine, a managed service that allows developers to deploy, manage, and scale AI agents.\n",
|
||||
"\n",
|
||||
"To use Memory Bank, you first need to create or get an existing Agent Engine instance. This provides the necessary APIs to store and retrieve memories associated with specific users.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "x5tUDbo8YVhe"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agent_engine = agent_engines.create()\n",
|
||||
"print(f\"Created Agent Engine: {agent_engine.resource_name}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "NXtRpE0KXvY6"
|
||||
},
|
||||
"source": [
|
||||
"## Building an Agent with Memory Bank using ADK\n",
|
||||
"\n",
|
||||
"Now we will construct the core components of our memory-enabled agent."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "q0hfdld4X9yg"
|
||||
},
|
||||
"source": [
|
||||
"### Define the Agent\n",
|
||||
"\n",
|
||||
"Define the ADK agent with a name, instructions (prompt), and specify the `load_memory` tool. This built-in ADK tool enables the agent to call our `VertexAiMemoryBankService` to search for information. The agent's prompt is crucial, as it instructs the LLM when and how to use this tool to recall user-specific context."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "qz6Zmy_L6HwY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agent = adk.Agent(\n",
|
||||
" model=MODEL_NAME,\n",
|
||||
" name=\"helpful_assistant\",\n",
|
||||
" instruction=\"\"\"You are a helpful assistant with perfect memory.\n",
|
||||
" Instructions:\n",
|
||||
" - Use the context to personalize responses\n",
|
||||
" - Naturally reference past conversations when relevant\n",
|
||||
" - Build upon previous knowledge about the user\n",
|
||||
" - If using semantic search, the memories shown are the most relevant to the current query\"\"\",\n",
|
||||
" tools=[adk.tools.preload_memory_tool.PreloadMemoryTool()],\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "TnmlrAZkYJZ2"
|
||||
},
|
||||
"source": [
|
||||
"### Configure the ADK Runner\n",
|
||||
"\n",
|
||||
"The `Runner` orchestrates the interaction between the user, the agent, and the various services. We will configure it with our `agent` and the `VertexAiSessionService` and `VertexAiMemoryBankService` to manage session state and long-term memory.\n",
|
||||
"\n",
|
||||
"> Notice that the `VertexAiMemoryBankService` released in ADK, you can also use a local session service."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "wPEy9Sgp7YLV"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"app_name = \"my_agent_\" + str(uuid.uuid4())[:6]\n",
|
||||
"agent_engine_id = agent_engine.name\n",
|
||||
"\n",
|
||||
"memory_bank_service = VertexAiMemoryBankService(\n",
|
||||
" project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"session_service = VertexAiSessionService(\n",
|
||||
" project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"runner = adk.Runner(\n",
|
||||
" agent=agent,\n",
|
||||
" app_name=app_name,\n",
|
||||
" session_service=session_service,\n",
|
||||
" memory_service=memory_bank_service,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PMfVSv9EYQu_"
|
||||
},
|
||||
"source": [
|
||||
"## Interacting with the Agent - Information gathering session\n",
|
||||
"\n",
|
||||
"Let's begin our first conversation. In this session, we will provide the agent with specific pieces of information that we expect it to remember later.\n",
|
||||
"\n",
|
||||
"The process is as follows:\n",
|
||||
"\n",
|
||||
"1. Create a new ADK session for the user.\n",
|
||||
"2. Send a series of messages to the agent.\n",
|
||||
"3. After the conversation, retrieve the completed session data.\n",
|
||||
"4. Explicitly add the session to our Memory Bank, which triggers the memory generation process.\n",
|
||||
"\n",
|
||||
"In a production application, adding the session to memory might be handled by a background process or through hooks in a custom `Runner` or `SessionService`. For this tutorial, we perform this step explicitly to clearly illustrate the process.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "tRqZ0Rz2YdXL"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"session1 = await runner.session_service.create_session(\n",
|
||||
" app_name=app_name,\n",
|
||||
" user_id=USER_ID,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "c4WVP3nuaIV-"
|
||||
},
|
||||
"source": [
|
||||
"Try these example phrases:\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"You: Hi, I work as an agent engineer\n",
|
||||
"You: I love hiking and have a dog named Max\n",
|
||||
"You: I'm working on a recommendation system project\n",
|
||||
"You: Bye\n",
|
||||
"```\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-tqSiFUzaKBC"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chat_loop(session1.id, USER_ID)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Bzfm-l9caJkf"
|
||||
},
|
||||
"source": [
|
||||
"Add the session to memory."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "55_1rmY2QufY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"completed_session = await runner.session_service.get_session(\n",
|
||||
" app_name=app_name, user_id=USER_ID, session_id=session1.id\n",
|
||||
")\n",
|
||||
"await memory_bank_service.add_session_to_memory(completed_session)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gF0AX5rzZYXd"
|
||||
},
|
||||
"source": [
|
||||
"## Interacting with the Agent - Memory Recall session\n",
|
||||
"\n",
|
||||
"Now, we will start a new session with the same user. This time, we will ask questions that require the agent to recall information from the first session. This is where the `load_memory` tool and our `VertexAIMemoryBankService.search_memory()` method will be invoked."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "vzBwJL8JZlwa"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"session2 = await runner.session_service.create_session(\n",
|
||||
" app_name=agent_engine.name,\n",
|
||||
" user_id=USER_ID,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "rzdQxPYuZuZb"
|
||||
},
|
||||
"source": [
|
||||
"Try these example phrases:\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"You: What do you remember about me?\n",
|
||||
"You: What is my dog's name?\n",
|
||||
"You: Bye\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "rbbqJEmDZqvv"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chat_loop(session2.id, USER_ID)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "c7jVCYtvaxbD"
|
||||
},
|
||||
"source": [
|
||||
"## Cleaning up\n",
|
||||
"\n",
|
||||
"To avoid charges, delete the agent engine when you're done experimenting.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "CAtT65gEcAFB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"delete_engine = True\n",
|
||||
"\n",
|
||||
"if delete_engine:\n",
|
||||
" agent_engines.delete(resource_name=agent_engine.resource_name, force=True)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "get_started_with_memory_bank_adk.ipynb",
|
||||
"toc_visible": true
|
||||
},
|
||||
"environment": {
|
||||
"kernel": "python3",
|
||||
"name": "workbench-notebooks.m130",
|
||||
"type": "gcloud",
|
||||
"uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": ""
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,670 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ur8xi4C7S06n"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright 2025 Google LLC\n",
|
||||
"#\n",
|
||||
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||
"# you may not use this file except in compliance with the License.\n",
|
||||
"# You may obtain a copy of the License at\n",
|
||||
"#\n",
|
||||
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||
"#\n",
|
||||
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||
"# See the License for the specific language governing permissions and\n",
|
||||
"# limitations under the License."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JAPoU8Sm5E6e"
|
||||
},
|
||||
"source": [
|
||||
"# Get started with Vertex AI Memory Bank - CrewAI\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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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%2Fgemini%2Fagent-engine%2Fmemory%2Fget_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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/gemini/agent-engine/memory/get_started_with_memory_bank_crewai.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": [
|
||||
"| Author(s) |\n",
|
||||
"| --- |\n",
|
||||
"| [Ivan Nardini](https://github.com/inardini) |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvgnzT1CKxrO"
|
||||
},
|
||||
"source": [
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"This tutorial demonstrates how to integrate **Vertex AI Memory Bank** with **CrewAI** agents to provide them with persistent, long-term memory across conversations.\n",
|
||||
"\n",
|
||||
"### What you'll learn\n",
|
||||
"\n",
|
||||
"By the end of this tutorial, you will be able to:\n",
|
||||
"\n",
|
||||
" * Provision a Vertex AI Agent Engine to serve as the backend for Memory Bank.\n",
|
||||
" * Implement a custom CrewAI `Storage` class that connects to the Memory Bank API.\n",
|
||||
" * Instantiate CrewAI agents and tasks that leverage this external, long-term memory.\n",
|
||||
" * Assemble and run a crew that can recall information from past interactions to provide personalized responses.\n",
|
||||
"\n",
|
||||
"### Why this is important\n",
|
||||
"\n",
|
||||
"CrewAI provides a powerful framework for orchestrating autonomous agents to solve complex tasks. However, by default, these agents are stateless. By integrating a dedicated long-term memory service like Vertex AI Memory Bank, you can empower your crews to:\n",
|
||||
"\n",
|
||||
" * **Maintain Context:** Remember user preferences, historical data, and previous interactions across multiple sessions.\n",
|
||||
" * **Enhance Collaboration:** Allow different agents in a crew to access a shared pool of knowledge about a user or topic.\n",
|
||||
" * **Deliver Personalization:** Build a deep, evolving understanding of users, leading to more effective and personalized assistance.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "61RBz8LLbxCR"
|
||||
},
|
||||
"source": [
|
||||
"## Get started"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "No17Cw5hgx12"
|
||||
},
|
||||
"source": [
|
||||
"### Install required packages\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "tFy3H3aPgx12"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet \"crewai\" \"google-cloud-aiplatform>=1.100.0\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"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": "0Q9k7TmH59fd"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"import os\n",
|
||||
"import uuid\n",
|
||||
"\n",
|
||||
"import vertexai\n",
|
||||
"\n",
|
||||
"# Set project configuration\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",
|
||||
" if not PROJECT_ID:\n",
|
||||
" raise ValueError(\"Project ID not found. Please set it in the form above or as the GOOGLE_CLOUD_PROJECT environment variable.\")\n",
|
||||
"\n",
|
||||
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
|
||||
"SERVICE_ACCOUNT_FILE = \"[your-service-account-file]\" # @param {type:\"string\"}\n",
|
||||
"\n",
|
||||
"# Set environment variables for CrewAI\n",
|
||||
"os.environ[\"DEFAULT_VERTEXAI_PROJECT\"] = PROJECT_ID\n",
|
||||
"os.environ[\"DEFAULT_VERTEXAI_LOCATION\"] = LOCATION\n",
|
||||
"\n",
|
||||
"# Agent configuration\n",
|
||||
"MODEL_NAME = \"gemini-2.5-flash\"\n",
|
||||
"USER_ID = f\"user_{uuid.uuid4()}\"\n",
|
||||
"\n",
|
||||
"print(f\"Project: {PROJECT_ID}\")\n",
|
||||
"print(f\"Location: {LOCATION}\")\n",
|
||||
"print(f\"Session User ID: {USER_ID}\")\n",
|
||||
"\n",
|
||||
"# Initialize Vertex AI client\n",
|
||||
"client = vertexai.Client(\n",
|
||||
" project=PROJECT_ID,\n",
|
||||
" location=LOCATION,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Load the service account credentials with Vertex AI User permission for the CrewAI LLM\n",
|
||||
"with open(SERVICE_ACCOUNT_FILE, \"r\") as file:\n",
|
||||
" vertex_credentials = json.load(file)\n",
|
||||
"vertex_credentials_json = json.dumps(vertex_credentials)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5303c05f7aa6"
|
||||
},
|
||||
"source": [
|
||||
"### Import libraries\n",
|
||||
"\n",
|
||||
"Import the necessary libraries for building our agent crew."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "6fc324893334"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import datetime\n",
|
||||
"from typing import Any, Dict, List\n",
|
||||
"from crewai import LLM, Agent, Crew, Task\n",
|
||||
"from crewai.memory.external.external_memory import ExternalMemory\n",
|
||||
"from crewai.memory.storage.interface import Storage"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5INaBikJ89-k"
|
||||
},
|
||||
"source": [
|
||||
"### Helper functions\n",
|
||||
"\n",
|
||||
"These functions create an interactive chat experience. The `run_single_turn` function encapsulates the logic for a single round of conversation, dynamically creating a new `Task` and `Crew` for each user input while reusing the same persistent memory object."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Axwi9upN8-4C"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def run_single_turn(query: str, agent: Agent, memory: ExternalMemory):\n",
|
||||
" \"\"\"Runs a single conversational turn with the CrewAI agent.\"\"\"\n",
|
||||
" task = Task(\n",
|
||||
" description=f\"Respond to the user's query: '{query}'\",\n",
|
||||
" expected_output=\"A helpful and personalized response based on your memory of the user.\",\n",
|
||||
" agent=agent,\n",
|
||||
" )\n",
|
||||
" crew = Crew(\n",
|
||||
" agents=[agent],\n",
|
||||
" tasks=[task],\n",
|
||||
" external_memory=memory,\n",
|
||||
" )\n",
|
||||
" return crew.kickoff()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def chat_loop(agent: Agent, memory: ExternalMemory) -> None:\n",
|
||||
" \"\"\"Main chat interface loop.\"\"\"\n",
|
||||
" print(\"\\nStarting chat. Type 'exit' or 'quit' to end.\")\n",
|
||||
" print(\"Every message will be automatically stored in memory.\\n\")\n",
|
||||
"\n",
|
||||
" while True:\n",
|
||||
" user_input = input(\"\\nYou: \")\n",
|
||||
" if user_input.lower() in [\"quit\", \"exit\", \"bye\"]:\n",
|
||||
" print(\"\\nAssistant: Thank you for chatting. Have a great day!\")\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
" response = run_single_turn(user_input, agent, memory)\n",
|
||||
" print(f\"\\nAssistant: {response}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "uGTgIVPMA0kA"
|
||||
},
|
||||
"source": [
|
||||
"## Creating the Agent Engine with Memory Bank\n",
|
||||
"\n",
|
||||
"Vertex AI Memory Bank is part of the Vertex AI Agent Engine. To use Memory Bank, you first create a new Agent Engine instance, which provides the APIs to store and retrieve user-specific memories."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "8RWs1bTXA2ZM"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agent_engine = client.agent_engines.create()\n",
|
||||
"print(f\"Created Agent Engine: {agent_engine.api_resource.name}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "y8oPi1rEO55v"
|
||||
},
|
||||
"source": [
|
||||
"## (Optional) Initialize User Memory\n",
|
||||
"\n",
|
||||
"You can optionally pre-seed the Memory Bank with initial facts about a user. This demonstrates how to store information that the agent can reference from the very first interaction.\n",
|
||||
"We'll create an initial memory entry for our user. This demonstrates how to store facts that the agent can reference in future conversations.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "B6xRZ6-sg_7L"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import_initial_memory_op = client.agent_engines.create_memory(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" fact=f\"The user's name is {USER_ID} and they have a passion for learning new things.\",\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Pretty print the created memory\n",
|
||||
"memory = import_initial_memory_op.response\n",
|
||||
"print(\"Created initial memory:\")\n",
|
||||
"print(f\" Fact: {memory.fact}\")\n",
|
||||
"print(f\" User ID: {memory.scope.get('user_id')}\")\n",
|
||||
"print(f\" Created at: {memory.create_time}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "VZwY1np4TyjF"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"for i, memory in enumerate(memories, 1):\n",
|
||||
" if hasattr(memory, \"memory\") and hasattr(memory.memory, \"fact\"):\n",
|
||||
" print(f\"{i}. {memory.memory.fact}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Ogkz7RTi6_ly"
|
||||
},
|
||||
"source": [
|
||||
"## Integrating Memory Bank with CrewAI\n",
|
||||
"\n",
|
||||
"To make Vertex AI Memory Bank compatible with CrewAI's memory system, we need to create a bridge between them.\n",
|
||||
"\n",
|
||||
"### Create a Custom Memory Storage Class\n",
|
||||
"\n",
|
||||
"The key to this integration is a custom storage class, `VertexAIMemoryBankStorage`, that inherits from CrewAI's `Storage` interface. This class translates CrewAI's generic `save` and `search` methods into specific API calls to the Vertex AI Memory Bank.\n",
|
||||
"\n",
|
||||
" * **`save(value, metadata)`**: This method takes a piece of text (`value`) and associated `metadata` from CrewAI. It then calls the `generate_memories` endpoint of the Memory Bank API to save the information.\n",
|
||||
" * **`search(query, **kwargs)`**: This method receives a search `query` from CrewAI. It calls the `retrieve_memories` endpoint, using semantic search if a query is provided, and formats the results into the dictionary structure that CrewAI expects.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "qjpgkilG7HkY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class VertexAIMemoryBankStorage(Storage):\n",
|
||||
" \"\"\"CrewAI storage implementation using Vertex AI Memory Bank.\"\"\"\n",
|
||||
"\n",
|
||||
" def __init__(\n",
|
||||
" self,\n",
|
||||
" project: str,\n",
|
||||
" location: str,\n",
|
||||
" agent_engine_name: str,\n",
|
||||
" use_semantic_search: bool = True,\n",
|
||||
" default_user_id: str = \"default_user\",\n",
|
||||
" ):\n",
|
||||
" \"\"\"Initialize Vertex AI Memory Bank storage.\"\"\"\n",
|
||||
" self.project = project\n",
|
||||
" self.location = location\n",
|
||||
" self.agent_engine_name = agent_engine_name\n",
|
||||
" self.use_semantic_search = use_semantic_search\n",
|
||||
" self.default_user_id = default_user_id\n",
|
||||
"\n",
|
||||
" # Validate required parameters early\n",
|
||||
" if not all([project, location, agent_engine_name]):\n",
|
||||
" raise ValueError(\"project, location, and agent_engine_name are required\")\n",
|
||||
"\n",
|
||||
" def save(self, value: str, metadata: Dict[str, Any]) -> None:\n",
|
||||
" \"\"\"Save a memory to Vertex AI Memory Bank.\"\"\"\n",
|
||||
"\n",
|
||||
" client = self._get_client()\n",
|
||||
"\n",
|
||||
" # Extract user_id from metadata or use default\n",
|
||||
" user_id = metadata.get(\"user_id\", self.default_user_id)\n",
|
||||
"\n",
|
||||
" # Prepare memory event\n",
|
||||
" event = {\n",
|
||||
" \"content\": {\n",
|
||||
" \"parts\": [{\"text\": value}],\n",
|
||||
" \"role\": metadata.get(\"role\", \"user\"),\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" # Add timestamp if provided\n",
|
||||
" if \"timestamp\" in metadata:\n",
|
||||
" event[\"timestamp\"] = metadata[\"timestamp\"]\n",
|
||||
"\n",
|
||||
" # Generate memory in Vertex AI\n",
|
||||
" client.agent_engines.generate_memories(\n",
|
||||
" name=self.agent_engine_name,\n",
|
||||
" direct_contents_source={\"events\": [event]},\n",
|
||||
" scope={\"user_id\": user_id},\n",
|
||||
" config={\"wait_for_completion\": True},\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" def search(self, query: str, limit: int = 5, **kwargs) -> List[Dict[str, Any]]:\n",
|
||||
" \"\"\"Search memories in Vertex AI Memory Bank.\n",
|
||||
"\n",
|
||||
" Args:\n",
|
||||
" query: Search query (used for semantic search if enabled)\n",
|
||||
" limit: Maximum number of results to return\n",
|
||||
" **kwargs: Additional parameters (e.g., user_id)\n",
|
||||
"\n",
|
||||
" Returns:\n",
|
||||
" List of memory dictionaries with 'value' and 'metadata' keys\n",
|
||||
" \"\"\"\n",
|
||||
" client = self._get_client()\n",
|
||||
"\n",
|
||||
" # Extract user_id from kwargs or use default\n",
|
||||
" user_id = kwargs.get(\"user_id\", self.default_user_id)\n",
|
||||
" scope = {\"user_id\": user_id}\n",
|
||||
"\n",
|
||||
" # Determine search method\n",
|
||||
" if self.use_semantic_search and query:\n",
|
||||
" retrieved_memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=self.agent_engine_name,\n",
|
||||
" scope=scope,\n",
|
||||
" similarity_search_params={\"search_query\": query, \"top_k\": limit},\n",
|
||||
" )\n",
|
||||
" else:\n",
|
||||
" retrieved_memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=self.agent_engine_name, scope=scope, simple_retrieval_params={}\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # Transform results to CrewAI format\n",
|
||||
" results = []\n",
|
||||
" for memory in retrieved_memories:\n",
|
||||
" if hasattr(memory, \"memory\") and hasattr(memory.memory, \"fact\"):\n",
|
||||
" result = {\n",
|
||||
" \"memory\": memory.memory.fact,\n",
|
||||
" \"metadata\": {\n",
|
||||
" \"user_id\": user_id,\n",
|
||||
" \"timestamp\": (\n",
|
||||
" memory.memory.update_time.isoformat()\n",
|
||||
" if hasattr(memory.memory, \"update_time\")\n",
|
||||
" else datetime.now().isoformat()\n",
|
||||
" ),\n",
|
||||
" },\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" results.append(result)\n",
|
||||
"\n",
|
||||
" return results\n",
|
||||
"\n",
|
||||
" def reset(self):\n",
|
||||
" client = self._get_client()\n",
|
||||
" client.agent_engines.delete(name=self.agent_engine_name)\n",
|
||||
"\n",
|
||||
" def _get_client(self):\n",
|
||||
" \"\"\"Create a Vertex AI client instance.\n",
|
||||
"\n",
|
||||
" Returns a new client for each operation to ensure proper\n",
|
||||
" event loop management in async contexts.\n",
|
||||
" \"\"\"\n",
|
||||
" return vertexai.Client(project=self.project, location=self.location)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Jh5pdQOb7Ook"
|
||||
},
|
||||
"source": [
|
||||
"## Define the CrewAI Components\n",
|
||||
"\n",
|
||||
"With the storage bridge in place, we can now define the elements of our crew.\n",
|
||||
"\n",
|
||||
"1. **Instantiate the LLM**: We'll configure the Gemini model that our agent will use.\n",
|
||||
"2. **Instantiate the Custom Storage and External Memory**: We create an instance of our `VertexAIMemoryBankStorage` class, passing in the necessary project and Agent Engine details. Then we initialize an `ExternalMemory` object configured with our custom `vertex_memory_bank` storage. This equips the entire crew with long-term memory.\n",
|
||||
"3. **Create the Agent**: We define a `Personal Assistant` agent, providing it with a role, goal, backstory, and the configured LLM."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "1aDYi0KQR1k6"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Instantiate the LLM\n",
|
||||
"llm = LLM(model=MODEL_NAME, temperature=0.7, vertex_credentials=vertex_credentials_json)\n",
|
||||
"\n",
|
||||
"# Instantiate your custom storage and the external memory\n",
|
||||
"vertex_memory_bank_storage = VertexAIMemoryBankStorage(\n",
|
||||
" project=PROJECT_ID,\n",
|
||||
" location=LOCATION,\n",
|
||||
" agent_engine_name=agent_engine.api_resource.name,\n",
|
||||
" default_user_id=USER_ID,\n",
|
||||
")\n",
|
||||
"external_vertex_memory_bank = ExternalMemory(storage=vertex_memory_bank_storage)\n",
|
||||
"\n",
|
||||
"# Create a personal assistant agent\n",
|
||||
"assistant = Agent(\n",
|
||||
" role=\"Personal Assistant\",\n",
|
||||
" goal=\"Remember user information and provide personalized help\",\n",
|
||||
" backstory=\"You are a helpful assistant with perfect memory of past conversations.\",\n",
|
||||
" llm=llm,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "c4WVP3nuaIV-"
|
||||
},
|
||||
"source": [
|
||||
"## Interacting with the Agent\n",
|
||||
"\n",
|
||||
"Now, let's interact with our memory-enabled agent. Try the following phrases to see how the agent builds and uses its memory.\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"You: Hi\n",
|
||||
"You: I work as an agent engineer\n",
|
||||
"You: I love hiking and have a dog named Max\n",
|
||||
"You: I'm working on a recommendation system project\n",
|
||||
"You: What do you remember about me?\n",
|
||||
"You: Bye\n",
|
||||
"```\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-tqSiFUzaKBC"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chat_loop(assistant, external_vertex_memory_bank)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "6nbPcSzu7-Ol"
|
||||
},
|
||||
"source": [
|
||||
"## Check for Stored Memories\n",
|
||||
"\n",
|
||||
"After the run, we can query the Memory Bank directly to confirm that the information from our latest interaction (\"software engineer,\" \"love hiking\") has been successfully saved."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "r-0g4fvI8BHZ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"for i, memory in enumerate(memories, 1):\n",
|
||||
" if hasattr(memory, \"memory\") and hasattr(memory.memory, \"fact\"):\n",
|
||||
" print(f\"{i}. {memory.memory.fact}\")\n",
|
||||
"\n",
|
||||
"print(f\"\\nTotal memories stored: {len(list(memories))}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2a4e033321ad"
|
||||
},
|
||||
"source": [
|
||||
"## Cleaning Up\n",
|
||||
"\n",
|
||||
"To avoid incurring ongoing charges to your Google Cloud account for the resources used, it is important to delete the Agent Engine instance when you are finished."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "3p_JZNCq8QEy"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"delete_engine = True\n",
|
||||
"\n",
|
||||
"if delete_engine:\n",
|
||||
" client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "get_started_with_memory_bank_crewai.ipynb",
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,670 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ur8xi4C7S06n"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Copyright 2025 Google LLC\n",
|
||||
"#\n",
|
||||
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
||||
"# you may not use this file except in compliance with the License.\n",
|
||||
"# You may obtain a copy of the License at\n",
|
||||
"#\n",
|
||||
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
||||
"#\n",
|
||||
"# Unless required by applicable law or agreed to in writing, software\n",
|
||||
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
||||
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
||||
"# See the License for the specific language governing permissions and\n",
|
||||
"# limitations under the License."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JAPoU8Sm5E6e"
|
||||
},
|
||||
"source": [
|
||||
"# Get started with Vertex AI Memory Bank - LangGraph\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/tree/main/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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%2Fgemini%2Fagent-engine%2Fmemory%2Fget_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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/gemini/agent-engine/memory/get_started_with_memory_bank_langgraph.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": [
|
||||
"| Author(s) |\n",
|
||||
"| --- |\n",
|
||||
"| [Ivan Nardini](https://github.com/inardini) |"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvgnzT1CKxrO"
|
||||
},
|
||||
"source": [
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"This tutorial demonstrates how to build a conversational agent using **LangGraph** with the **Vertex AI Memory Bank API** to provide the agent with long-term memory. This enables the agent to personalize responses based on information learned from prior conversations.\n",
|
||||
"\n",
|
||||
"### What you'll learn\n",
|
||||
"\n",
|
||||
"By the end of this tutorial, you will be able to:\n",
|
||||
"\n",
|
||||
" * Provision a Vertex AI Agent Engine to use Memory Bank.\n",
|
||||
" * Define a LangGraph agent state and graph structure.\n",
|
||||
" * Create a graph node that retrieves memories, generates personalized responses, and stores new information.\n",
|
||||
" * Build and run a conversational agent that maintains context across sessions.\n",
|
||||
"\n",
|
||||
"### Why this is important\n",
|
||||
"\n",
|
||||
"While standard chatbot frameworks can manage short-term conversation history, they often fail to retain information across different sessions. LangGraph allows you to build robust, stateful applications by defining agent workflows as graphs. Integrating it with a dedicated long-term memory service like Vertex AI Memory Bank unlocks the ability to create truly personalized and context-aware agents that build a continuous understanding of their users over time.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "61RBz8LLbxCR"
|
||||
},
|
||||
"source": [
|
||||
"## Get started"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "No17Cw5hgx12"
|
||||
},
|
||||
"source": [
|
||||
"### Install required packages\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "tFy3H3aPgx12"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet google-cloud-aiplatform>=1.100.0 langgraph langchain-core langchain-google-vertexai"
|
||||
]
|
||||
},
|
||||
{
|
||||
"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 and enviroment\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",
|
||||
"import uuid\n",
|
||||
"import vertexai\n",
|
||||
"\n",
|
||||
"# Project configuration\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",
|
||||
" if not PROJECT_ID:\n",
|
||||
" raise ValueError(\"Project ID not found. Please set it in the form above or as the GOOGLE_CLOUD_PROJECT environment variable.\")\n",
|
||||
"\n",
|
||||
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
|
||||
"\n",
|
||||
"# Agent configuration\n",
|
||||
"MODEL_NAME = \"gemini-2.5-flash\"\n",
|
||||
"USER_ID = f\"user_{uuid.uuid4()}\"\n",
|
||||
"\n",
|
||||
"print(f\"Project: {PROJECT_ID}\")\n",
|
||||
"print(f\"Location: {LOCATION}\")\n",
|
||||
"print(f\"Session User ID: {USER_ID}\")\n",
|
||||
"\n",
|
||||
"# Initialize Vertex AI client\n",
|
||||
"client = vertexai.Client(\n",
|
||||
" project=PROJECT_ID,\n",
|
||||
" location=LOCATION,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5303c05f7aa6"
|
||||
},
|
||||
"source": [
|
||||
"### Import libraries\n",
|
||||
"\n",
|
||||
"Import the necessary libraries for building our conversational agent with LangGraph and Vertex AI."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "6fc324893334"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from typing import TypedDict, Annotated\n",
|
||||
"from functools import partial\n",
|
||||
"from langchain_core.messages import SystemMessage, HumanMessage\n",
|
||||
"from langchain_google_vertexai import ChatVertexAI\n",
|
||||
"from langgraph.graph import StateGraph, START, END, add_messages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZChEL1y9BIU7"
|
||||
},
|
||||
"source": [
|
||||
"### Helper functions\n",
|
||||
"\n",
|
||||
"These helper functions will manage the user-facing chat loop and process a single turn of the conversation with our LangGraph agent.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "F7mGHiul1NQe"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def run_single_turn(user_input: str, user_id: str) -> str:\n",
|
||||
" \"\"\"Run a single conversation turn.\"\"\"\n",
|
||||
" config = {\"configurable\": {\"thread_id\": user_id}}\n",
|
||||
" state = {\n",
|
||||
" \"messages\": [HumanMessage(content=user_input)],\n",
|
||||
" \"user_id\": user_id\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" response_content = None\n",
|
||||
" for event in compiled_graph.stream(state, config):\n",
|
||||
" for value in event.values():\n",
|
||||
" if value.get(\"messages\"):\n",
|
||||
" response_content = value[\"messages\"][-1].content\n",
|
||||
"\n",
|
||||
" return response_content\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def chat_loop(user_id: str) -> None:\n",
|
||||
" \"\"\"Main chat interface loop.\"\"\"\n",
|
||||
" print(\"\\nStarting chat. Type 'exit' or 'quit' to end.\")\n",
|
||||
" print(\"Every message will be automatically stored in memory.\\n\")\n",
|
||||
"\n",
|
||||
" while True:\n",
|
||||
" user_input = input(\"\\nYou: \")\n",
|
||||
" if user_input.lower() in ['quit', 'exit', 'bye']:\n",
|
||||
" print(\"\\nAssistant: Thank you for chatting. Have a great day!\")\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
" response = run_single_turn(user_input, user_id)\n",
|
||||
" if response:\n",
|
||||
" print(f\"\\nAssistant: {response}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EdvJRUWRNGHE"
|
||||
},
|
||||
"source": [
|
||||
"## Creating the Agent Engine with Memory Bank\n",
|
||||
"\n",
|
||||
"Vertex AI Memory Bank is part of the Vertex AI Agent Engine, a managed service that enables developers to deploy, manage, and scale AI agents.\n",
|
||||
"\n",
|
||||
"To use Memory Bank, you first create a new Agent Engine instance, which provides the APIs to store and retrieve user-specific memories.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "x5tUDbo8YVhe"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agent_engine = client.agent_engines.create()\n",
|
||||
"print(f\"Created Agent Engine: {agent_engine.api_resource.name}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "y8oPi1rEO55v"
|
||||
},
|
||||
"source": [
|
||||
"## (Optional) Initialize user memory\n",
|
||||
"\n",
|
||||
"You can optionally pre-seed the Memory Bank with initial facts about a user. This demonstrates how to store information that the agent can reference from the very first interaction.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "B6xRZ6-sg_7L"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import_initial_memory_op = client.agent_engines.create_memory(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" fact=f\"The user's name is {USER_ID} and they have a passion for learning new things.\",\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"memory = import_initial_memory_op.response\n",
|
||||
"print(\"Created initial memory:\")\n",
|
||||
"print(f\" Fact: {memory.fact}\")\n",
|
||||
"print(f\" User ID: {memory.scope.get('user_id')}\")\n",
|
||||
"print(f\" Created at: {memory.create_time}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "VZwY1np4TyjF"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"for i, memory in enumerate(memories, 1):\n",
|
||||
" if hasattr(memory, 'memory') and hasattr(memory.memory, 'fact'):\n",
|
||||
" print(f\"{i}. {memory.memory.fact}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ntRLeQdtCoj1"
|
||||
},
|
||||
"source": [
|
||||
"## Building the LangGraph Agent\n",
|
||||
"\n",
|
||||
"Now, we will construct our conversational agent. The agent will be designed to:\n",
|
||||
"\n",
|
||||
"1. Retrieve relevant memories before generating a response.\n",
|
||||
"2. Use these memories to personalize its conversation.\n",
|
||||
"3. Automatically store each conversation turn to build its long-term memory.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "opi14Ez06rNU"
|
||||
},
|
||||
"source": [
|
||||
"### Define the Agent State\n",
|
||||
"\n",
|
||||
"The `State` object is a central concept in LangGraph. It holds the data that persists and is passed between the nodes in our graph. For this agent, the state will track the conversation `messages` and the `user_id`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Z12_8xnH01WJ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class State(TypedDict):\n",
|
||||
" \"\"\"Agent state definition.\"\"\"\n",
|
||||
" messages: Annotated[list, add_messages]\n",
|
||||
" user_id: str"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KvZhpAbr6sWn"
|
||||
},
|
||||
"source": [
|
||||
"### Create the Core Logic: The Chatbot Node\n",
|
||||
"\n",
|
||||
"This node contains the main logic of our agent. For each turn, it performs a sequence of actions:\n",
|
||||
"\n",
|
||||
"1. **Retrieve Memories**: It queries the Memory Bank API to find memories relevant to the user's latest message. It can use either semantic search for relevance or chronological retrieval for a complete history.\n",
|
||||
"2. **Construct Prompt**: It injects the retrieved memories into the system prompt, providing the LLM with personalized context.\n",
|
||||
"3. **Generate Response**: It calls the Gemini model with the context-enhanced prompt to generate a response.\n",
|
||||
"4. **Store New Memories**: It sends the latest user message and agent response back to Memory Bank to be processed and stored for future conversations.\n",
|
||||
"\n",
|
||||
"This \"always store\" approach ensures the agent continuously learns from every interaction."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "1c3RfqnE0zte"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def chatbot_node(state: State, agent_engine_name: str, use_semantic_search: bool = True) -> dict:\n",
|
||||
" \"\"\"Main agent logic node - stores every interaction.\"\"\"\n",
|
||||
"\n",
|
||||
" user_id = state[\"user_id\"]\n",
|
||||
" user_message = state[\"messages\"][-1].content\n",
|
||||
"\n",
|
||||
" # Check user_message is string\n",
|
||||
" if not isinstance(user_message, str):\n",
|
||||
" raise ValueError(\"User message must be a string\")\n",
|
||||
"\n",
|
||||
" # Retrieve memories based on the selected method\n",
|
||||
" if use_semantic_search:\n",
|
||||
" # Use semantic search to find relevant memories based on the user's current message\n",
|
||||
" memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine_name,\n",
|
||||
" scope={\"user_id\": user_id},\n",
|
||||
" similarity_search_params={\n",
|
||||
" \"search_query\": user_message,\n",
|
||||
" \"top_k\": 10 # Retrieve top 10 most relevant memories\n",
|
||||
" }\n",
|
||||
" )\n",
|
||||
" else:\n",
|
||||
" # Retrieve all memories\n",
|
||||
" memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine_name,\n",
|
||||
" scope={\"user_id\": user_id},\n",
|
||||
" simple_retrieval_params={}\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # Extract facts from memory objects\n",
|
||||
" memory_facts = [\n",
|
||||
" memory.memory.fact\n",
|
||||
" for memory in memories\n",
|
||||
" if hasattr(memory, 'memory') and hasattr(memory.memory, 'fact')\n",
|
||||
" ]\n",
|
||||
"\n",
|
||||
" # Format memory context based on retrieval method\n",
|
||||
" if memory_facts:\n",
|
||||
" if use_semantic_search:\n",
|
||||
" memory_context = \"Relevant memories from previous conversations:\\n\" + \"\\n\".join(\n",
|
||||
" f\"- {fact}\" for fact in memory_facts\n",
|
||||
" )\n",
|
||||
" else:\n",
|
||||
" memory_context = \"Previous conversation history:\\n\" + \"\\n\".join(\n",
|
||||
" f\"- {fact}\" for fact in memory_facts\n",
|
||||
" )\n",
|
||||
" else:\n",
|
||||
" memory_context = \"No previous conversations found.\"\n",
|
||||
"\n",
|
||||
" # Create system prompt with memory context\n",
|
||||
" system_message = SystemMessage(\n",
|
||||
" content=f\"\"\"You are a helpful assistant with perfect memory.\n",
|
||||
"\n",
|
||||
" {memory_context}\n",
|
||||
"\n",
|
||||
" Instructions:\n",
|
||||
" - Use the context to personalize responses\n",
|
||||
" - Naturally reference past conversations when relevant\n",
|
||||
" - Build upon previous knowledge about the user\n",
|
||||
" - If using semantic search, the memories shown are the most relevant to the current query\"\"\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # Generate response using LLM\n",
|
||||
" messages = [system_message] + state[\"messages\"]\n",
|
||||
" llm = ChatVertexAI(model=MODEL_NAME, project=PROJECT_ID, location=LOCATION)\n",
|
||||
" response = llm.invoke(messages)\n",
|
||||
"\n",
|
||||
" # Store the conversation turn in memory\n",
|
||||
" operation = client.agent_engines.generate_memories(\n",
|
||||
" name=agent_engine_name,\n",
|
||||
" direct_contents_source={\n",
|
||||
" \"events\": [{\n",
|
||||
" \"content\": {\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"parts\": [{\"text\": user_message}],\n",
|
||||
" }\n",
|
||||
" }, {\n",
|
||||
" \"content\": {\n",
|
||||
" \"role\": \"model\",\n",
|
||||
" \"parts\": [{\"text\": response.content}],\n",
|
||||
" }\n",
|
||||
" }]\n",
|
||||
" },\n",
|
||||
" scope={\"user_id\": user_id},\n",
|
||||
" config={\"wait_for_completion\": True},\n",
|
||||
" )\n",
|
||||
" return {\"messages\": [response]}\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "apId5Vvi6_QA"
|
||||
},
|
||||
"source": [
|
||||
"### Build and Compile the Conversation Graph\n",
|
||||
"\n",
|
||||
"LangGraph uses a state graph to define the conversational flow. Our graph is simple: it has a single `chatbot` node that processes each turn. We connect the `START` of the graph to our node and the node's output to the `END`.\n",
|
||||
"\n",
|
||||
"We use `functools.partial` to bind the `agent_engine_name` and `use_semantic_search` arguments to our node function ahead of time."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "3G6Pax34R8tB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"graph_builder = StateGraph(State)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Create chatbot node with agent engine name and search preference bound\n",
|
||||
"chatbot_with_memory = partial(\n",
|
||||
" chatbot_node,\n",
|
||||
" agent_engine_name=agent_engine.api_resource.name,\n",
|
||||
" use_semantic_search=True\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Assemble and compile the graph\n",
|
||||
"graph_builder.add_node(\"chatbot\", chatbot_with_memory)\n",
|
||||
"graph_builder.add_edge(START, \"chatbot\")\n",
|
||||
"graph_builder.add_edge(\"chatbot\", END)\n",
|
||||
"\n",
|
||||
"compiled_graph = graph_builder.compile()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "c4WVP3nuaIV-"
|
||||
},
|
||||
"source": [
|
||||
"## Interacting with the Agent\n",
|
||||
"\n",
|
||||
"Now, let's interact with our memory-enabled agent. Try the following phrases to see how the agent builds and uses its memory.\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"You: Hi, I work as an agent engineer\n",
|
||||
"You: I love hiking and have a dog named Max\n",
|
||||
"You: I'm working on a recommendation system project\n",
|
||||
"You: What do you remember about me?\n",
|
||||
"You: Bye\n",
|
||||
"```\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-tqSiFUzaKBC"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chat_loop(USER_ID)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vLG-474VXfzy"
|
||||
},
|
||||
"source": [
|
||||
"## Check for Stored Memories\n",
|
||||
"\n",
|
||||
"After the conversation, we can retrieve all stored memories to see what the agent has learned about the user. This demonstrates that the information has been successfully processed and persisted by Vertex AI Memory Bank.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "8hpsMHAIOrxN"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(\"Stored Memories\")\n",
|
||||
"memories = client.agent_engines.retrieve_memories(\n",
|
||||
" name=agent_engine.api_resource.name,\n",
|
||||
" scope={\"user_id\": USER_ID},\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"for memory in memories:\n",
|
||||
" if hasattr(memory, 'memory') and hasattr(memory.memory, 'fact'):\n",
|
||||
" print(f\"{memory.memory.fact}\")\n",
|
||||
"\n",
|
||||
"# Show total memory count\n",
|
||||
"print(f\"\\nTotal memories stored: {len(list(memories))}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2a4e033321ad"
|
||||
},
|
||||
"source": [
|
||||
"## Cleaning up\n",
|
||||
"\n",
|
||||
"To avoid charges, delete the agent engine when you're done experimenting.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "CAtT65gEcAFB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"delete_engine = True\n",
|
||||
"\n",
|
||||
"if delete_engine:\n",
|
||||
" client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "get_started_with_memory_bank_langgraph.ipynb",
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
Reference in New Issue
Block a user