{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3YcBnq20nC6r"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "xU0F5ObiGgF4"
      },
      "source": [
        "# Building a Multi-Agent RAG Application with LangGraph and Agent Engine\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/tutorial_langgraph_rag_agent.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Run 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%2Ftutorial_langgraph_rag_agent.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>      \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_langgraph_rag_agent.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",
        "  <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/tutorial_langgraph_rag_agent.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\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/tutorial_langgraph_rag_agent.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/tutorial_langgraph_rag_agent.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/tutorial_langgraph_rag_agent.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/tutorial_langgraph_rag_agent.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/tutorial_langgraph_rag_agent.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": "4sA9r45YMz_O"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Xiaolong Yang](https://github.com/shawn-yang-google) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZft-jYpHmYv"
      },
      "source": [
        "## Overview\n",
        "\n",
        "[Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview) is a managed service in Vertex AI that helps you to build and deploy agent frameworks. It gives you the flexibility to choose how much reasoning you want to delegate to the LLM and how much you want to handle with customized code.\n",
        "\n",
        "RAG (Retrieval-Augmented Generation) is an AI framework that combines the strengths of traditional information retrieval systems (such as databases) with the capabilities of generative large language models (LLMs).  \n",
        "\n",
        "[LangGraph](https://langchain-ai.github.io/langgraph/) is a library for building stateful, multi-actor applications with LLMs, used to create agent and multi-agent workflows.\n",
        "\n",
        "This notebook demonstrates how to build, deploy, and test a LangGraph + RAG application using [Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview) in Vertex AI.\n",
        "\n",
        "\n",
        "## Context\n",
        "\n",
        "In previous tutorials:\n",
        "* [LangGraph application with Agent Engine](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_langgraph.ipynb?)\n",
        "You have learned how to combine LangGraph's workflow orchestration with the scalability of Vertex AI, which enables you to build custom generative AI applications.\n",
        "* [RAG application with Agent Engine and Cloud SQL for PostgreSQL](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_cloud_sql_pg_rag_agent.ipynb)\n",
        "By combining this extra knowledge with its own language skills, the AI can write text that is more accurate, up-to-date, and relevant to your specific needs.\n",
        "Your [LangChain](https://python.langchain.com/docs/get_started/introduction) agent uses an [Postgres Vector Store](https://github.com/googleapis/langchain-google-cloud-sql-pg-python/tree/main) to perform a similary search and retrieve related data to ground the LLM response.\n",
        "\n",
        "## Objectives\n",
        "\n",
        "In this tutorial, you will learn how to build and deploy an agent (model, tools, and agent) using the Vertex AI SDK for Python and Cloud SQL for PostgreSQL LangGraph integration.\n",
        "\n",
        "We're using the `Multi Agent Collaboration` [approach](https://blog.langchain.dev/langgraph-multi-agent-workflows/). \n",
        "This sample notebook could be adapted to use other multi-agent implementations described in the [link](https://blog.langchain.dev/langgraph-multi-agent-workflows/), such as the `Agent Supervisor` or other approaches. \n",
        "\n",
        "You will develop a LangGraph Application like: <img src=\"https://storage.googleapis.com/github-repo/generative-ai/gemini/agent-engine/images/tutorial_langgraph_rag.png\" alt=\"Image description\"> \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QL58mPu9Hw7g"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "1. In the Google Cloud console, on the project selector page, select or [create a Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects).\n",
        "2. [Make sure that billing is enabled for your Google Cloud project](https://cloud.google.com/billing/docs/how-to/verify-billing-enabled#console).\n",
        "3. Follow the instruction in [RAG application with Agent Engine and Cloud SQL for PostgreSQL](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_cloud_sql_pg_rag_agent.ipynb), set up Cloud SQL for PostgreSQL.\n",
        "\n",
        "### Required roles\n",
        "\n",
        "To get the permissions that you need to complete the tutorial, ask your administrator to grant you the [Owner](https://cloud.google.com/iam/docs/understanding-roles#owner) (`roles/owner`) IAM role on your project. For more information about granting roles, see [Manage access](https://cloud.google.com/iam/docs/granting-changing-revoking-access).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-RYpMytsZ882"
      },
      "source": [
        "### Install and import dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w_94DKOCX5pG"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet \\\n",
        "    \"google-cloud-aiplatform[agent_engines,langchain]\"==1.60.0 \\\n",
        "    langchain-google-cloud-sql-pg==0.6.1 \\\n",
        "    cloud-sql-python-connector>=1.9.0 \\\n",
        "    langchain-google-vertexai==1.0.4 \\\n",
        "    cloudpickle==3.0.0 \\\n",
        "    pydantic==2.7.4 \\\n",
        "    langgraph==0.0.51 \\\n",
        "    httpx==0.27.2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bThFamq9351N"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Import the necessary Python libraries. These libraries provide the tools we need to interact with LangGraph, Vertex AI, and other components of our application."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "KNt0YeKaMz_Q"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from typing import Literal\n",
        "import uuid\n",
        "\n",
        "from google.cloud import storage\n",
        "from langchain_core.documents import Document\n",
        "from langchain_core.messages import BaseMessage, HumanMessage\n",
        "from langchain_core.tools import tool\n",
        "from langchain_google_cloud_sql_pg import PostgresEngine, PostgresVectorStore\n",
        "from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings\n",
        "from langgraph.graph import END, MessageGraph\n",
        "from langgraph.prebuilt import ToolNode\n",
        "import vertexai\n",
        "from vertexai import agent_engines"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yPKXjZrFZuUZ"
      },
      "source": [
        "### Authenticate to Google Cloud\n",
        "\n",
        "Authenticate to Google Cloud as the IAM user logged into this notebook in order to access your Google Cloud Project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "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": "9aGBuLA7aQ6O"
      },
      "source": [
        "### Define project information\n",
        "\n",
        "Initialize `gcloud` with your Project ID and resource location. At this time, only `us-central1` is supported."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\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": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = \"gs://[your-staging-bucket]\"  # @param {type:\"string\"}\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)\n",
        "!gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_yG0kddIvr7"
      },
      "source": [
        "## Set up Cloud SQL\n",
        "\n",
        "You should have already set up Cloud SQL in [RAG application with Agent Engine and Cloud SQL for PostgreSQL](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_cloud_sql_pg_rag_agent.ipynb):\n",
        "* Enable APIs.\n",
        "* Create a Cloud SQL instance.\n",
        "* Create a database.\n",
        "* Initialize multiple vector store tables.\n",
        "* Create a user.\n",
        "\n",
        "\n",
        "In this Colab, we will create two new vector store tables: Book and Movie.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtiB5-LVVkv0"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "INSTANCE = \"langgraph-rag-instance\"  # @param {type:\"string\"}\n",
        "DATABASE = \"harry_potter_data\"  # @param {type:\"string\"}\n",
        "MOVIE_TABLE_NAME = \"my-movie\"  # @param {type:\"string\"}\n",
        "BOOK_TABLE_NAME = \"my-book\"  # @param {type:\"string\"}\n",
        "PASSWORD = input(\"Please provide a password to be used for 'postgres' database user: \")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dqDjyLpS5zCm"
      },
      "source": [
        "### Grant access to vector store table to IAM users"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "24NjnjF95ySA"
      },
      "outputs": [],
      "source": [
        "engine = await PostgresEngine.afrom_instance(\n",
        "    PROJECT_ID, REGION, INSTANCE, DATABASE, user=\"postgres\", password=PASSWORD\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OaP1LRhPi0y7"
      },
      "source": [
        "### Initialize multiple vector store tables\n",
        "\n",
        "The `PostgresEngine` has a helper method `init_vectorstore_table()` that can be used to create a table with the proper schema to store vector embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "GGd89YWIi2qg"
      },
      "outputs": [],
      "source": [
        "for table_name in [MOVIE_TABLE_NAME, BOOK_TABLE_NAME]:\n",
        "    engine = await PostgresEngine.afrom_instance(\n",
        "        PROJECT_ID, REGION, INSTANCE, DATABASE, user=\"postgres\", password=PASSWORD\n",
        "    )\n",
        "\n",
        "    await engine.ainit_vectorstore_table(\n",
        "        table_name=table_name,\n",
        "        vector_size=768,  # Vector size for VertexAI model(text-embedding-005)\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sQ1MI8ARi5Rr"
      },
      "source": [
        "### Add embeddings to the vector store\n",
        "\n",
        "Load data from a CSV file to generate and insert embeddings to the vector store.\n",
        "\n",
        "We will use two datasets:\n",
        "\n",
        "* Harry Potter Movie\n",
        "  - Intro: https://www.kaggle.com/datasets/rounakbanik/the-movies-dataset\n",
        "  - Data: gs://github-repo/generative-ai/gemini/agent-engine/sample_data/harry_potter_movies.json\n",
        "* Harry Potter Book\n",
        "  - Intro: https://www.kaggle.com/datasets/shubhammaindola/harry-potter-books\n",
        "  - Data: gs://github-repo/generative-ai/gemini/agent-engine/sample_data/harry_potter_books.json\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "Kcawj2ILdNmN"
      },
      "outputs": [],
      "source": [
        "def get_docs_from_gcs(bucket_name, gcs_dir, blob_name):\n",
        "    \"\"\"Fetches a JSON file from GCS, deserializes it, and returns the data.\n",
        "\n",
        "    Args:\n",
        "      bucket_name: Name of the GCS bucket (e.g., 'my-bucket').\n",
        "      gcs_dir: Directory within the bucket where the JSON file is located.\n",
        "      blob_name: Path and filename within the bucket\n",
        "                 (e.g., 'my_data.json').\n",
        "\n",
        "    Returns:\n",
        "      A Python object representing the Document, or None if the file\n",
        "      is not found or an error occurs.\n",
        "    \"\"\"\n",
        "\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    blob = bucket.blob(f\"{gcs_dir}/{blob_name}\")\n",
        "\n",
        "    if not blob.exists():\n",
        "        print(f\"File not found: gs://{bucket_name}/{gcs_dir}/{blob_name}\")\n",
        "        return None\n",
        "\n",
        "    try:\n",
        "        with blob.open(\"r\") as f:\n",
        "            json_docs = json.loads(f.read())\n",
        "    except json.JSONDecodeError:\n",
        "        print(f\"Error: Invalid JSON format in gs://{bucket_name}/{gcs_dir}/{blob_name}\")\n",
        "        return None\n",
        "\n",
        "    docs = []\n",
        "    for json_doc in json_docs:\n",
        "        docs.append(Document(**(json_doc[\"kwargs\"])))\n",
        "\n",
        "    return docs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pg5k6FyykfzW"
      },
      "source": [
        "#### Movies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dkMjEXEmi4ro"
      },
      "outputs": [],
      "source": [
        "# Initialize the vector store for movies\n",
        "vector_store = await PostgresVectorStore.create(\n",
        "    engine,\n",
        "    table_name=MOVIE_TABLE_NAME,\n",
        "    embedding_service=VertexAIEmbeddings(\n",
        "        model_name=\"text-embedding-005\", project=PROJECT_ID\n",
        "    ),\n",
        ")\n",
        "docs = get_docs_from_gcs(\n",
        "    \"github-repo\",\n",
        "    \"generative-ai/gemini/agent-engine/sample_data\",\n",
        "    \"harry_potter_movies.json\",\n",
        ")\n",
        "# Add data to the vector store\n",
        "ids = [str(uuid.uuid4()) for i in range(len(docs))]\n",
        "await vector_store.aadd_documents(docs, ids=ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GBdIVxqVkjT-"
      },
      "source": [
        "#### Books"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PWAElYbEkyzB"
      },
      "outputs": [],
      "source": [
        "# Initialize the vector store for books\n",
        "vector_store = await PostgresVectorStore.create(\n",
        "    engine,\n",
        "    table_name=BOOK_TABLE_NAME,\n",
        "    embedding_service=VertexAIEmbeddings(\n",
        "        model_name=\"text-embedding-005\", project=PROJECT_ID\n",
        "    ),\n",
        ")\n",
        "docs = get_docs_from_gcs(\n",
        "    \"github-repo\",\n",
        "    \"generative-ai/gemini/agent-engine/sample_data\",\n",
        "    \"harry_potter_books.json\",\n",
        ")\n",
        "# Add data to the vector store\n",
        "ids = [str(uuid.uuid4()) for i in range(len(docs))]\n",
        "await vector_store.aadd_documents(docs, ids=ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XCra5kJVKyg5"
      },
      "source": [
        "## Define the retriever tool\n",
        "\n",
        "Tools are interfaces that an agent, chain, or LLM can use to enable the Gemini model to interact with external systems, databases, document stores, and other APIs so that the model can get the most up-to-date information or take action with those systems.\n",
        "\n",
        "In this example, you'll define a function that will retrieve similar documents from the vector store using semantic search."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "vLx7O_UdqDyr"
      },
      "outputs": [],
      "source": [
        "@tool\n",
        "def movie_similarity_search(query: str) -> str:\n",
        "    \"\"\"\n",
        "    Perform a similarity search for movies based on the user's last message.\n",
        "\n",
        "    Args:\n",
        "        query str: The current conversation state, where the last message contains the query.\n",
        "\n",
        "    Returns:\n",
        "        str: A list of BaseMessage containing the search results.\n",
        "    \"\"\"\n",
        "    engine = PostgresEngine.from_instance(\n",
        "        PROJECT_ID,\n",
        "        REGION,\n",
        "        INSTANCE,\n",
        "        DATABASE,\n",
        "        quota_project=PROJECT_ID,\n",
        "        user=\"postgres\",\n",
        "        password=PASSWORD,\n",
        "    )\n",
        "\n",
        "    vector_store = PostgresVectorStore.create_sync(\n",
        "        engine,\n",
        "        table_name=MOVIE_TABLE_NAME,\n",
        "        embedding_service=VertexAIEmbeddings(\n",
        "            model_name=\"text-embedding-005\", project=PROJECT_ID\n",
        "        ),\n",
        "    )\n",
        "    retriever = vector_store.as_retriever()\n",
        "    return str([doc for doc in retriever.invoke(query)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "XVjf2fT_bIGa"
      },
      "outputs": [],
      "source": [
        "@tool\n",
        "def book_similarity_search(query: str) -> str:\n",
        "    \"\"\"\n",
        "    Perform a similarity search for books based on the user's last message.\n",
        "\n",
        "    Args:\n",
        "        state (List[BaseMessage]): The current conversation state, where the last message contains the query.\n",
        "\n",
        "    Returns:\n",
        "        List[BaseMessage]: A list of BaseMessage containing the search results.\n",
        "    \"\"\"\n",
        "    engine = PostgresEngine.from_instance(\n",
        "        PROJECT_ID,\n",
        "        REGION,\n",
        "        INSTANCE,\n",
        "        DATABASE,\n",
        "        quota_project=PROJECT_ID,\n",
        "        # Uncomment to use built-in authentication instead of IAM authentication\n",
        "        user=\"postgres\",\n",
        "        password=PASSWORD,\n",
        "    )\n",
        "\n",
        "    vector_store = PostgresVectorStore.create_sync(\n",
        "        engine,\n",
        "        table_name=BOOK_TABLE_NAME,\n",
        "        embedding_service=VertexAIEmbeddings(\n",
        "            model_name=\"text-embedding-005\", project=PROJECT_ID\n",
        "        ),\n",
        "    )\n",
        "    retriever = vector_store.as_retriever()\n",
        "    return str([doc for doc in retriever.invoke(query)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF8xqd84351O"
      },
      "source": [
        "### Define router\n",
        "\n",
        "We're using the `Multi Agent Collaboration` [approach](https://blog.langchain.dev/langgraph-multi-agent-workflows/). \n",
        "This sample notebook could be adapted to use other multi-agent implementations described in the [link](https://blog.langchain.dev/langgraph-multi-agent-workflows/), such as the `Agent Supervisor` or other approaches. \n",
        "\n",
        "Then, you'll define a router to control the flow of the conversation, determining which tool to use based on user input or the state of the interaction. Here we'll use a simple router setup, and you can customize the behavior of your router to handle multiple tools, custom logic, or multi-agent workflows.\n",
        "\n",
        "In this example, the router will invoke different nodes in the graph based on whether the user prompt contains the word 'book' or 'movie'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "7m9XRSn3351O"
      },
      "outputs": [],
      "source": [
        "def router(\n",
        "    state: list[BaseMessage],\n",
        ") -> Literal[\"book_similarity_search\", \"movie_similarity_search\", \"__end__\"]:\n",
        "    if not state[0].content or len(state[1].tool_calls) == 0:\n",
        "        return \"__end__\"\n",
        "    if \"book\" in state[0].content:\n",
        "        return \"book_similarity_search\"\n",
        "    if \"movie\" in state[0].content:\n",
        "        return \"movie_similarity_search\"\n",
        "    return \"__end__\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CFeOIbed351O"
      },
      "source": [
        "## Define LangGraph application\n",
        "\n",
        "Now you'll bring everything together to define your LangGraph application as a custom template in Agent Engine.\n",
        "\n",
        "This application will use the tool and router that you just defined. LangGraph provides a powerful way to structure these interactions and leverage the capabilities of LLMs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZYtR6-zqudb"
      },
      "source": [
        "#### Multi stage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "WWjFaLeW351O"
      },
      "outputs": [],
      "source": [
        "class MultiStageLangGraphApp:\n",
        "    def __init__(self, project: str, location: str) -> None:\n",
        "        self.project_id = project\n",
        "        self.location = location\n",
        "\n",
        "    # The set_up method is used to define application initialization logic\n",
        "    def set_up(self) -> None:\n",
        "        model = ChatVertexAI(model=\"gemini-2.0-flash\")\n",
        "        builder = MessageGraph()\n",
        "\n",
        "        # Checker node\n",
        "        def checker(state: list[BaseMessage]):\n",
        "            if not state[0].content:\n",
        "                return \"__end__\"\n",
        "            user_question = state[0].content\n",
        "            response = model.invoke(\n",
        "                [\n",
        "                    HumanMessage(\n",
        "                        content=(\n",
        "                            f\"What is the type of the question? {user_question}\"\n",
        "                            \"Think step by step, then answer one of the following:\"\n",
        "                            \"* <Answer>movie</Answer>\"\n",
        "                            \"* <Answer>book</Answer>\"\n",
        "                            \"* <Answer>no</Answer>\"\n",
        "                        )\n",
        "                    )\n",
        "                ]\n",
        "            )\n",
        "            table_name = response.content.split(\"<Answer>\")[1].split(\"</Answer>\")[0]\n",
        "            # Multiturn requests alternate between user and model.\n",
        "            state[0].content = f\"query:{state[0].content},table_name:{table_name}\"\n",
        "\n",
        "        builder.add_node(\"checker\", checker)\n",
        "        # Set entry point to checker node so it is reachable\n",
        "        builder.set_entry_point(\"checker\")\n",
        "\n",
        "        # Tool node.\n",
        "        model_with_tools = model.bind_tools(\n",
        "            [book_similarity_search, movie_similarity_search]\n",
        "        )\n",
        "        builder.add_node(\"tools\", model_with_tools)\n",
        "        # Add edge from tools to checker so the flow is checker->tools->router...\n",
        "        builder.add_edge(\"checker\", \"tools\")\n",
        "\n",
        "        # Summerize node.\n",
        "        # node\n",
        "        def summerizar(state: list[BaseMessage]):\n",
        "            question = state[0].content\n",
        "            related_docs = state[-1].content\n",
        "            response = model.invoke(\n",
        "                [\n",
        "                    HumanMessage(\n",
        "                        content=(\n",
        "                            f\"\"\"\n",
        "                            Use the docs: {related_docs} to answer question:{question}.\n",
        "                            The answer format should be json dict.\n",
        "                            \"\"\"\n",
        "                        )\n",
        "                    )\n",
        "                ]\n",
        "            )\n",
        "            # Multiturn requests alternate between user and model.\n",
        "            state.append(response)\n",
        "\n",
        "        builder.add_node(\"summerizar_node\", summerizar)\n",
        "        builder.add_edge(\"summerizar_node\", END)\n",
        "        # Book retrieval node\n",
        "        book_node = ToolNode([book_similarity_search])\n",
        "        builder.add_node(\"book_similarity_search\", book_node)\n",
        "        builder.add_edge(\"book_similarity_search\", \"summerizar_node\")\n",
        "\n",
        "        # Movie retrieval node\n",
        "        movie_node = ToolNode([movie_similarity_search])\n",
        "        builder.add_node(\"movie_similarity_search\", movie_node)\n",
        "        builder.add_edge(\"movie_similarity_search\", \"summerizar_node\")\n",
        "\n",
        "        # Router to check condition.\n",
        "        builder.add_conditional_edges(\"tools\", router)\n",
        "\n",
        "        self.runnable = builder.compile()\n",
        "\n",
        "    # The query method will be used to send inputs to the agent\n",
        "    def query(self, message: str):\n",
        "        \"\"\"Query the application.\n",
        "\n",
        "        Args:\n",
        "            message: The user message.\n",
        "\n",
        "        Returns:\n",
        "            str: The LLM response.\n",
        "        \"\"\"\n",
        "        chat_history = self.runnable.invoke(HumanMessage(message))\n",
        "\n",
        "        return chat_history[-1].content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sEfQYtgSm9ol"
      },
      "source": [
        "### Local test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "IcWux9IskE-c"
      },
      "outputs": [],
      "source": [
        "agent = MultiStageLangGraphApp(project=PROJECT_ID, location=LOCATION)\n",
        "agent.set_up()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0aaf11c1677a"
      },
      "source": [
        "Expect a JSON format answer like \n",
        "```json\n",
        "{\"company\": [\"Warner Bros.\", \"Heyday Films\"]}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J9yUujSokJpQ"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "```json\n",
              "{\n",
              "  'answer': 'Warner Bros and Heyday Films produce Harry Potter and the Deathly Hallows: Part 2.'\n",
              "}\n",
              "```"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent.query(message=\"Which company produces and distributes Harry Potter films\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "077f1396f641"
      },
      "source": [
        "Expect a JSON format answer like \n",
        "```json\n",
        "{\n",
        "    \"answer\": [\"Daniel Radcliffe\", \"Darren Criss\"]\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vn1wBUEyLGSG"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "```json\n",
            "[\n",
            " {\n",
            "  \"actor\": \"Darren Criss\",\n",
            "  \"movie\": \"A Very Potter Musical\"\n",
            " },\n",
            " {\n",
            "  \"actor\": \"Daniel Radcliffe\",\n",
            "  \"movie\": \"Harry Potter and the Deathly Hallows: Part 2\"\n",
            " }\n",
            "]\n",
            "```\n"
          ]
        }
      ],
      "source": [
        "agent.query(message=\"Who acts as Harry Potter\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df468170bc6d"
      },
      "source": [
        "Expect a JSON format answer like \n",
        "```json\n",
        "{\n",
        "    \"answer\": \"Harry Potter and the Chamber of Secrets.\"\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EPGRJjdEb228"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'```json\n",
              "{\n",
              " \"book\": \"Harry Potter and the Chamber of Secrets\"\n",
              "}\n",
              "```'"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "agent.query(message=\"In which book Harry Potter drives car\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Building and deploying a LangGraph app on Agent Engine\n",
        "\n",
        "In the following sections, we'll walk through the process of building and deploying a LangGraph application using Agent Engine in Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ERxxgFTcI3DC"
      },
      "source": [
        "## Deploy the service\n",
        "\n",
        "Now that you've specified a model, tools, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!\n",
        "\n",
        "Here, you'll use the LangChain agent template provided in the Vertex AI SDK for Agent Engine, which brings together the model, tools, and reasoning that you've built up so far."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k2nGSr2_JWcc"
      },
      "outputs": [],
      "source": [
        "remote_app = agent_engines.create(\n",
        "    MultiStageLangGraphApp(project=PROJECT_ID, location=LOCATION),\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[agent_engines,langchain]==1.60.0\",\n",
        "        \"langchain-google-cloud-sql-pg==0.6.1\",\n",
        "        \"cloud-sql-python-connector>=1.9.0\",\n",
        "        \"langchain-google-vertexai==1.0.4\",\n",
        "        \"cloudpickle==3.0.0\",\n",
        "        \"pydantic==2.7.4\",\n",
        "        \"langgraph==0.0.51\",\n",
        "        \"httpx==0.27.2\",\n",
        "    ],\n",
        "    display_name=\"Agent Engine with LangGraph RAG Agent\",\n",
        "    description=\"This is a sample custom application in Agent Engine that uses LangGraph and sql pg rag\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TYqMpB16I4iH"
      },
      "source": [
        "## Try it out\n",
        "\n",
        "Query the remote app directly or retrieve the application endpoint via the resource ID or display name. The endpoint can be used from any Python environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5A-5oNnQjePC"
      },
      "source": [
        "### Ask question that can only be answered by the movie."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P9-7ZcQugWkJ"
      },
      "outputs": [],
      "source": [
        "response = remote_app.query(message=\"Who acts as Harry Potter\")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UuACdm6zgaLm"
      },
      "outputs": [],
      "source": [
        "response = remote_app.query(\n",
        "    message=\"Which company produces and distributes Harry Potter film\"\n",
        ")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oXVavFQfi9Hz"
      },
      "source": [
        "### Ask question that can only be answered by the book."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NxiCmdMnalY-"
      },
      "outputs": [],
      "source": [
        "response = remote_app.query(message=\"In which book Harry Potter drives car\")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MrZ9IjnAI5v9"
      },
      "source": [
        "## Clean up\n",
        "\n",
        "If you created a new project for this tutorial, delete the project. If you used an existing project and wish to keep it without the changes added in this tutorial, delete resources created for the tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tBc48ZHOJS6J"
      },
      "source": [
        "### Deleting the project\n",
        "\n",
        "The easiest way to eliminate billing is to delete the project that you created for the tutorial.\n",
        "\n",
        "1. In the Google Cloud console, go to the [Manage resources](https://console.cloud.google.com/iam-admin/projects?_ga=2.235586881.1783688455.1719351858-1945987529.1719351858) page.\n",
        "1. In the project list, select the project that you want to delete, and then click Delete.\n",
        "1. In the dialog, type the project ID, and then click Shut down to delete the project.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ed-BFtW-JPbI"
      },
      "source": [
        "### Deleting tutorial resources\n",
        "\n",
        "Delete the Agent Engine instance(s) and Cloud SQL instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgNlHrxkb6c-"
      },
      "outputs": [],
      "source": [
        "# Delete the AgentEngine instance\n",
        "remote_app.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goyrqS2_I8Hs"
      },
      "outputs": [],
      "source": [
        "# Or delete all Agent Engine apps\n",
        "apps = agent_engines.list()\n",
        "for app in apps:\n",
        "    app.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "odvj8aKpb3Wi"
      },
      "outputs": [],
      "source": [
        "# Delete the Cloud SQL instance\n",
        "!gcloud sql instances delete {INSTANCE} \\\n",
        "  --project={PROJECT_ID}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "-RYpMytsZ882",
        "R5Xep4W9lq-Z",
        "OaP1LRhPi0y7",
        "GBdIVxqVkjT-"
      ],
      "name": "tutorial_langgraph_rag_agent.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
