{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Run RAG Pipelines in BigQuery with BQML and Vector Search\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/use-cases/retrieval-augmented-generation/rag_with_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_with_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_with_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_with_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_with_bigquery.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_with_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_with_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_with_bigquery.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/use-cases/retrieval-augmented-generation/rag_with_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_with_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": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Jeff Nelson](https://github.com/jeffonelson/), Eric Hao |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates a basic end-to-end retrieval-augmented generation (RAG) pipeline using [BigQuery](https://cloud.google.com/bigquery/) and [BigQuery ML](https://cloud.google.com/bigquery/docs/bqml-introduction) functions. To do so, we:\n",
        "\n",
        "* Complete setup steps to download sample data and access [Vertex AI](https://cloud.google.com/vertex-ai) from BigQuery\n",
        "* Generate [object table](https://cloud.google.com/bigquery/docs/object-table-introduction) to access unstructured PDFs that reside in [Cloud Storage](https://cloud.google.com/storage)\n",
        "* Create a remote model, so BigQuery can call [Document AI](https://cloud.google.com/document-ai) to parse the PDF inputs\n",
        "* Parse response from Document AI into chunks and metadata, then generate vector embeddings for the chunks\n",
        "* Run a [vector search](https://cloud.google.com/bigquery/docs/vector-search) against embeddings in BigQuery, return relevant chunks, and summarize them with Gemini"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dc949afc1f08"
      },
      "source": [
        "## How to open this notebook in BigQuery Studio\n",
        "\n",
        "This notebook was written to be compatible for use within BigQuery Studio. To open this notebook in BigQuery, click to [Run in Colab Enterprise](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_with_bigquery.ipynb). This will open a new window in the Cloud Console and prompt you to confirm import. Then, navigate to BigQuery, where you will find the notebook available in the Explorer pane under Notebooks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ba5c12e483d"
      },
      "source": [
        "## About the dataset\n",
        "\n",
        "This example uses [the Federal Reserve's 2023 Survey of Consumer Finances](https://www.federalreserve.gov/publications/files/scf23.pdf) (SCF) report. The document contains information around US family income, net worth, credit use, and other common household financial indicators."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ce33dbc8fde"
      },
      "source": [
        "## Services and Costs\n",
        "\n",
        "This tutorial uses the following Google Cloud data analytics and ML services, they are billable components of Google Cloud:\n",
        "\n",
        "* BigQuery & BigQuery ML [(pricing)](https://cloud.google.com/bigquery/pricing)\n",
        "* Vertex AI Generative AI models [(pricing)](https://cloud.google.com/vertex-ai/generative-ai/pricing)\n",
        "* Document AI [(pricing)](https://cloud.google.com/document-ai/pricing)\n",
        "* Cloud Storage [(pricing)](https://cloud.google.com/storage/pricing)\n",
        "\n",
        "Use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "# Setup Steps to access Vertex AI models from BigQuery and enable APIs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff210a6d4d21"
      },
      "source": [
        "### Install Document AI SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2e9e2b9e1b1f"
      },
      "outputs": [],
      "source": [
        "%pip install --quiet google-cloud-documentai==2.31.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ed31279f009"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "567212ff53a6"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b96b39fd4d7b"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fa362c2ef5b5"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "9a07a9f9a4a9"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()\n",
        "    print(\"Authenticated\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Define your Google Cloud project"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"your-project-id\"  # @param {type: \"string\"}\n",
        "PROJECT_NUMBER = \"your-project-number\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04deeb11bbca"
      },
      "source": [
        "### Enable Data Table Display\n",
        "\n",
        "This makes it easier to visualize tabular data within a Notebook environment later on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "af9974f04f9f"
      },
      "outputs": [],
      "source": [
        "%load_ext google.colab.data_table"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4256d07d596"
      },
      "source": [
        "### Create a new dataset in BigQuery\n",
        "\n",
        "This will house any tables created throughout this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8a4c1a356d10"
      },
      "outputs": [],
      "source": [
        "!bq mk --location=us --dataset --project_id={PROJECT_ID} docai_demo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a100b689816b"
      },
      "source": [
        "### Create a Cloud resource connection\n",
        "\n",
        "[Cloud resource connections](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) enable BigQuery to access other Cloud services, like Cloud Storage and Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "885da43402b1"
      },
      "outputs": [],
      "source": [
        "!bq mk --connection --connection_type=CLOUD_RESOURCE --location=us --project_id={PROJECT_ID} \"demo_conn\"\n",
        "!bq show --location=us --connection --project_id={PROJECT_ID} \"demo_conn\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd9f6cbe4393"
      },
      "source": [
        "### Add permissions to Cloud resource connection service account\n",
        "\n",
        "The Cloud resource connection is associated with a service account. The following cell enables the service account to access services like Document AI, Cloud Storage, and Vertex AI.\n",
        "\n",
        "**Note:** Copy the service account ID from the prior cell and input it below. It will look like `your-copied-service-account@gcp-sa-bigquery-condel.iam.gserviceaccount.com`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "16b193a840cd"
      },
      "outputs": [],
      "source": [
        "connection_service_account = \"your-copied-service-account@gcp-sa-bigquery-condel.iam.gserviceaccount.com\"  # @param {type: \"string\"}\n",
        "connection_member = f\"serviceAccount:{connection_service_account}\"\n",
        "\n",
        "\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/documentai.viewer' --condition=None --quiet\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/storage.objectViewer' --condition=None --quiet\n",
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member={connection_member} --role='roles/aiplatform.user' --condition=None --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ba09d9393559"
      },
      "source": [
        "### Download the sample PDF used for this notebook and store it in a new Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4605453a6675"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "# Create a unique Cloud Storage bucket name\n",
        "bucket_name = f\"{PROJECT_ID}-{random.randint(10000, 99999)}\"\n",
        "\n",
        "# Create the bucket\n",
        "!gsutil mb -l US -p {PROJECT_ID} gs://{bucket_name}\n",
        "\n",
        "# Download the PDF sample\n",
        "!wget scf23.pdf \"https://www.federalreserve.gov/publications/files/scf23.pdf\"\n",
        "\n",
        "# Upload the PDF sample to the newly created Cloud Storage bucket\n",
        "!gsutil cp scf23.pdf gs://{bucket_name}/\n",
        "\n",
        "# Print confirmation\n",
        "print(f\"PDF uploaded to gs://{bucket_name}/scf23.pdf\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6b0a0bd5c4fd"
      },
      "source": [
        "## Create an object table\n",
        "\n",
        "An object table allows BigQuery to read unstructured data in Google Cloud Storage. This uses the BigQuery Python client library to continue using the `bucket_name` variable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "94cc075094c4"
      },
      "outputs": [],
      "source": [
        "from google.cloud import bigquery\n",
        "\n",
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "query = f\"\"\"\n",
        "CREATE OR REPLACE EXTERNAL TABLE `docai_demo.object_table`\n",
        "WITH CONNECTION `us.demo_conn`  -- Replace with your connection ID\n",
        "OPTIONS (\n",
        "  uris = ['gs://{bucket_name}/scf23.pdf'],\n",
        "  object_metadata = 'DIRECTORY'\n",
        ");\n",
        "\"\"\"\n",
        "\n",
        "query_job = client.query(query)  # API request\n",
        "query_job.result()  # Waits for the query to complete\n",
        "\n",
        "print(\"External table docai_demo.object_table created or replaced successfully.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c17ec8736188"
      },
      "source": [
        "### Show the object table\n",
        "\n",
        "Confirm that the results display the PDF document in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9f471aa348b2"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "SELECT * \n",
        "FROM `docai_demo.object_table`;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ec9d2c49fd34"
      },
      "source": [
        "## Use BQML and Document AI to parse documents"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afc0a5902ef3"
      },
      "source": [
        "### Create a Layout Parser Processor in Document AI\n",
        "\n",
        "[Create a new processor](https://cloud.google.com/document-ai/docs/create-processor#documentai_fetch_processor_types-python) in Document AI with the type `LAYOUT_PARSER_PROCESSOR`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "519ea8a55496"
      },
      "outputs": [],
      "source": [
        "from google.api_core.client_options import ClientOptions\n",
        "from google.cloud import documentai\n",
        "\n",
        "location = \"us\"\n",
        "processor_display_name = \"layout_parser_processor\"\n",
        "processor_type = \"LAYOUT_PARSER_PROCESSOR\"\n",
        "\n",
        "\n",
        "def create_processor_sample(\n",
        "    PROJECT_ID: str, location: str, processor_display_name: str, processor_type: str\n",
        ") -> None:\n",
        "    opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
        "\n",
        "    client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
        "\n",
        "    # The full resource name of the location\n",
        "    parent = client.common_location_path(PROJECT_ID, location)\n",
        "\n",
        "    # Create a processor\n",
        "    processor = client.create_processor(\n",
        "        parent=parent,\n",
        "        processor=documentai.Processor(\n",
        "            display_name=processor_display_name, type_=processor_type\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    # Return the processor ID needed for creating a BigQuery connection\n",
        "    return processor.name.split(\"/\")[-1]\n",
        "\n",
        "\n",
        "# Call this function to create the processor and return its ID\n",
        "processor_id = create_processor_sample(\n",
        "    PROJECT_ID, location, processor_display_name, processor_type\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cfcaa4c4584"
      },
      "source": [
        "### Create a remote model in BigQuery that connects with your Document AI Layout Parser Processor\n",
        "\n",
        "This one-time setup step allows BigQuery to reference the Document AI Processor you just created.\n",
        "\n",
        "**Note:** If if you receive an 400 GET error \"permission denied for document processor\", you may need to wait a minute for permissions to propagate from earlier steps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "57233716c232"
      },
      "outputs": [],
      "source": [
        "query = f\"\"\"\n",
        "CREATE OR REPLACE MODEL `docai_demo.layout_parser` \n",
        "REMOTE WITH CONNECTION `us.demo_conn`\n",
        "OPTIONS(remote_service_type=\"CLOUD_AI_DOCUMENT_V1\", document_processor=\"{processor_id}\")\n",
        "\"\"\"\n",
        "\n",
        "query_job = client.query(query)  # API request\n",
        "query_job.result()  # Waits for the query to complete\n",
        "\n",
        "print(\"Remote model docai_demo.layout_parser created or replaced successfully.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fc821b6f845d"
      },
      "source": [
        "### Process the document using BigQuery ML\n",
        "\n",
        "Use the [`ML.PROCESS_DOCUMENT` function](https://cloud.google.com/bigquery/docs/process-document) from BigQuery to call your Document AI processor and pass through the PDF. This uses the Layout Parser configuration and chunks your document.\n",
        "\n",
        "**Note:** this may take a minute or so to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a489b3cb9e1d"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID --location us\n",
        "\n",
        "CREATE or REPLACE TABLE docai_demo.demo_result AS (\n",
        "  SELECT * FROM ML.PROCESS_DOCUMENT(\n",
        "  MODEL docai_demo.layout_parser,\n",
        "  TABLE docai_demo.object_table,\n",
        "  PROCESS_OPTIONS => (JSON '{\"layout_config\": {\"chunking_config\": {\"chunk_size\": 250}}}')\n",
        "  )\n",
        ");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfb0a1fa3266"
      },
      "source": [
        "### Parse the JSON results returned to BigQuery\n",
        "\n",
        "The `ML.PROCESS_DOCUMENT` function parses the PDF from Cloud Storage and returns a JSON blob to BigQuery. In this step, we'll parse the JSON, extract document chunks and metadata, and return it to a new BigQuery table."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2bc4dad2e399"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID --location us\n",
        "\n",
        "CREATE OR REPLACE TABLE docai_demo.demo_result_parsed AS (\n",
        "SELECT\n",
        "  uri,\n",
        "  JSON_EXTRACT_SCALAR(json , '$.chunkId') AS id,\n",
        "  JSON_EXTRACT_SCALAR(json , '$.content') AS content,\n",
        "  JSON_EXTRACT_SCALAR(json , '$.pageFooters[0].text') AS page_footers_text,\n",
        "  JSON_EXTRACT_SCALAR(json , '$.pageSpan.pageStart') AS page_span_start,\n",
        "  JSON_EXTRACT_SCALAR(json , '$.pageSpan.pageEnd') AS page_span_end\n",
        "FROM docai_demo.demo_result, UNNEST(JSON_EXTRACT_ARRAY(ml_process_document_result.chunkedDocument.chunks, '$')) json\n",
        ");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "adca53cc55d8"
      },
      "source": [
        "### Display the parsed document chunks\n",
        "\n",
        "Show a preview of the parsed results and metadata."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c60bcdc388c4"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID --location us\n",
        "\n",
        "SELECT *\n",
        "FROM docai_demo.demo_result_parsed\n",
        "ORDER BY id\n",
        "LIMIT 5;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a980e66443bc"
      },
      "source": [
        "## Connect to Vertex AI embedding generation and Gemini access"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eae56fa8c74c"
      },
      "source": [
        "### Connect to a text embedding model\n",
        "\n",
        "[Create a remote model](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-remote-model) allowing BigQuery access to a text embedding model hosted in Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3c53a24e59a1"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "CREATE OR REPLACE MODEL `docai_demo.embedding_model` \n",
        "REMOTE WITH CONNECTION `us.demo_conn` OPTIONS(endpoint=\"text-embedding-005\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45d1ccc016c8"
      },
      "source": [
        "### Generate embeddings\n",
        "\n",
        "Use the [`ML.GENERATE_EMBEDDING` function](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-generate-embedding) in BigQuery to generate embeddings for all text chunks in the document."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63bf77f48b8c"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "CREATE OR REPLACE TABLE `docai_demo.embeddings` AS\n",
        "SELECT * FROM ML.GENERATE_EMBEDDING(\n",
        "  MODEL `docai_demo.embedding_model`,\n",
        "  TABLE `docai_demo.demo_result_parsed`\n",
        ");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1ce3b78e01a"
      },
      "source": [
        "### Connect to a Gemini LLM endpoint\n",
        "\n",
        "Create a remote model allowing BigQuery access to a Gemini foundation model hosted in Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b760c54502e"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "CREATE OR REPLACE MODEL `docai_demo.gemini_flash` REMOTE\n",
        "WITH CONNECTION `us.demo_conn` OPTIONS(endpoint=\"gemini-2.0-flash\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "879593c348c4"
      },
      "source": [
        "## Run vector search, return results, and pass them to Gemini for text generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2eb640b369a3"
      },
      "source": [
        "### Sample BigQuery vector search\n",
        "\n",
        "Run a sample BigQuery vector search against your chunks. This query takes your text input, creates an embedding using the `ML.GENERATE_EMBEDDING` function, and then passes the embedding through to the [`VECTOR_SEARCH` function](https://cloud.google.com/bigquery/docs/reference/standard-sql/search_functions#vector_search). The results are the top ten chunks that are most semantically related to your input.\n",
        "\n",
        "In the search query below, the input text asks \"Did the typical family net worth increase? If so, by how much?\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf9fa689905d"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "SELECT query.query, base.uri, base.id, base.content, distance\n",
        "    FROM\n",
        "      VECTOR_SEARCH( TABLE `docai_demo.embeddings`,\n",
        "        'ml_generate_embedding_result',\n",
        "        (\n",
        "        SELECT\n",
        "          ml_generate_embedding_result,\n",
        "          content AS query\n",
        "        FROM\n",
        "          ML.GENERATE_EMBEDDING( MODEL `docai_demo.embedding_model`,\n",
        "            ( SELECT 'Did the typical family net worth increase? If so, by how much?' AS content)\n",
        "          ) \n",
        "        ),\n",
        "        top_k => 10,\n",
        "        OPTIONS => '{\"fraction_lists_to_search\": 0.01}') \n",
        "ORDER BY distance DESC;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "623765bd9154"
      },
      "source": [
        "## Generate text augmented by vector search results\n",
        "\n",
        "This step builds upon the prior one - but instead of simply returning the top text chunks, it calls the `ML.GENERATE_TEXT` function to summarize them alongside the question we input.\n",
        "\n",
        "In this query you:\n",
        "* **Retrieve** the closest chunks semantically using the `VECTOR_SEARCH` function (this is what was done in the prior query)\n",
        "* **Augment** the Gemini LLM with this knowledge\n",
        "* **Generate** a succinct answer using the `ML.GENERATE_TEXT` function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2f6f83f2eca7"
      },
      "outputs": [],
      "source": [
        "%%bigquery --project $PROJECT_ID\n",
        "\n",
        "SELECT\n",
        "  ml_generate_text_llm_result AS generated,\n",
        "  -- prompt -- Commented out, but please feel free to uncomment if you would like to see the full context passed to the Gemini model\n",
        "FROM\n",
        "  ML.GENERATE_TEXT( MODEL `docai_demo.gemini_flash`,\n",
        "    (\n",
        "    SELECT\n",
        "    CONCAT( 'Did the typical family net worth change? How does this compare the SCF survey a decade earlier? Be concise and use the following context:',\n",
        "    STRING_AGG(FORMAT(\"context: %s and reference: %s\", base.content, base.uri), ',\\n')) AS prompt,\n",
        "    FROM\n",
        "      VECTOR_SEARCH( TABLE \n",
        "        `docai_demo.embeddings`,\n",
        "        'ml_generate_embedding_result',\n",
        "        (\n",
        "        SELECT\n",
        "          ml_generate_embedding_result,\n",
        "          content AS query\n",
        "        FROM\n",
        "          ML.GENERATE_EMBEDDING( MODEL `docai_demo.embedding_model`,\n",
        "            (\n",
        "            SELECT\n",
        "              'Did the typical family net worth change? How does this compare the SCF survey a decade earlier?' AS content\n",
        "            )\n",
        "          ) \n",
        "        ),\n",
        "        top_k => 10,\n",
        "        OPTIONS => '{\"fraction_lists_to_search\": 0.01}') \n",
        "      ),\n",
        "      STRUCT(512 AS max_output_tokens, TRUE AS flatten_json_output)\n",
        "  );"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "945bbb4357b2"
      },
      "source": [
        "### Sample questions to try out:\n",
        "\n",
        "Here are a list of a few other questions to spark your imagination. Feel free to try your own too!\n",
        "* Did the amount of debt families own on their home increase between 2019 and 2022?\n",
        "* Did younger or older families see their net worth increase more?\n",
        "* How much did the median family income change between 2018 and 2021?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial by uncommenting the below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ab59128be6c"
      },
      "outputs": [],
      "source": [
        "# # Deletes the BigQuery assets and Google Cloud Storage bucket\n",
        "\n",
        "# !bq rm -r -f $PROJECT_ID:docai_demo\n",
        "# !bq rm --connection --project_id=$PROJECT_ID --location=us demo_conn\n",
        "# !gsutil rm -r gs://{bucket_name}\n",
        "\n",
        "\n",
        "# # Deletes the Document AI processor\n",
        "# def delete_processor_sample(\n",
        "#     PROJECT_ID: str, location: str, processor_id: str\n",
        "# ) -> None:\n",
        "#     \"\"\"Deletes a processor.\"\"\"\n",
        "\n",
        "#     opts = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
        "\n",
        "#     client = documentai.DocumentProcessorServiceClient(client_options=opts)\n",
        "\n",
        "#     # The full resource name of the processor\n",
        "#     name = f\"projects/{PROJECT_ID}/locations/{location}/processors/{processor_id}\"\n",
        "\n",
        "#     try:\n",
        "#         client.delete_processor(name=name)\n",
        "#         print(f\"Processor {processor_id} deleted successfully.\")\n",
        "#     except Exception as e:\n",
        "#         print(f\"Error deleting processor: {e}\")\n",
        "\n",
        "\n",
        "# # Call this function to delete the processor\n",
        "# delete_processor_sample(PROJECT_ID, location, processor_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aefa89207b70"
      },
      "source": [
        "# Wrap up\n",
        "\n",
        "This notebook demonstrates an example of how to achieve a basic end-to-end retrieval-augmented generation pipeline using BigQuery. It integrates BigQuery ML functions like `ML.PROCESS_DOCUMENT` to call Document AI and parse PDFs, `ML.GENERATE_EMBEDDING` to generate embeddings on text chunks and input queries, and `ML.GENERATE_TEXT` to provide a concise answer. It also uses the `VECTOR_SEARCH` function to identify similar text (using embeddings) in BigQuery using familiar SQL syntax.\n",
        "\n",
        "To continue learn more, check out our documentation on [BigQuery ML](https://cloud.google.com/bigquery/docs/bqml-introduction) and [BigQuery Vector Search](https://cloud.google.com/bigquery/docs/vector-search)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "rag_with_bigquery.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
