language/use-cases/sql-code-generation/sql_code_generation.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "2713adfe27b4"
      },
      "source": [
        "# SQL Code Generation on Vertex AI\n",
        "\n",
        "**NOTE:** This notebook uses the PaLM generative model, which will reach its [discontinuation date in October 2024](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text#model_versions). Please refer to [this updated notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/c86e9da59015a269894bc5ccf91ff08f33cdee44/language/code/code_generation.ipynb) for a version which uses the latest Gemini model.\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/language/use-cases/sql-code-generation/sql_code_generation.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%2Flanguage%2Fuse-cases%2Fsql-code-generation%2Fsql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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/language/use-cases/sql-code-generation/sql_code_generation.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": "70a9ba8b7182"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s): | [Roy Arsan](https://www.linkedin.com/in/arsan) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3PN8SlH5Kyp"
      },
      "source": [
        "## Overview\n",
        "Large language models can be used for generating code, including SQL. In particular, models can convert natural language text into SQL queries. One common purpose is to enable users to query data without requiring knowledge of tables' names, data schema nor the specific SQL dialect or query engine of the underlying data warehouse like BigQuery.\n",
        "\n",
        "This notebook covers prompt engineering best practices for SQL code generation, and puts in practice learnings from [SQL-PaLM: Improve Large Language Model Adaptation for text-to-SQL](https://arxiv.org/pdf/2306.00739.pdf). For example, the BigQuery dataset schema is retrieved and provided dynamically as context to the prompt, for grounding the LLM and personalizing its output. The notebook also demonstrates simple model evaluation whereby the generated SQL queries are evaluated by executing them against the BigQuery dataset, and by comparing them with ground truth queries and corresponding results.\n",
        "\n",
        "For this notebook, you generate SQL queries to analyze Cloud Audit Logs and answer critical security questions around activity in your own Google Cloud project. While this notebook uses BigQuery logs dataset, the concepts and approach presented here can be applied to other databases and datasets.\n",
        "\n",
        "![NL2SQL flow](https://services.google.com/fh/files/misc/nl2sql_for_log_analytics2.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "czFRNoBSoBjW"
      },
      "source": [
        "### Objective\n",
        "\n",
        "By the end of the notebook, you should be able to:\n",
        "\n",
        "* Use model to generate SQL queries based on Natural Language questions:\n",
        "  * Using few-shot prompting\n",
        "  * Providing custom dataset schemas as context\n",
        "  * Formatting model output\n",
        "\n",
        "* Evaluate model-generated queries by:\n",
        "  * Executing sanitized queries against live dataset\n",
        "  * Comparing queries (and their results) to ground truth queries using simple fuzzy string matching\n",
        "  * Calculating model accuracy score\n",
        "\n",
        "In addition, you can use this notebook to answer your own security questions from your own audit logs, such as:\n",
        "\n",
        "- Any unusually high cloud API usage by any user identity over the last month?\n",
        "- Any destructive actions by an unapproved identity over the last 7 days?\n",
        "- Any unusual day-to-day spike in data volume accessed by any user this week?\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8b1nifEeds-m"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "499edc0841c5"
      },
      "source": [
        "### Prerequisite\n",
        " If you haven't already done so, the only requirement is to [upgrade your existing log bucket](https://cloud.google.com/logging/docs/buckets#upgrade-bucket) to use Log Analytics which provides you with a linked BigQuery dataset with your own queryable logs data. This is a **one-click step without incurring additional costs**. By default, Cloud Audit Admin Activity logs are enabled, ingested and stored in every project's `_Required` bucket without any charges.\n",
        "\n",
        "![one click prerequisite](https://services.google.com/fh/files/misc/upgrade_log_bucket.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENtD4TZZgoOM"
      },
      "source": [
        "### Install SDKs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sZd6ei_YgEnn"
      },
      "outputs": [],
      "source": [
        "# Install Vertex AI SDK to use for model predictions\n",
        "%pip install google-cloud-aiplatform google-cloud-bigquery --upgrade --user\n",
        "\n",
        "# Install fuzzy string comparison modules for model output evaluation\n",
        "%pip install -q python-Levenshtein --upgrade --user\n",
        "%pip install -q fuzzywuzzy --upgrade --user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e4931c081d52"
      },
      "source": [
        "Install Python SDK for Google Sheets only if you wish to later save the model-generated SQL queries and their results into a Google Sheet for subsequent troubleshooting or to expand your few-shot examples dataset. This is **not applicable for Vertex AI Workbench or Colab Enterprise** because notebooks in those environments cannot access Google Drive or Google Sheets for security purposes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c5c67900517c"
      },
      "outputs": [],
      "source": [
        "# [Optional] Install Python SDK for Google Sheets\n",
        "%pip install gspread --upgrade --user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1c0c19e21e0a"
      },
      "source": [
        "**Colab only:** Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "87988e3ac23d"
      },
      "outputs": [],
      "source": [
        "# # Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KbA0Abs_4YMg"
      },
      "source": [
        "### Authenticating your notebook environment\n",
        "* If you are using **Colab** to run this notebook, run the cell below and continue.\n",
        "* If you are using **Vertex AI Workbench**, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oc466AEHwGWG"
      },
      "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": "code",
      "execution_count": null,
      "metadata": {
        "id": "b96HGMSdtyWh"
      },
      "outputs": [],
      "source": [
        "# For debug only\n",
        "!gcloud config list --format 'value(core.account)'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cypUyEoaR8ye"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1db41b42c089"
      },
      "source": [
        "**Colab only:** Uncomment the following cell to initialize the Vertex AI SDK. For Vertex AI Workbench, you don't need to run this.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0caada39c488"
      },
      "outputs": [],
      "source": [
        "# PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
        "# LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
        "\n",
        "# from google.cloud import aiplatform\n",
        "# aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rsuh0WZFSGrR"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "from vertexai.language_models import TextGenerationModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vGwqGMPo04sv"
      },
      "source": [
        "### Set project and datasets for BigQuery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k9jwzkY11k0h"
      },
      "source": [
        "This is the project containing:\n",
        " - The linked BigQuery dataset `BQ_LINKED_DATASET` with your raw logs, and,\n",
        " - A new BigQuery dataset `BQ_PROCESSED_DATASET` you'll create to store the processed logs.\n",
        "\n",
        "This project could be the same or a separate project than the one you're using for Vertex AI.\n",
        "\n",
        "Make sure you have **BigQuery Data Viewer** role over `BQ_LINKED_DATASET` dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_spkCtnc1EQ-"
      },
      "outputs": [],
      "source": [
        "BQ_PROJECT_ID = \"[bq-project-id]\"  # @param {type:\"string\"}\n",
        "BQ_LINKED_DATASET = \"[linked-bq-dataset]\"  # @param {type:\"string\"}\n",
        "BQ_PROCESSED_DATASET = \"[new-bq-dataset]\"  # @param {type:\"string\"}\n",
        "\n",
        "from google.cloud import bigquery\n",
        "\n",
        "client = bigquery.Client(project=BQ_PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mkU4reYbeS4A"
      },
      "source": [
        "### Import models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t9bn8UGDfGoy"
      },
      "source": [
        "We will interact with Vertex AI LLM model `text-bison`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z6LSaXI_e0co"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"text-bison\"  # @param {type:\"string\"}\n",
        "\n",
        "model = TextGenerationModel.from_pretrained(MODEL_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "np2ZLxlC8trs"
      },
      "source": [
        "## Prepare the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KUIu1o3lB5PR"
      },
      "source": [
        "> You can skip this section if your raw logs are already processed and normalized in curated tables using [Dataform as part of Community Security Analytics](https://github.com/GoogleCloudPlatform/security-analytics/tree/main/dataform) (CSA). For more information on CSA and how to automatically and continuously build post-processed tables out of your raw logs, see this [Google Cloud blog post](https://cloud.google.com/blog/products/data-analytics/deploy-community-security-analytics-with-dataform)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CUIvePZJScBW"
      },
      "source": [
        "Like any other AI/ML project, first thing is to prepare your data including datasets for few-shot prompting and subsequent evaluation. You'll preprocess the raw logs that reside in your BigQuery linked dataset into a summary table into your new BigQuery dataset. This table will contain the logs in aggregated form and also normalized into a simple schema. This allows you to unlock and scale ML analysis:\n",
        "- From a computation point of view because the dataset is smaller and simple.\n",
        "- From a talent point of view because researchers and analysts are not required to be familiar with the complex schema of raw logs ([LogEntry definition](https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry)).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "19a3ff726f47"
      },
      "source": [
        "### Create new dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6da8207ff470"
      },
      "outputs": [],
      "source": [
        "!bq --location=US mk --dataset {BQ_PROJECT_ID}:{BQ_PROCESSED_DATASET}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qqFx8jPSyTEa"
      },
      "source": [
        "### Build table of user actions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tsPQNHhGyqBJ"
      },
      "source": [
        "Let's search and process the audit logs to create a table of user actions aggregated by day. This summary table reduces the set to relevant records and simplifies the structure which in turn simplifies exploratory and advanced analytics. For those interested, we use the same SQL query as defined in CSA repo, specifically [`csa_4_01_summary_daily`](https://github.com/GoogleCloudPlatform/security-analytics/blob/main/dataform/definitions/summary/csa_4_01_summary_daily.sqlx)  Dataform definition file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2yk-WLtorL5b"
      },
      "outputs": [],
      "source": [
        "TABLE_NAME = \"csa_4_01_summary_daily\"\n",
        "TABLE_ID = f\"{BQ_PROJECT_ID}.{BQ_PROCESSED_DATASET}.{TABLE_NAME}\"\n",
        "SUMMARY_LOOKBACK_DAYS = 90\n",
        "\n",
        "job_config = bigquery.QueryJobConfig(\n",
        "    destination=TABLE_ID, write_disposition=\"WRITE_TRUNCATE\"\n",
        ")\n",
        "\n",
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  EXTRACT(DATE FROM timestamp) AS day,\n",
        "  proto_payload.audit_log.authentication_info.principal_email,\n",
        "  ARRAY_AGG(DISTINCT proto_payload.audit_log.method_name IGNORE NULLS) AS actions,\n",
        "  COUNT(*) AS counter\n",
        "FROM `{BQ_PROJECT_ID}.{BQ_LINKED_DATASET}._AllLogs`\n",
        "WHERE\n",
        "  timestamp >=  TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {SUMMARY_LOOKBACK_DAYS} DAY)\n",
        "  AND proto_payload.audit_log.authentication_info.principal_email IS NOT NULL\n",
        "  AND proto_payload.audit_log.method_name NOT LIKE \"storage.%.get\"\n",
        "  AND proto_payload.audit_log.method_name NOT LIKE \"v1.compute.%.list\"\n",
        "  AND proto_payload.audit_log.method_name NOT LIKE \"beta.compute.%.list\"\n",
        "GROUP BY\n",
        "  day,\n",
        "  proto_payload.audit_log.authentication_info.principal_email\n",
        "\"\"\"\n",
        "\n",
        "# Start the query and save results in new table\n",
        "query_job = client.query(sql, job_config=job_config)\n",
        "result = query_job.result()  # Wait for the job to complete.\n",
        "\n",
        "print(f\"{result.total_rows} user action records loaded to table {TABLE_ID}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UipY2H4OnArZ"
      },
      "source": [
        "### Import sample queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A7SxIuTJFQJs"
      },
      "source": [
        "You will now retrieve a list of 15 sample security questions and corresponding SQL queries from a CSV file. These security questions are variations from the open-source [Community Security Analytics](https://github.com/GoogleCloudPlatform/security-analytics). CSA provides a set of security questions and corresponding queries for BigQuery, Log Analytics and Chronicle.\n",
        "\n",
        "We will use a subset of these queries as few-shot examples as part of the model prompt, and the remaining set for model evaluation.\n",
        "\n",
        "Run the following to read the CSV file from a GCS bucket and load all records into an in-memory pandas DataFrame:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gNbOxsMXwWz_"
      },
      "outputs": [],
      "source": [
        "BUCKET_ID = \"csa-datasets-public\"  # @param {type:\"string\"}\n",
        "FILENAME = \"SQL_Generator_Example_Queries.csv\"  # @param {type:\"string\"}\n",
        "\n",
        "df = pd.read_csv(f\"gs://{BUCKET_ID}/{FILENAME}\", header=0)\n",
        "df.head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_a7IbRewsOee"
      },
      "source": [
        "Retrieve table name referenced by the sample questions. This should be the same as the table of post-processed logs we created above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QOwkhv24ltod"
      },
      "outputs": [],
      "source": [
        "# Retrieve unique table names excluding null values and empty string\n",
        "BQ_TABLES = df[\"Qualified table name\"].replace(\"\", np.nan).dropna().unique()\n",
        "print(BQ_TABLES)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0hucGJI05c6i"
      },
      "source": [
        "### Extract train & eval datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JUtPNLdyG36Y"
      },
      "source": [
        "Extract train & eval datasets and store in respective dataframes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lcpSM5P35isB"
      },
      "outputs": [],
      "source": [
        "train_df = df.loc[df[\"Dataset\"] == \"Train\", [\"Question\", \"SQL Query\"]]\n",
        "eval_df = df.loc[df[\"Dataset\"] == \"Eval\", [\"Question\", \"SQL Query\"]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OqavOVDgjzAG"
      },
      "source": [
        "Take a peek at a few records from each set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oBW3dGefj3do"
      },
      "outputs": [],
      "source": [
        "train_df.head(5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_hzWuBV4j8gn"
      },
      "outputs": [],
      "source": [
        "eval_df.head(5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqeEuAV3Fr-R"
      },
      "source": [
        "## Prepare few-shot prompt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JxTajIC0Xlm"
      },
      "source": [
        "### Define prompt template"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7z_rB3XJ2eeU"
      },
      "source": [
        "The model prompt will include the following components:\n",
        "1. Concise statement to specify the task\n",
        "1. Schema definition to describe existing dataset\n",
        "1. A few shot examples of questions in natural language and corresponding SQL statements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fxQcNsjakJHE"
      },
      "source": [
        "This is the template we will later use to generate the prompt using these 3 components."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MrhnLpqjhKoV"
      },
      "outputs": [],
      "source": [
        "# This string template takes three arguments:\n",
        "# - schema definition\n",
        "# - few shot examples\n",
        "# - question for which query needs to be generated\n",
        "\n",
        "prompt_template = \"\"\"\\\n",
        "This is a task converting text into GoogleSQL statement.\n",
        "We will first give you the dataset schema and then ask a question in text.\n",
        "You are asked to generate SQL statement which is valid for BigQuery.\n",
        "Remove any delimiters around answer such as \"```\"\n",
        "\n",
        "BigQuery tables schema definition:\n",
        "{schema_definition}\n",
        "Here are a few shot examples:\n",
        "{few_examples}\n",
        "Write GoogleSQL query for following question: {question}\n",
        "Answer: \"Query here\"\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t03J_4G740C4"
      },
      "source": [
        "### Build schema definition (compact version)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A9NsA3rp3bGC"
      },
      "source": [
        "First, we need to build a concise schema definition of your dataset. As mentioned earlier, we'll use that as part of our prompt's context for grounding the results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eFeeXr3Svuce"
      },
      "source": [
        "Retrieve table and column definitions from the `INFORMATION_SCHEMA` of your BigQuery dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qEUiZ8BtSBxA"
      },
      "outputs": [],
      "source": [
        "# Following SQL query will generate schema definition of your dataset\n",
        "QUERY = f\"\"\"\\\n",
        "SELECT\n",
        "    '[Schema (values)]: ' || '| log_summary | ' || STRING_AGG(table_values, ' | ') || ';' AS tables_definition,\n",
        "    '[Column names (type)]: ' || STRING_AGG(column_names_types) || ';' AS columns_definition\n",
        "FROM (\n",
        "    SELECT\n",
        "      table_name,\n",
        "      table_name || ' : ' || STRING_AGG(column_name, ' , ') as table_values,\n",
        "      STRING_AGG(table_name || ' : ' || column_name || ' (' || data_type || ')', ' | ') as column_names_types\n",
        "    FROM {BQ_PROJECT_ID}.{BQ_PROCESSED_DATASET}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS\n",
        "    WHERE table_name IN {'(' + \",\".join(map(lambda x: f\"'{x}'\", BQ_TABLES)) + ')'}\n",
        "    GROUP BY table_name\n",
        "    ORDER BY table_name\n",
        ")\n",
        "\"\"\"\n",
        "\n",
        "# Create query job\n",
        "query_job = client.query(QUERY)\n",
        "# Get first row\n",
        "schema = next(query_job.result())\n",
        "\n",
        "# Build schema definition\n",
        "schema_definition = f\"\"\"\\\n",
        "{schema.tables_definition}\n",
        "\n",
        "{schema.columns_definition}\n",
        "\"\"\"\n",
        "\n",
        "print(schema_definition)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6F79laX04lV7"
      },
      "source": [
        "### Add queries as few-shot examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5H-RyYs7yC_T"
      },
      "source": [
        "This is a helper function to format example input and output within the prompt. We will use this helper function as we add shots to the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IlmSk-j6yCeO"
      },
      "outputs": [],
      "source": [
        "one_shot_template = \"\"\"\n",
        "Question: {question}\n",
        "\n",
        "Answer: {query}\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hZCf0wyOidE4"
      },
      "source": [
        "Let's add queries from our dataset as examples to our prompt:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HmfAgdcsicZo"
      },
      "outputs": [],
      "source": [
        "few_examples = \"\"\n",
        "for index, row in train_df.iterrows():\n",
        "    few_examples += one_shot_template.format(\n",
        "        question=row[\"Question\"], query=row[\"SQL Query\"]\n",
        "    )\n",
        "\n",
        "print(f\"Added {str(train_df.shape[0])} pairs as few-shot examples\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nAresy6SkwqF"
      },
      "source": [
        "### Review full prompt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6gnG1Ns17KLL"
      },
      "source": [
        "Using the schema definition, the few shot examples, and a sample question to be answered, let's generate a full prompt as an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0rQDO2-ElSmE"
      },
      "outputs": [],
      "source": [
        "question = \"This is a sample question\"\n",
        "\n",
        "prompt = prompt_template.format(\n",
        "    schema_definition=schema_definition, few_examples=few_examples, question=question\n",
        ")\n",
        "print(\"Prompt:\")\n",
        "print(prompt)\n",
        "\n",
        "print(\"Number of input tokens: \" + str(len(prompt)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QrUvSBGWc1JJ"
      },
      "source": [
        "## Generate SQL queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7oH5AXF2pZG"
      },
      "source": [
        "### Define helper function to generate SQL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hG1uHw9Loyhf"
      },
      "source": [
        "The following helper function `generate_sql()` is used to retrieve a SQL query from the Vertex AI LLM model using the prompt template we have built thus far.\n",
        "\n",
        "Notice how `generate_sql()` uses `sanitize_output()` function to strip the response down to the SQL query itself before returning the results. Even though the model prompt includes instructions to tune the model output, there may still be enclosing quotes or code block backticks which need to be stripped out to avoid a subsequent SQL syntax error."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lYwkEtikErBt"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "\n",
        "\n",
        "# Strip text to include only the SQL code block with\n",
        "def sanitize_output(text: str) -> str:\n",
        "    # Strip whitespace and any potential backticks enclosing the code block\n",
        "    text = text.strip()\n",
        "    regex = re.compile(r\"^\\s*```(\\w+)?|```\\s*$\")\n",
        "    text = regex.sub(\"\", text).strip()\n",
        "\n",
        "    # Find and remove any trailing quote without corresponding opening quote\n",
        "    if re.search(r'^[^\"]*\"$', text):\n",
        "        text = text[:-1]\n",
        "    # Find and remove any leading quote without corresponding closing quote\n",
        "    if re.search(r'^\"[^\"]*$', text):\n",
        "        text = text[1:]\n",
        "\n",
        "    return text\n",
        "\n",
        "\n",
        "# Call model using prompt and pre-defined parameters\n",
        "def generate_sql(\n",
        "    model,\n",
        "    prompt: str,\n",
        "    temperature: float = 0.2,\n",
        "    max_output_tokens: int = 1024,\n",
        "    top_k: int = 40,\n",
        "    top_p: float = 0.8,\n",
        ") -> str:\n",
        "    print(\"Generating SQL...\")\n",
        "    print(\"Number of input tokens: \" + str(len(prompt)))\n",
        "\n",
        "    response = model.predict(\n",
        "        prompt,\n",
        "        temperature=temperature,\n",
        "        max_output_tokens=max_output_tokens,\n",
        "        top_k=top_k,\n",
        "        top_p=top_p,\n",
        "    )\n",
        "\n",
        "    text = response.text\n",
        "    print(\"Number of output tokens: \" + str(len(text)))\n",
        "    print(\"Response:\")\n",
        "    print(text)\n",
        "\n",
        "    # Strip text to include only the SQL code block\n",
        "    text = sanitize_output(text)\n",
        "    print(\"Response stripped:\")\n",
        "    print(text)\n",
        "\n",
        "    return text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67g1gIfiov0Z"
      },
      "source": [
        "### Define helper function to execute SQL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l51jX7b_v4AJ"
      },
      "source": [
        "The following helper function `execute_sql()` is used to execute a SQL query against the live BigQuery dataset, and returning results as a dataframe.\n",
        "\n",
        "Notice how `execute_sql()` ensures to qualify table names with the project and BigQuery dataset you specified above, before executing the SQL query. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HcwXLPgCv4kf"
      },
      "outputs": [],
      "source": [
        "# Limit number of bytes processed as a guardrail for cost control\n",
        "BQ_MAX_BYTES_BILLED = pow(2, 30)  # 1GB\n",
        "\n",
        "\n",
        "def execute_sql(query: str):\n",
        "    print(\"Executing SQL...\")\n",
        "\n",
        "    # Qualify table names with your project and dataset ID\n",
        "    for table_name in BQ_TABLES:\n",
        "        query = query.replace(\n",
        "            table_name, f\"{BQ_PROJECT_ID}.{BQ_PROCESSED_DATASET}.{table_name}\"\n",
        "        )\n",
        "\n",
        "    print(\"Query:\")\n",
        "    print(query)\n",
        "\n",
        "    # Validate the query by performing a dry run without incurring a charge\n",
        "    job_config = bigquery.QueryJobConfig(use_query_cache=False, dry_run=True)\n",
        "    try:\n",
        "        response = client.query(query, job_config=job_config)\n",
        "    except Exception as e:\n",
        "        print(\"Error validating query:\")\n",
        "        print(e)\n",
        "        return e\n",
        "\n",
        "    print(f\"Query will process {response.total_bytes_processed / 1024:.2f} KB.\")\n",
        "\n",
        "    # Execute the query\n",
        "    job_config = bigquery.QueryJobConfig(\n",
        "        use_query_cache=False, maximum_bytes_billed=BQ_MAX_BYTES_BILLED\n",
        "    )\n",
        "    try:\n",
        "        response = client.query(query)\n",
        "        df = response.to_dataframe()\n",
        "    except Exception as e:\n",
        "        print(\"Error executing query:\")\n",
        "        print(e)\n",
        "        return e\n",
        "\n",
        "    return df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fUXRWeJgq046"
      },
      "source": [
        "### Example 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZv20YXOp8Cu"
      },
      "source": [
        "Let's generate the SQL to answer this sample question:\n",
        "\n",
        "*List all user actions that contains the word 'delete' or 'remove' over the last month. Include the user and the day in the results.*\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JgjytBJ1prXB"
      },
      "outputs": [],
      "source": [
        "question = \"List all user actions that contains the word 'delete' or 'remove' over the last month. Include the user and the day in the results.\"\n",
        "\n",
        "query = generate_sql(\n",
        "    model,\n",
        "    prompt_template.format(\n",
        "        schema_definition=schema_definition,\n",
        "        few_examples=few_examples,\n",
        "        question=question,\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pcjPwYyGMHoc"
      },
      "source": [
        "Let's test the generated query with the live dataset in BigQuery.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_70ym0AuMWd6"
      },
      "outputs": [],
      "source": [
        "# Execute the query\n",
        "query_result = execute_sql(query)\n",
        "print(query_result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aRaFdtf92QD"
      },
      "source": [
        "### Example 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fqs_0oe492QO"
      },
      "source": [
        "Let's generate the SQL to answer this sample question:\n",
        "\n",
        "*List any action containing IAM case-insensitive by any unapproved user over the last 7 days, where approved user include 'admin@example.com'.*\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-C50Vy392QP"
      },
      "outputs": [],
      "source": [
        "question = \"List any action containing IAM case-insensitive by any unapproved user over the last 7 days, where approved user include 'admin@example.com'\"\n",
        "\n",
        "query = generate_sql(\n",
        "    model,\n",
        "    prompt_template.format(\n",
        "        schema_definition=schema_definition,\n",
        "        few_examples=few_examples,\n",
        "        question=question,\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aTZUCI9h92QP"
      },
      "source": [
        "Let's test the generated query against your BigQuery dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zbd67RAB92QQ"
      },
      "outputs": [],
      "source": [
        "# Execute the query\n",
        "query_result = execute_sql(query)\n",
        "print(query_result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcEvAO8LFQUw"
      },
      "source": [
        "## Evaluate model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hh3xQUD9FW-s"
      },
      "source": [
        "### Run model on evaluation dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oh4geBloBuwW"
      },
      "source": [
        "Let's generate SQL queries for all questions in our evaluation dataset. That dataset includes both `Question` and the ground truth `SQL Query`. Run the following code to automatically call the model for each question in the dataset and record the response in a new column `Generated SQL Query`. This may take few minutes as model calls are done serially.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3uUHpyKHFXIr"
      },
      "outputs": [],
      "source": [
        "eval_df[\"Generated SQL Query\"] = eval_df[\"Question\"].apply(\n",
        "    lambda x: generate_sql(\n",
        "        model,\n",
        "        prompt_template.format(\n",
        "            schema_definition=schema_definition, few_examples=few_examples, question=x\n",
        "        ),\n",
        "    )\n",
        ")\n",
        "\n",
        "eval_df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3RUz6u0CeE7"
      },
      "outputs": [],
      "source": [
        "len(eval_df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0L56bIRm0eSu"
      },
      "source": [
        "### Execute ground truth queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IyWegUGQlVh7"
      },
      "source": [
        "Before we evaluate our generated queries, let's run the \"ground truth\" queries as part of our dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pihnnb0D0rk3"
      },
      "outputs": [],
      "source": [
        "eval_df[\"SQL Query Result\"] = eval_df[\"SQL Query\"].apply(execute_sql)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "It3GagA8mKsy"
      },
      "source": [
        "Let's peak into the results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-in-uRDimNeN"
      },
      "outputs": [],
      "source": [
        "eval_df.loc[:, [\"SQL Query\", \"SQL Query Result\"]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epWMpdnOFFcl"
      },
      "source": [
        "### Execute generated queries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dg_fO3uKInZf"
      },
      "outputs": [],
      "source": [
        "eval_df[\"Generated SQL Query Result\"] = eval_df[\"Generated SQL Query\"].apply(\n",
        "    execute_sql\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fP1Cu74Q1rdy"
      },
      "source": [
        "Let's peek into some of the results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uG830P_0hSbL"
      },
      "outputs": [],
      "source": [
        "eval_df.loc[:, [\"Generated SQL Query\", \"Generated SQL Query Result\"]].head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jrpMs3OkU4s0"
      },
      "source": [
        "### Serialize results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g41_fLqL2WXI"
      },
      "source": [
        " Since the results of each successful query are represented as a nested DataFrame within the evaluation DataFrame, we need to first serialize them so it's easier to compare the results, and to optionally save them in a spreadsheet or a CSV file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9-4PHg_D1lrA"
      },
      "outputs": [],
      "source": [
        "def format_query_result(query_result):\n",
        "    if isinstance(query_result, pd.DataFrame):\n",
        "        if query_result.shape[0] == 0:\n",
        "            return \"No results found\"\n",
        "        return query_result.to_csv(index=False)\n",
        "    elif isinstance(query_result, Exception):\n",
        "        return query_result.message\n",
        "    else:\n",
        "        return query_result\n",
        "\n",
        "\n",
        "eval_df[\"Generated SQL Query Result Formatted\"] = eval_df[\n",
        "    \"Generated SQL Query Result\"\n",
        "].apply(format_query_result)\n",
        "eval_df[\"SQL Query Result Formatted\"] = eval_df[\"SQL Query Result\"].apply(\n",
        "    format_query_result\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBeV8NvK13iB"
      },
      "source": [
        "Inspect the results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mxn9EbNz15_T"
      },
      "outputs": [],
      "source": [
        "eval_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dXMRFuoogNEu"
      },
      "source": [
        "### Calculate match score"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GCUknr-04Hgt"
      },
      "source": [
        "Let's evaluate the model accuracy by calculating a match score for each pair of queries. In our case, we'll calculate score based on fuzzy string matching of each generated query with the corresponding ground truth query."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ruyw56kmyHX6"
      },
      "source": [
        "Let's import and use `fuzzywuzzy` library which we already installed along with `Levenshtein` module for fast computation of Levenshtein distance between two strings:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F8W26_ORLuie"
      },
      "outputs": [],
      "source": [
        "from fuzzywuzzy import fuzz\n",
        "\n",
        "\n",
        "def get_match_of_queries(df):\n",
        "    return fuzz.partial_ratio(df[\"SQL Query\"], df[\"Generated SQL Query\"])\n",
        "\n",
        "\n",
        "eval_df[\"match_score_queries\"] = eval_df.apply(get_match_of_queries, axis=1)\n",
        "eval_df.loc[:, [\"SQL Query\", \"Generated SQL Query\", \"match_score_queries\"]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NHxlbLyMSil"
      },
      "source": [
        "Now calculate the mean score across all questions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyvMEN_YMSin"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"The average match score based on raw generated queries is: \",\n",
        "    round(eval_df[\"match_score_queries\"].mean(), 2),\n",
        "    \"%\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmfyDMMy2Vrw"
      },
      "source": [
        "The generated queries might have a different SQL implementation or variables naming than the 'ground truth' query, yet still yield the correct answer to the security question. So calculating the match score of each query string with its corresponding 'ground truth' query string is not a sufficient evaluation method. Therefore, we'll also calculate the match score between the actual results returned from your dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VyTuVS0L0z_T"
      },
      "source": [
        "Let's run the fuzzy match logic on the formatted version of the results. That formatted column is already stringified and ready for string comparison. This may take several minutes to complete depending on size of actual results being compared."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wIImFcc-47bV"
      },
      "outputs": [],
      "source": [
        "def get_match_of_results(df):\n",
        "    return fuzz.partial_ratio(\n",
        "        df[\"SQL Query Result Formatted\"], df[\"Generated SQL Query Result Formatted\"]\n",
        "    )\n",
        "\n",
        "\n",
        "# This may take several minutes to complete\n",
        "eval_df[\"match_score\"] = eval_df.apply(get_match_of_results, axis=1)\n",
        "eval_df.loc[\n",
        "    :,\n",
        "    [\n",
        "        \"SQL Query\",\n",
        "        \"Generated SQL Query\",\n",
        "        \"SQL Query Result Formatted\",\n",
        "        \"Generated SQL Query Result Formatted\",\n",
        "        \"match_score_queries\",\n",
        "        \"match_score\",\n",
        "    ],\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8XJ3Not74fu-"
      },
      "source": [
        "Now calculate the mean score across all questions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "slz0gcg14jJl"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"The average match score of all generated queries is: \",\n",
        "    round(eval_df[\"match_score\"].mean(), 2),\n",
        "    \"%\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yYfmxEznN21D"
      },
      "source": [
        "## [Optional] Save results in Google Sheets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JtBb7x36_Ij_"
      },
      "source": [
        "You may want to save all generated queries, results and scores into a Google Sheet for visual inspection and for future reference like model accuracy tracking.\n",
        "\n",
        "Skip this section if using **Vertex AI Workbench** or **Colab Enterprise** because notebooks in those environments cannot access Google Drive or Google Sheets for security purposes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aTShdbsVQ0Id"
      },
      "source": [
        "Create a new spreadsheet in Google Sheets (https://sheets.new) and copy over your unique spreadsheet ID into `QUERIES_SHEET_ID` parameter.  You can find your spreadsheet ID in the Google Sheets URL: docs.google.com/spreadsheets/d/*spreadsheetId*/edit#gid=0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "juEtvr6im9xh"
      },
      "outputs": [],
      "source": [
        "QUERIES_SHEET_ID = \"\"  # @param {type:\"string\"}\n",
        "QUERIES_WORKSHEET_NAME = \"Evaluation Dataset\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wXtDBgDQN21Q"
      },
      "outputs": [],
      "source": [
        "from google.auth import default\n",
        "import gspread\n",
        "\n",
        "# Authenticate with Google Sheets\n",
        "creds, _ = default()\n",
        "gc = gspread.authorize(creds)\n",
        "\n",
        "wks_results = gc.open_by_key(QUERIES_SHEET_ID).worksheet(QUERIES_WORKSHEET_NAME)\n",
        "\n",
        "# Drop Query Result column which may contain non-serializable objects\n",
        "eval_df_copy = eval_df.drop(columns=[\"SQL Query Result\", \"Generated SQL Query Result\"])\n",
        "\n",
        "\n",
        "def limit_cell_length(cell) -> str:\n",
        "    if len(cell) >= 50000:\n",
        "        return cell[:49990] + \"...\"\n",
        "    return cell\n",
        "\n",
        "\n",
        "eval_df_copy[\"Generated SQL Query Result Formatted\"] = eval_df_copy[\n",
        "    \"Generated SQL Query Result Formatted\"\n",
        "].apply(limit_cell_length)\n",
        "eval_df_copy[\"SQL Query Result Formatted\"] = eval_df_copy[\n",
        "    \"SQL Query Result Formatted\"\n",
        "].apply(limit_cell_length)\n",
        "\n",
        "wks_results.update(\n",
        "    [eval_df_copy.columns.values.tolist()] + eval_df_copy.values.tolist()\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uizDMHaGVYFA"
      },
      "source": [
        "## Summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p7V58R6pSdyI"
      },
      "source": [
        "In this notebook, you were able to:\n",
        "- Prepare datasets including writing logs into a summary table for easier and faster log analysis.\n",
        "- Build a prompt template using an existing dataset of Text:SQL pairs as few-shot examples combined with a dynamically retrieved context, that is the database schema, for grounding the model.\n",
        "- Convert NL question into SQL query using the language model and a few-shot prompt.\n",
        "- Sanitize and validate model output.\n",
        "- Evaluate model output by executing generated queries on BigQuery\n",
        "- Run model on an entire evaluation dataset\n",
        "- Calculate match score based on model output, that is the generated queries compared to the ground truth queries\n",
        "- Calculate match score based on actual results of generated queries, compared to the results of the ground truth queries.\n",
        "- Save results for future reference or for tracking model accuracy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tKHcKv3HVb4V"
      },
      "source": [
        "## Cleanup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8992fa1f9e2d"
      },
      "source": [
        "To clean up all Google Cloud resources used in this notebook, 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, namely the BigQuery dataset `BQ_PROCESSED_DATASET` with the processed data: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "341f3d460d32"
      },
      "outputs": [],
      "source": [
        "# Delete the created BigQuery dataset\n",
        "!bq rm -r -f {BQ_PROJECT_ID}:{BQ_PROCESSED_DATASET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sql_code_generation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}