gemini/reasoning-engine/tutorial_alloydb_rag_agent.ipynb

{
  "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": [
        "# Deploying a RAG Application with AlloyDB to LangChain on Vertex AI\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/reasoning-engine/tutorial_alloydb_rag_agent.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" 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%2Freasoning-engine%2Ftutorial_alloydb_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/reasoning-engine/tutorial_alloydb_rag_agent.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" 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/reasoning-engine/tutorial_alloydb_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/reasoning-engine/tutorial_alloydb_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/reasoning-engine/tutorial_alloydb_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/reasoning-engine/tutorial_alloydb_rag_agent.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/53/X_logo_2023_original.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/reasoning-engine/tutorial_alloydb_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/reasoning-engine/tutorial_alloydb_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": "5988a2fe325a"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Averi Kitsch](https://github.com/averikitsch) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZft-jYpHmYv"
      },
      "source": [
        "## Overview\n",
        "\n",
        "[LangChain on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview)\n",
        "is a managed service that helps you to build and deploy LangChain apps to a managed Reasoning Engine runtime.\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).  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",
        "\n",
        "## Objectives\n",
        "\n",
        "In this tutorial, you will learn how to build and deploy an agent (model, tools, and reasoning) using the Vertex AI SDK for Python and AlloyDB for PostgreSQL LangChain integration.\n",
        "\n",
        "Your [LangChain](https://python.langchain.com/docs/get_started/introduction) agent will use an [AlloyDB Vector Store](https://github.com/googleapis/langchain-google-alloydb-pg-python/tree/main) to perform a similary search and retrieve related data to ground the LLM response.\n",
        "\n",
        "* Install and set up the AlloyDB for PostgreSQL for LangChain and the Vertex AI Python SDKs\n",
        "* Create an AlloyDB cluster and instance\n",
        "* Create an AlloyDB database user\n",
        "* Define a retriever to perform similarity searches\n",
        "* Use the LangChain agent template provided in the Vertex AI SDK for Reasoning Engine\n",
        "* Deploy and test your agent on Reasoning Engine in Vertex AI\n",
        "\n",
        "⚠️ Note: LangChain on Vertex AI currently only supports public IP for AlloyDB due to VPC network restrictions. "
      ]
    },
    {
      "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",
        "1. [Make sure that billing is enabled for your Google Cloud project](https://cloud.google.com/billing/docs/how-to/verify-billing-enabled#console).\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 --quiet \"langchain-google-alloydb-pg>=0.7.0\" \"google-cloud-aiplatform[reasoningengine,langchain]\" langchain-google-vertexai langchain-community"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2520bbf159a9"
      },
      "outputs": [],
      "source": [
        "import uuid\n",
        "\n",
        "from langchain_community.document_loaders.csv_loader import CSVLoader\n",
        "from langchain_core.documents import Document\n",
        "from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore\n",
        "from langchain_google_vertexai import VertexAIEmbeddings\n",
        "from sqlalchemy import text  # noqa: F401\n",
        "import vertexai\n",
        "from vertexai.preview import reasoning_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": 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": "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": "code",
      "execution_count": null,
      "metadata": {
        "id": "vIeI4T_XVcDA"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"my-project\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"\n",
        "\n",
        "!gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esrdZL3IL6LN"
      },
      "source": [
        "## Create a Cloud Storage bucket\n",
        "\n",
        "Create or reuse and existing Cloud Storage bucket. Reasoning engine stages the artifacts of your applications in a Cloud Storage bucket as part of the deployment process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sptkevO4aUT1"
      },
      "outputs": [],
      "source": [
        "STAGING_BUCKET_NAME = \"my-project-bucket\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = f\"gs://{STAGING_BUCKET_NAME}\"\n",
        "\n",
        "# Create a Cloud Storage bucket, if it doesn't already exist\n",
        "!gsutil mb -c standard {STAGING_BUCKET}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "veGoLZBYZjxY"
      },
      "source": [
        "### Enable APIs\n",
        "\n",
        "This tutorial uses the following billable components of Google Cloud, which you'll need to enable for this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PcKjP3PiXDIi"
      },
      "outputs": [],
      "source": [
        "!gcloud services enable aiplatform.googleapis.com alloydb.googleapis.com servicenetworking.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_yG0kddIvr7"
      },
      "source": [
        "## Set up AlloyDB\n",
        "\n",
        "Use the provided variable names or update the values to use a pre-exisiting AlloyDB cluster and instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtiB5-LVVkv0"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "CLUSTER = \"my-cluster\"  # @param {type:\"string\"}\n",
        "INSTANCE = \"my-instance\"  # @param {type:\"string\"}\n",
        "DATABASE = \"my_database\"  # @param {type:\"string\"}\n",
        "TABLE_NAME = \"my_test_table\"  # @param {type:\"string\"}\n",
        "PASSWORD = input(\"Please provide a password to be used for 'postgres' database user: \")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WwwSQ2ZUf51F"
      },
      "source": [
        "### Create an AlloyDB cluster and primary instance\n",
        "\n",
        "This tutorial requires a AlloyDB cluster and instance with public IP and IAM authentication enabled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LdbQ2Q0wdNiy"
      },
      "outputs": [],
      "source": [
        "# Create a cluster\n",
        "!gcloud alloydb clusters create {CLUSTER} \\\n",
        "  --database-version=POSTGRES_15 \\\n",
        "  --password={PASSWORD} \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --enable-private-service-connect"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8HU5u9lLgQSr"
      },
      "outputs": [],
      "source": [
        "# Create an instance\n",
        "!gcloud alloydb instances create {INSTANCE} \\\n",
        "  --instance-type=PRIMARY \\\n",
        "  --cpu-count=2 \\\n",
        "  --region={REGION} \\\n",
        "  --cluster={CLUSTER} \\\n",
        "  --project={PROJECT_ID}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nR1sw5iAh-7w"
      },
      "outputs": [],
      "source": [
        "# Enable Public IP and IAM authentication\n",
        "!gcloud beta alloydb instances update {INSTANCE} \\\n",
        "  --cluster={CLUSTER} \\\n",
        "  --region={REGION} \\\n",
        "  --database-flags=password.enforce_complexity=yes,alloydb.iam_authentication=on \\\n",
        "  --assign-inbound-public-ip=ASSIGN_IPV4"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "maqL163MOPu3"
      },
      "source": [
        "**Wait for the update operation to complete!**\n",
        "\n",
        "The operation status can be checked using the following `gcloud` command and replacing `<OPERATION_ID>` with the ID from the above output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eXfwXE6BCFmt"
      },
      "outputs": [],
      "source": [
        "# Wait for the update operation to complete!\n",
        "# Copy the Operation ID from above e.g. operation-1719355575046-61bbeaa4b10a5-82ddd8fb-a6368d07\n",
        "!gcloud alloydb operations describe --region {REGION} operation-1719869774125-61c3662f257cc-a89c23c4-e59e5f40"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZIjNIc71ixye"
      },
      "source": [
        "### Create a database\n",
        "\n",
        "Create a new database for the application using the AlloyDB for LangChain library to establish a connection pool using the `AlloyDBEngine`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnNy1NFVk9F7"
      },
      "source": [
        "By default, [IAM database authentication](https://cloud.google.com/alloydb/docs/connect-iam) will be used as the method of database authentication. This library uses the IAM principal belonging to the [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials) sourced from the environment.\n",
        "\n",
        "However, to smooth the onboarding process this tutorial will use the [built-in database authentication](https://cloud.google.com/alloydb/docs/database-users/about) using a username and password to access the AlloyDB database can also be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pn4c8oZtixLI"
      },
      "outputs": [],
      "source": [
        "engine = await AlloyDBEngine.afrom_instance(\n",
        "    PROJECT_ID,\n",
        "    REGION,\n",
        "    CLUSTER,\n",
        "    INSTANCE,\n",
        "    database=\"postgres\",\n",
        "    user=\"postgres\",\n",
        "    password=PASSWORD,\n",
        ")\n",
        "\n",
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(\"COMMIT\"))\n",
        "    await conn.execute(text(f\"CREATE DATABASE {DATABASE}\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OaP1LRhPi0y7"
      },
      "source": [
        "### Initialize a vector store table\n",
        "\n",
        "The `AlloyDBEngine` 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": null,
      "metadata": {
        "id": "GGd89YWIi2qg"
      },
      "outputs": [],
      "source": [
        "engine = await AlloyDBEngine.afrom_instance(\n",
        "    PROJECT_ID, REGION, CLUSTER, 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(textembedding-gecko@latest)\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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BqXo2nNYPTXV"
      },
      "outputs": [],
      "source": [
        "# Retrieve the CSV file\n",
        "!gsutil cp gs://github-repo/generative-ai/gemini/reasoning-engine/movies.csv ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5tGaHva7r4Kc"
      },
      "outputs": [],
      "source": [
        "# Load the CSV file\n",
        "metadata = [\n",
        "    \"show_id\",\n",
        "    \"type\",\n",
        "    \"country\",\n",
        "    \"date_added\",\n",
        "    \"release_year\",\n",
        "    \"rating\",\n",
        "    \"duration\",\n",
        "    \"listed_in\",\n",
        "]\n",
        "loader = CSVLoader(file_path=\"/content/movies.csv\", metadata_columns=metadata)\n",
        "docs = loader.load()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dkMjEXEmi4ro"
      },
      "outputs": [],
      "source": [
        "# Initialize the vector store\n",
        "vector_store = await AlloyDBVectorStore.create(\n",
        "    engine,\n",
        "    table_name=TABLE_NAME,\n",
        "    embedding_service=VertexAIEmbeddings(\n",
        "        model_name=\"textembedding-gecko@latest\", project=PROJECT_ID\n",
        "    ),\n",
        ")\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": "e9ZfPaG9FGj9"
      },
      "source": [
        "### Create a user\n",
        "\n",
        "Set up the AI Platform Reasoning Engine Service Agent service account (`service-PROJECT_NUMBER@gcp-sa-aiplatform-re.iam.gserviceaccount.com`) as a database user - to log into the database, a database client - to connect to the database, and an AI Platform user - to connect to Vertex AI models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iH39YSf_iZle"
      },
      "outputs": [],
      "source": [
        "# Define service account\n",
        "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"value(projectNumber)\"\n",
        "SERVICE_ACCOUNT = f\"service-{PROJECT_NUMBER[0]}@gcp-sa-aiplatform-re.iam.gserviceaccount.com\"\n",
        "IAM_USER = SERVICE_ACCOUNT.replace(\".gserviceaccount.com\", \"\")\n",
        "\n",
        "# Force the creation of the AI Platform service accounts\n",
        "# The service accounts will be created at deploy time if not pre-created\n",
        "!gcloud beta services identity create --service=aiplatform.googleapis.com --project={PROJECT_ID}\n",
        "\n",
        "# Add a service account as database IAM user\n",
        "# For an IAM service account, supply the service account's address without the .gserviceaccount.com\n",
        "!gcloud alloydb users create {IAM_USER} \\\n",
        "  --region={REGION} \\\n",
        "  --cluster={CLUSTER} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --type=IAM_BASED\n",
        "\n",
        "# Grant IAM Permissions for database-user authentication\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "    --role=roles/alloydb.databaseUser\n",
        "\n",
        "# Grant IAM permissions to access AlloyDB instances\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "    --role=roles/alloydb.client\n",
        "\n",
        "# Grant IAM permissions to access AI Platform services\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT}  \\\n",
        "    --role=roles/aiplatform.user\n",
        "\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT}  \\\n",
        "    --role=roles/serviceusage.serviceUsageConsumer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ojqmIHx8IWdB"
      },
      "outputs": [],
      "source": [
        "# Grant access to vector store table to IAM users\n",
        "engine = await AlloyDBEngine.afrom_instance(\n",
        "    PROJECT_ID, REGION, CLUSTER, INSTANCE, DATABASE, user=\"postgres\", password=PASSWORD\n",
        ")\n",
        "\n",
        "async with engine._pool.connect() as conn:\n",
        "    await conn.execute(text(f'GRANT SELECT ON {TABLE_NAME} TO \"{IAM_USER}\";'))"
      ]
    },
    {
      "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.\n",
        "\n",
        "For improved security measures, the tool will use IAM-based authentication to authenticate to the databases instead of using the built-in user/password authentication."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MpUzFTt2K3Ky"
      },
      "outputs": [],
      "source": [
        "def similarity_search(query: str) -> list[Document]:\n",
        "    \"\"\"Searches and returns movies.\n",
        "\n",
        "    Args:\n",
        "      query: The user query to search for related items\n",
        "\n",
        "    Returns:\n",
        "      List[Document]: A list of Documents\n",
        "    \"\"\"\n",
        "    engine = AlloyDBEngine.from_instance(\n",
        "        PROJECT_ID,\n",
        "        REGION,\n",
        "        CLUSTER,\n",
        "        INSTANCE,\n",
        "        DATABASE,\n",
        "        # Uncomment to use built-in authentication instead of IAM authentication\n",
        "        # user=\"postgres\",\n",
        "        # password=PASSWORD,\n",
        "    )\n",
        "\n",
        "    vector_store = AlloyDBVectorStore.create_sync(\n",
        "        engine,\n",
        "        table_name=TABLE_NAME,\n",
        "        embedding_service=VertexAIEmbeddings(\n",
        "            model_name=\"textembedding-gecko@latest\", project=PROJECT_ID\n",
        "        ),\n",
        "    )\n",
        "    retriever = vector_store.as_retriever()\n",
        "    return retriever.invoke(query)"
      ]
    },
    {
      "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 Reasoning 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": [
        "vertexai.init(project=PROJECT_ID, location=\"us-central1\", staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "remote_app = reasoning_engines.ReasoningEngine.create(\n",
        "    reasoning_engines.LangchainAgent(\n",
        "        model=\"gemini-pro\",\n",
        "        tools=[similarity_search],\n",
        "        model_kwargs={\n",
        "            \"temperature\": 0.1,\n",
        "        },\n",
        "    ),\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[reasoningengine,langchain]==1.68.0\",\n",
        "        \"langchain-google-alloydb-pg==0.7.0\",\n",
        "        \"langchain-google-vertexai==1.0.4\",\n",
        "    ],\n",
        "    display_name=\"PrebuiltAgent\",\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": "code",
      "execution_count": null,
      "metadata": {
        "id": "45hiyNyfJ9Zo"
      },
      "outputs": [],
      "source": [
        "response = remote_app.query(input=\"Find movies about engineers\")\n",
        "print(response[\"output\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rLO17Uv9Xha-"
      },
      "outputs": [],
      "source": [
        "# Retrieve the application endpoint via the display name\n",
        "app_list = reasoning_engines.ReasoningEngine.list(filter='display_name=\"PrebuiltAgent\"')\n",
        "RESOURCE_ID = app_list[0].name\n",
        "\n",
        "# Retrieve the application endpoint via the resource ID\n",
        "remote_app = reasoning_engines.ReasoningEngine(\n",
        "    f\"projects/{PROJECT_ID}/locations/{LOCATION}/reasoningEngines/{RESOURCE_ID}\"\n",
        ")"
      ]
    },
    {
      "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 reasoning engine instance(s) and AlloyDB cluster and instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgNlHrxkb6c-"
      },
      "outputs": [],
      "source": [
        "# Delete the ReasoningEngine instance\n",
        "remote_app.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goyrqS2_I8Hs"
      },
      "outputs": [],
      "source": [
        "# Or delete all Reasoning Engine apps\n",
        "apps = reasoning_engines.ReasoningEngine.list()\n",
        "for app in apps:\n",
        "    app.delete()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "odvj8aKpb3Wi"
      },
      "outputs": [],
      "source": [
        "# Delete the AlloyDB cluster and instance\n",
        "!gcloud alloydb clusters delete {CLUSTER} \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --force \\\n",
        "  --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8NbUPwEfI62R"
      },
      "source": [
        "## What's next\n",
        "\n",
        "* Dive deeper into [LangChain on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview).\n",
        "* Learn more about the [AlloyDB for LangChain library](https://github.com/googleapis/langchain-google-alloydb-pg-python).\n",
        "* Explore other [Reasoning Engine samples](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/gemini/reasoning-engine)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_alloydb_rag_agent.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}