{
  "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 Cloud SQL for PostgreSQL to 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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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_cloud_sql_pg_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",
        "[Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview)\n",
        "is a managed service that helps you to build and deploy agent frameworks to a managed 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 Cloud SQL for PostgreSQL LangChain integration.\n",
        "\n",
        "Your [LangChain](https://python.langchain.com/docs/get_started/introduction) agent will use 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",
        "* Install and set up the Cloud SQL for PostgreSQL for LangChain and the Vertex AI Python SDKs\n",
        "* Create a Cloud SQL instance\n",
        "* Create a Cloud SQL database user\n",
        "* Define a retriever to perform similarity searches\n",
        "* Use the LangChain agent template provided in the Vertex AI SDK for Agent Engine\n",
        "* Deploy and test your agent on Agent Engine in Vertex AI\n",
        "\n",
        "⚠️ Note: Agent Engine currently only supports public IP for Cloud SQL for PostgreSQL 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-cloud-sql-pg>=0.10.0\" \"google-cloud-aiplatform[agent_engines,langchain]\" langchain-google-vertexai langchain-community"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f871fca8cbe5"
      },
      "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_cloud_sql_pg import PostgresEngine, PostgresVectorStore\n",
        "from langchain_google_vertexai import VertexAIEmbeddings\n",
        "from sqlalchemy import text  # noqa: F401\n",
        "import vertexai\n",
        "from vertexai import agent_engines\n",
        "from vertexai.preview.reasoning_engines import LangchainAgent"
      ]
    },
    {
      "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": "TqSfLzpdNG9J"
      },
      "source": [
        "## Create a Cloud Storage bucket\n",
        "\n",
        "Create or reuse and existing Cloud Storage bucket. Agent 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",
        "!gcloud storage buckets create --default-storage-class=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 sqladmin.googleapis.com servicenetworking.googleapis.com"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_yG0kddIvr7"
      },
      "source": [
        "## Set up Cloud SQL\n",
        "\n",
        "Use the provided variable names or update the values to use a pre-exisiting Cloud SQL instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtiB5-LVVkv0"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @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 a Cloud SQL instance\n",
        "\n",
        "This tutorial requires a Cloud SQL instance with public IP enabled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LdbQ2Q0wdNiy"
      },
      "outputs": [],
      "source": [
        "# Create Cloud SQL instance\n",
        "!gcloud sql instances create {INSTANCE} \\\n",
        "  --database-version=POSTGRES_15 \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --root-password={PASSWORD} \\\n",
        "  --cpu=1 \\\n",
        "  --memory=4GB \\\n",
        "  --assign-ip \\\n",
        "  --database-flags=cloudsql.iam_authentication=On"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZIjNIc71ixye"
      },
      "source": [
        "### Create a database\n",
        "\n",
        "Create a new database for the application using the Cloud SQL for LangChain library to establish a connection pool using the `PostgresEngine`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnNy1NFVk9F7"
      },
      "source": [
        "By default, [IAM database authentication](https://cloud.google.com/sql/docs/mysql/iam-logins) 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/sql/docs/mysql/built-in-authentication) using a username and password to access the Cloud SQL database can also be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pn4c8oZtixLI"
      },
      "outputs": [],
      "source": [
        "engine = await PostgresEngine.afrom_instance(\n",
        "    PROJECT_ID,\n",
        "    REGION,\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 `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": null,
      "metadata": {
        "id": "GGd89YWIi2qg"
      },
      "outputs": [],
      "source": [
        "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zX73buKxPXL8"
      },
      "outputs": [],
      "source": [
        "# Retrieve the CSV file\n",
        "!gcloud storage cp gs://github-repo/generative-ai/gemini/agent-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 PostgresVectorStore.create(\n",
        "    engine,\n",
        "    table_name=TABLE_NAME,\n",
        "    embedding_service=VertexAIEmbeddings(\n",
        "        model_name=\"text-embedding-005\", 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 Agent 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 sql users create {IAM_USER} \\\n",
        "  --instance={INSTANCE} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --type=cloud_iam_service_account\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/cloudsql.instanceUser\n",
        "\n",
        "# Grant IAM permissions to access Cloud SQL instances\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "    --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "    --role=roles/cloudsql.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 PostgresEngine.afrom_instance(\n",
        "    PROJECT_ID, REGION, 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 = 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=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 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 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": [
        "vertexai.init(project=PROJECT_ID, location=\"us-central1\", staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "remote_app = agent_engines.create(\n",
        "    LangchainAgent(\n",
        "        model=\"gemini-2.0-flash\",\n",
        "        tools=[similarity_search],\n",
        "        model_kwargs={\n",
        "            \"temperature\": 0.1,\n",
        "        },\n",
        "    ),\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[agent_engines,langchain]\",\n",
        "        \"langchain-google-cloud-sql-pg==0.10.0\",\n",
        "        \"langchain-google-vertexai==1.0.10\",\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 = agent_engines.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 = agent_engines.get(\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 Agent Engine instance(s) and Cloud SQL instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgNlHrxkb6c-"
      },
      "outputs": [],
      "source": [
        "# Delete the Agent Engine 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}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8NbUPwEfI62R"
      },
      "source": [
        "## What's next\n",
        "\n",
        "* Dive deeper into [Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview).\n",
        "* Learn more about the [Cloud SQL for LangChain library](https://github.com/googleapis/langchain-google-cloud-sql-pg-python).\n",
        "* Explore other [Agent Engine samples](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/gemini/agent-engine)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_cloud_sql_pg_rag_agent.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
