gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0dFfn6IQRjwL"
      },
      "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": "6VKtsO2cPx0c"
      },
      "source": [
        "# Augment Gemini Output with Vector Embeddings from BigQuery\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68HXNikaOXh8"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fretrieval-augmented-generation%2Frag_vector_embedding_in_bigquery.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.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/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.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/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.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/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.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/use-cases/retrieval-augmented-generation/rag_vector_embedding_in_bigquery.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": "0ExfeGcm9Iu5"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Logan Ramalingam](https://github.com/logan-google) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ja124yztPhGw"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook shows how to use BigQuery to create generate embeddings from text in a BigQuery table, store them within BigQuery, and then use the embeddings to augment the results from LLM in Vector Search.\n",
        "\n",
        "In this notebook, we create text embeddings for publicly available abstracts from [patents data](https://console.cloud.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) and use them in our LLM search. Google Patents Public Data, provided by IFI CLAIMS Patent Services, is a worldwide bibliographic and US full-text dataset of patent publications.\n",
        "\n",
        "\n",
        "```patents-public-data.google_patents_research.publications```\n",
        "\n",
        "This notebook references the steps mentioned in [Perform semantic search and retrieval-augmented generation](https://cloud.google.com/bigquery/docs/vector-index-text-search-tutorial)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q1U3CtPWTGTO"
      },
      "source": [
        "## Required roles and permissions\n",
        "\n",
        "To create a connection, you need membership in the following Identity and Access Management (IAM) role:\n",
        "\n",
        "* ```roles/bigquery.connectionAdmin```\n",
        "\n",
        "To grant permissions to the connection's service account, you need the following permission:\n",
        "\n",
        "* ```resourcemanager.projects.setIamPolicy```\n",
        "\n",
        "The IAM permissions needed in this tutorial for the remaining BigQuery operations are included in the following two roles:\n",
        "\n",
        "\n",
        "*   BigQuery Data Editor (```roles/bigquery.dataEditor```) to create models, tables, and indexes.\n",
        "\n",
        "*   BigQuery User (```roles/bigquery.user```) to run BigQuery jobs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9UyVG5euTqOc"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k9DON_o2Tvvx"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oqy-9gVcO8KW"
      },
      "outputs": [],
      "source": [
        "pip install --upgrade --user --quiet google-cloud-aiplatform google-cloud-bigquery lxml google-cloud-bigquery-connection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8y3BBD2mT-Yj"
      },
      "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 its restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZHuV1OBCUC7e"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "obVlQPPTT_-G"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gtpBwMszUFuj"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4FnXQw5HUFd_"
      },
      "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": "BDAU3-1SUMwM"
      },
      "source": [
        "## Set Google Cloud project information and initialize BigQuery Connect\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": "CBTVwIiHO_NB"
      },
      "outputs": [],
      "source": [
        "# Google Cloud Project ID\n",
        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
        "\n",
        "# BigQuery Dataset for storing embeddings and model\n",
        "DATASET_ID = \"bq_vector_embeddings\"  # @param {type:\"string\"}\n",
        "\n",
        "# BigQuery Region\n",
        "REGION = \"US\"  # @param {type: \"string\"}\n",
        "\n",
        "# BigQuery Connection name\n",
        "CONN_NAME = \"bqml_llm_conn\"\n",
        "\n",
        "# Embeddings Remote Model name in BigQuery\n",
        "EMBEDDINGS_MODEL_ID = \"llm_gecko\"  # @param {type:\"string\"}\n",
        "\n",
        "# Embeddings Table name in BigQuery\n",
        "EMBEDDINGS_TABLE_ID = \"embeddings\"  # @param {type:\"string\"}\n",
        "\n",
        "# LLM Remote Model name in BigQuery\n",
        "LLM_MODEL_ID = \"llm_gemini\"  # @param {type:\"string\"}\n",
        "\n",
        "# Embeddings Model to use\n",
        "EMBEDDINGS_ENDPOINT_TYPE = \"text-embedding-004\"  # @param {type:\"string\"}\n",
        "\n",
        "# LLM Model to use\n",
        "LLM_ENDPOINT_TYPE = \"gemini-1.5-pro\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2SUO-szePFC_"
      },
      "outputs": [],
      "source": [
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X3YexMilEk7A"
      },
      "source": [
        "## Import libraries\n",
        "\n",
        "Let's start by importing the libraries that we will need for this tutorial"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "khvtJTNwPLQB"
      },
      "outputs": [],
      "source": [
        "# if in Colab, enable data_table format\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import data_table\n",
        "\n",
        "    data_table.enable_dataframe_formatter()\n",
        "\n",
        "from google.cloud import bigquery\n",
        "from google.cloud import bigquery_connection_v1 as bq_connection\n",
        "from google.cloud.exceptions import NotFound"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J4JHJsIgWAYY"
      },
      "source": [
        "## Setup BigQuery Environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sd-CU5BtvzYI"
      },
      "source": [
        "### Initialize Google BigQuery Client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aXveDrBtvyr-"
      },
      "outputs": [],
      "source": [
        "client = bigquery.Client(project=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AFSqGN9JWK25"
      },
      "source": [
        "### Wrapper to use BigQuery client to run query and return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6hYtolHcPOgz"
      },
      "outputs": [],
      "source": [
        "def run_bq_query(sql: str):\n",
        "    \"\"\"\n",
        "    Input: SQL query, as a string, to execute in BigQuery\n",
        "    Returns the query results or error, if any\n",
        "    \"\"\"\n",
        "    try:\n",
        "        query_job = client.query(sql)\n",
        "        result = query_job.result()\n",
        "        print(f\"JOB ID: {query_job.job_id} STATUS: {query_job.state}\")\n",
        "        return result\n",
        "\n",
        "    except Exception as e:\n",
        "        raise Exception(str(e))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jem8DhWWMx0e"
      },
      "source": [
        "### Create BigQuery Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ew-5kjpJMw00"
      },
      "outputs": [],
      "source": [
        "# Set dataset_id to the ID of the dataset to create.\n",
        "dataset = f\"{PROJECT_ID}.{DATASET_ID}\"\n",
        "\n",
        "# Construct a full Dataset object to send to the API.\n",
        "dataset_object = bigquery.Dataset(dataset)\n",
        "\n",
        "# Specify the geographic location where the dataset should reside.\n",
        "dataset_object.location = \"US\"\n",
        "\n",
        "# Send the dataset to the API for creation, with an explicit timeout.\n",
        "# Raises google.api_core.exceptions.Conflict if the Dataset already\n",
        "# exists within the project.\n",
        "try:\n",
        "    client.get_dataset(dataset_object)  # Make an API request.\n",
        "    print(f\"Dataset {dataset} already exists\")\n",
        "except NotFound:\n",
        "    dataset = client.create_dataset(dataset_object, timeout=30)  # Make an API request.\n",
        "    print(f\"Created dataset {client.project}.{dataset_object.dataset}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0qek5a1W2Pl"
      },
      "source": [
        "### Create BigQuery Cloud resource connection\n",
        "\n",
        "You will need to create a [Cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) to enable BigQuery to interact with Vertex AI services:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AYZb179IPQG6"
      },
      "outputs": [],
      "source": [
        "conn_client = bq_connection.ConnectionServiceClient()\n",
        "new_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n",
        "exists_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}\"\n",
        "cloud_resource_properties = bq_connection.CloudResourceProperties({})\n",
        "\n",
        "# Try to use an existing connection if one already exists. If not, create a new one.\n",
        "try:\n",
        "    request = conn_client.get_connection(\n",
        "        request=bq_connection.GetConnectionRequest(name=exists_conn_parent)\n",
        "    )\n",
        "    conn_service_account = f\"serviceAccount:{request.cloud_resource.service_account_id}\"\n",
        "except Exception:\n",
        "    connection = bq_connection.types.Connection(\n",
        "        {\"friendly_name\": CONN_NAME, \"cloud_resource\": cloud_resource_properties}\n",
        "    )\n",
        "    request = bq_connection.CreateConnectionRequest(\n",
        "        {\n",
        "            \"parent\": new_conn_parent,\n",
        "            \"connection_id\": CONN_NAME,\n",
        "            \"connection\": connection,\n",
        "        }\n",
        "    )\n",
        "    response = conn_client.create_connection(request)\n",
        "    conn_service_account = (\n",
        "        f\"serviceAccount:{response.cloud_resource.service_account_id}\"\n",
        "    )\n",
        "print(conn_service_account)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jieSbg0jXAUK"
      },
      "source": [
        "### Set permissions for Service Account\n",
        "The resource connection service account requires certain project-level permissions which are outlined in the <a href=\"https://cloud.google.com/bigquery/docs/bigquery-ml-remote-model-tutorial#set_up_access\" target=\"_blank\">Vertex AI function documentation</a>.\n",
        "\n",
        "<br>\n",
        "\n",
        "**Note:** If you are using Vertex AI Workbench, the service account used by Vertex AI may not have sufficient permissions to add IAM policy bindings.\n",
        "\n",
        "The [IAM Grant Access](https://cloud.google.com/iam/docs/granting-changing-revoking-access#grant-single-role) page gives instructions on how these policy bindings can be added using Cloud Shell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q5aNt3ZAPSJI"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={conn_service_account} --role='roles/serviceusage.serviceUsageConsumer'\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={conn_service_account} --role='roles/bigquery.connectionUser'\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={conn_service_account} --role='roles/aiplatform.user'\n",
        "!gcloud services enable bigqueryconnection.googleapis.com\n",
        "# wait 60 seconds, give IAM updates time to propagate, otherwise, following cells will fail\n",
        "time.sleep(60)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LUe5waT9XS0l"
      },
      "source": [
        "# Configure Vertex AI Embeddings Model in BigQuery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__1U_CJkYTNS"
      },
      "source": [
        "## Create the remote model for text embeddings generation\n",
        "Create a remote model that represents a hosted Vertex AI text embeddings generation model.\n",
        "\n",
        "The query takes several seconds to complete, after which the model ```EMBEDDINGS_MODEL_ID``` appears in the ```DATASET_ID``` in the Explorer pane."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rOuKE6z9GPIJ"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"CREATE OR REPLACE MODEL\n",
        "            `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`\n",
        "          REMOTE WITH CONNECTION\n",
        "            `{PROJECT_ID}.{REGION}.{CONN_NAME}`\n",
        "          OPTIONS (ENDPOINT = '{EMBEDDINGS_ENDPOINT_TYPE}');\"\"\"\n",
        "result = run_bq_query(sql)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Baocs6-vYzBf"
      },
      "source": [
        "## Generate text embeddings\n",
        "Generate text embeddings from patent abstracts using the ```ML.GENERATE_TEXT_EMBEDDING``` function, and then write them to a BigQuery table so that they can be searched.\n",
        "\n",
        "**Note: Query might take up to 10 minutes to run.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i6_BDJn9GRAi"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "      CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}` AS\n",
        "      SELECT * FROM ML.GENERATE_TEXT_EMBEDDING(\n",
        "        MODEL `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`,\n",
        "        (\n",
        "          SELECT *, abstract AS content\n",
        "          FROM `patents-public-data.google_patents_research.publications`\n",
        "          WHERE LENGTH(abstract) > 0 AND LENGTH(title) > 0 AND country = 'Singapore'\n",
        "        )\n",
        "      )\n",
        "      WHERE ARRAY_LENGTH(text_embedding) > 0;\n",
        "      \"\"\"\n",
        "result = run_bq_query(sql)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "725vcaYAY_KA"
      },
      "source": [
        "## Create Vector index\n",
        "\n",
        "A vector index is a data structure designed to let the ```VECTOR_SEARCH``` function perform a more efficient vector search of embeddings. When ```VECTOR_SEARCH``` is able to use a vector index, the function uses the Approximate Nearest Neighbor search technique to help improve search performance, with the trade-off of reducing recall and thus returning more approximate results.\n",
        "\n",
        "**NOTE: Query might take up to 5 minutes to run.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u_9ndU0BPbPv"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"CREATE OR REPLACE VECTOR INDEX my_index ON `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`(text_embedding) OPTIONS(index_type = 'IVF',\n",
        "distance_type = 'COSINE',   ivf_options = '{{\"num_lists\":500}}')\"\"\"\n",
        "result = run_bq_query(sql)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B_JRpw5uZb79"
      },
      "source": [
        "### Verify vector index creation\n",
        "\n",
        "The vector index is populated asynchronously. You can check whether the index is ready to be used by querying the ```INFORMATION_SCHEMA.VECTOR_INDEXES``` view and verifying that the coverage_percentage column value is greater than 0 and the ```last_refresh_time``` column value isn't ```NULL```."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PWe97Sk5PdSe"
      },
      "outputs": [],
      "source": [
        "# Check vector index creation status, 'coverage_percentage' should be 100\n",
        "sql = f\"\"\"\n",
        "\n",
        "    SELECT table_name, index_name, index_status, coverage_percentage, last_refresh_time, disable_reason\n",
        "    FROM `{PROJECT_ID}.{DATASET_ID}.INFORMATION_SCHEMA.VECTOR_INDEXES`\n",
        "    WHERE table_name = '{EMBEDDINGS_TABLE_ID}'\n",
        "    \"\"\"\n",
        "\n",
        "result = run_bq_query(sql).to_dataframe()\n",
        "print(result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uAcxPdClaRMQ"
      },
      "source": [
        "## Perform a text similarity search using the vector index\n",
        "\n",
        "Use the ```VECTOR_SEARCH``` function to search for the top 5 relevant patents that match embeddings generated from a text query. The model you use to generate the embeddings in this query must be the same as the one you use to generate the embeddings in the table you are comparing against, otherwise the search results won't be accurate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6-5r-0q7Pern"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "  SELECT\n",
        "    query.query,\n",
        "    base.content,\n",
        "    distance\n",
        "  FROM\n",
        "    VECTOR_SEARCH( TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`,\n",
        "      'text_embedding',\n",
        "      (\n",
        "      SELECT\n",
        "        text_embedding,\n",
        "        content AS query\n",
        "      FROM\n",
        "        ML.GENERATE_TEXT_EMBEDDING( MODEL `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`,\n",
        "          (\n",
        "          SELECT 'improving password security' AS content))\n",
        "      ),\n",
        "      top_k => 5,\n",
        "      OPTIONS => '{{\"fraction_lists_to_search\":0.01}}');\"\"\"\n",
        "\n",
        "result = run_bq_query(sql).to_dataframe()\n",
        "print(result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CQCg_qDzbGV3"
      },
      "source": [
        "# Generate text using embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0A4OZzdRbaDK"
      },
      "source": [
        "## Create the remote model for text generation\n",
        "\n",
        "Create a remote model that represents a hosted Gemini Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fWrK3FH9H99l"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\n",
        "      CREATE OR REPLACE MODEL\n",
        "        `{PROJECT_ID}.{DATASET_ID}.{LLM_MODEL_ID}`\n",
        "        REMOTE WITH CONNECTION\n",
        "          `{PROJECT_ID}.{REGION}.{CONN_NAME}`\n",
        "        OPTIONS (ENDPOINT = '{LLM_ENDPOINT_TYPE}');\n",
        "      \"\"\"\n",
        "result = run_bq_query(sql)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dziHo0LQbrHv"
      },
      "source": [
        "## Generate text augmented by vector search results\n",
        "\n",
        "Feed the search results as prompts to generate text with the ```ML.GENERATE_TEXT``` function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uqH_qvy_cKot"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"SELECT ml_generate_text_llm_result AS generated, prompt\n",
        "FROM ML.GENERATE_TEXT(\n",
        "  MODEL `{PROJECT_ID}.{DATASET_ID}.{LLM_MODEL_ID}`,\n",
        "  (\n",
        "    SELECT CONCAT(\n",
        "      'Propose some project ideas to improve user password security using the context below. Add the patent title and url to each idea: ',\n",
        "      STRING_AGG(\n",
        "        FORMAT(\"patent title: %s, patent abstract: %s\", base.title, base.abstract))\n",
        "      ) AS prompt,\n",
        "    FROM VECTOR_SEARCH(\n",
        "      TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`, 'text_embedding',\n",
        "      (\n",
        "        SELECT text_embedding, content AS query\n",
        "        FROM ML.GENERATE_TEXT_EMBEDDING(\n",
        "          MODEL `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`,\n",
        "         (SELECT 'improving password security' AS content)\n",
        "        )\n",
        "      ),\n",
        "    top_k => 5, options => '{{\"fraction_lists_to_search\": 0.01}}')\n",
        "  ),\n",
        "  STRUCT(600 AS max_output_tokens, TRUE AS flatten_json_output));\"\"\"\n",
        "\n",
        "query_job = client.query(sql)\n",
        "rows = query_job.result()\n",
        "\n",
        "for row in rows:\n",
        "    print(row[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qyVndOMsw7CU"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Clean up resources created in this notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UZOryN4Nw6vW"
      },
      "outputs": [],
      "source": [
        "# Delete Vector Index\n",
        "sql = f\"\"\"DROP VECTOR INDEX my_index ON `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`\"\"\"\n",
        "result = run_bq_query(sql)\n",
        "\n",
        "# Delete Gemini Model\n",
        "sql = f\"\"\"DROP MODEL `{PROJECT_ID}.{DATASET_ID}.{LLM_MODEL_ID}`\"\"\"\n",
        "result = run_bq_query(sql)\n",
        "\n",
        "# Delete Embeddings Model\n",
        "sql = f\"\"\"DROP MODEL `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_MODEL_ID}`\"\"\"\n",
        "result = run_bq_query(sql)\n",
        "\n",
        "# Delete Embeddings Table\n",
        "sql = f\"\"\"DROP TABLE `{PROJECT_ID}.{DATASET_ID}.{EMBEDDINGS_TABLE_ID}`\"\"\"\n",
        "result = run_bq_query(sql)\n",
        "\n",
        "# Delete BigQuery Connection\n",
        "request = bq_connection.DeleteConnectionRequest({\"name\": exists_conn_parent})\n",
        "response = conn_client.delete_connection(request)\n",
        "\n",
        "# Delete Dataset\n",
        "client.delete_dataset(dataset_object, delete_contents=True, not_found_ok=True)\n",
        "\n",
        "# Close BigQuery Connection\n",
        "client.close()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "rag_vector_embedding_in_bigquery.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}