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

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "47bgjiSzeJnA"
      },
      "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": "DF9B5dFQeLu7"
      },
      "source": [
        "# SQL Code Generation on Vertex AI using LangChain 🦜🔗\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/main/gemini/use-cases/retrieval-augmented-generation/NLP2SQL_using_dynamic_RAG.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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_langchain.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": "80103b348da3"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Authors: | [Shubham Chawla](https://www.github.com/shubhamgoogle), [Roy Arsan](https://www.linkedin.com/in/arsan) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvryK2kGeWfP"
      },
      "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 with [LangChain Google Cloud Vertex AI](https://python.langchain.com/docs/integrations/llms/google_vertex_ai_palm) implementation, 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 Retrieval Augmented Generation (RAG) by using [SemanticSimilarityExampleSelector](https://python.langchain.com/docs/modules/model_io/prompts/example_selector_types/similarity) from LangChain Example Selector to dynamically retrieve and pass the most relevant few shot examples to enrich the LLM prompt. This helps ensure most accurate and relevant LLM output, that is the generated SQL query, while limiting number of required LLM input tokens and thereby cost. 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": "ool643Tzg7W2"
      },
      "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",
        "  * Use VertexAIEmbeddings to create embeddings\n",
        "  * Use LangChain Example Selector to automatically select relevant examples for 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 pandas dataframe equals check\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": "iqxB1023g_fM"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R2Vfi5W0hNto"
      },
      "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": "pIZYzBgHhT6R"
      },
      "source": [
        "### Install SDKs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z3EARCs46U-I"
      },
      "outputs": [],
      "source": [
        "# Install Vertex AI SDK to use for model predictions\n",
        "%pip install google-cloud-aiplatform google-cloud-bigquery pandas --upgrade --user\n",
        "%pip install --upgrade --quiet langchain langchain-core langchain-google-vertexai\n",
        "%pip install langchain-community faiss-cpu"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbScDyV5hf1D"
      },
      "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": "8eJdvWP-Dq35"
      },
      "outputs": [],
      "source": [
        "# # Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "akTXroE0ChYP"
      },
      "source": [
        "### Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gBMjDTw59cFt"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "import sys\n",
        "\n",
        "from IPython.display import display\n",
        "from google.cloud import aiplatform, bigquery\n",
        "from langchain.prompts.example_selector import SemanticSimilarityExampleSelector\n",
        "from langchain.prompts.few_shot import FewShotPromptTemplate\n",
        "from langchain.prompts.prompt import PromptTemplate\n",
        "from langchain_community.vectorstores import FAISS\n",
        "from langchain_google_vertexai import VertexAI, VertexAIEmbeddings\n",
        "import numpy as np\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "akW4u3cIjLG7"
      },
      "source": [
        "### Set project and datasets for BigQuery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZWnTdzIjLM6"
      },
      "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": "Q-mYUuAS7CeX"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "LOCATION_US = \"US\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\n",
        "BQ_LINKED_DATASET = \"\"  # @param {type:\"string\"}\n",
        "BQ_PROCESSED_DATASET = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "66PtfF1PhnEV"
      },
      "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": "WPy4O74t_LU2"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EhcMpyDBKexo"
      },
      "source": [
        "### Set Embedding and Vertex AI LLM Model\n",
        "\n",
        "In the current example we are using text-bison@002 large language model but you can use other Google provided models gemini, gemini-pro, ulta, etc. For embedding we are using textembedding-gecko with the latest version."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q4vG3p5oKd92"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"text-bison@002\"  # @param {type:\"string\"}\n",
        "EMBEDDING_MODEL_ID = \"textembedding-gecko@latest\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sipwSsYKCEMe"
      },
      "source": [
        "### Import sample queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KIbiuWUTCHDD"
      },
      "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": "O_9ZBVhjIXFI"
      },
      "outputs": [],
      "source": [
        "BUCKET_ID = \"csa-datasets-public\"  # @param {type:\"string\"}\n",
        "FILENAME = \"SQL_Generator_Example_Queries.csv\"  # @param {type:\"string\"}\n",
        "df = pd.read_csv(f\"gs://{BUCKET_ID}/{FILENAME}\", header=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4QBZwTUnDnuC"
      },
      "source": [
        "### Extract train & eval datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bPoHff6CDslh"
      },
      "source": [
        "Extract train & eval datasets and store in respective dataframes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x24U6jDzDo0E"
      },
      "outputs": [],
      "source": [
        "train_df = df.loc[df[\"Dataset\"] == \"Train\", [\"Question\", \"SQL Query\"]]\n",
        "eval_df = df.loc[df[\"Dataset\"] == \"Eval\", [\"Question\", \"SQL Query\"]]\n",
        "train_dict = (\n",
        "    train_df[[\"Question\", \"SQL Query\"]]\n",
        "    .rename(columns={\"SQL Query\": \"answer\"})\n",
        "    .rename(columns={\"Question\": \"question\"})\n",
        "    .to_dict(orient=\"records\")\n",
        ")\n",
        "train_df.head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_6AtPphSBjxn"
      },
      "source": [
        "## Prepare the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GC6Wbd6EBnM4"
      },
      "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": "HB1Ki-d5Bp75"
      },
      "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": "gR65cweHLbA_"
      },
      "source": [
        "### Create new dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yJ4ufY5pLcEK"
      },
      "outputs": [],
      "source": [
        "!bq --location=LOCATION_US mk --dataset {BQ_PROJECT_ID}:{BQ_PROCESSED_DATASET}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pX3XEZ4MLMac"
      },
      "source": [
        "Create new csa_4_01_summary_daily using log analytics BigQuery table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWvGJw0C7YK2"
      },
      "outputs": [],
      "source": [
        "TABLE_NAME = \"csa_4_01_summary_daily\"\n",
        "TABLE_ID = f\"{PROJECT_ID}.{BQ_PROCESSED_DATASET}.{TABLE_NAME}\"\n",
        "SUMMARY_LOOKBACK_DAYS = 90\n",
        "\n",
        "client = bigquery.Client(project=PROJECT_ID, location=LOCATION_US)\n",
        "client.create_dataset(dataset=BQ_PROCESSED_DATASET, exists_ok=True)\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 `{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": "TRA74CdhEODS"
      },
      "source": [
        "### Build schema definition (compact version)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TJneRks9ERIt"
      },
      "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": "MyZHLd-oEozF"
      },
      "source": [
        "Retrieve table and column definitions from the `INFORMATION_SCHEMA` of your BigQuery dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n49TC0PaJ23L"
      },
      "outputs": [],
      "source": [
        "# Following SQL query will generate schema definition of your dataset\n",
        "\n",
        "BQ_TABLES = df[\"Qualified table name\"].replace(\"\", np.nan).dropna().unique()\n",
        "print(BQ_TABLES)\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 {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": "0aaf6279479d"
      },
      "source": [
        "### Build Prompt\n",
        "\n",
        "Lets create prompt using user input and few shots extracted dynamically using LangChain FewShotPromptTemplate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "419c91c9b825"
      },
      "source": [
        "#### Create Vertex AI embeddings to create a text embedding  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3632901bda94"
      },
      "outputs": [],
      "source": [
        "embedding = VertexAIEmbeddings(model_name=EMBEDDING_MODEL_ID, project=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eh0vG61jR6Bb"
      },
      "source": [
        "#### Create Few Shot Prompt\n",
        "\n",
        "The `SemanticSimilarityExampleSelector` selects examples based on a combination of which examples are most similar to the inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IftnEU2SAiJ8"
      },
      "outputs": [],
      "source": [
        "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
        "    # This is the list of examples available to select from.\n",
        "    train_dict,\n",
        "    # This is the embedding class used to produce embeddings which are used to measure semantic similarity.\n",
        "    embedding,\n",
        "    # This is the VectorStore class that is used to store the embeddings and do a similarity search over.\n",
        "    FAISS,\n",
        "    # This is the number of examples to produce.\n",
        "    k=2,\n",
        ")\n",
        "\n",
        "# Select the most similar example to the input.\n",
        "question = \"select user actions that contains the word 'delete' or 'remove'\"\n",
        "selected_examples = example_selector.select_examples({\"question\": question})\n",
        "print(f\"Examples most similar to the input: {question}\")\n",
        "for example in selected_examples:\n",
        "    print(\"\\n\")\n",
        "    for k, v in example.items():\n",
        "        print(f\"{k}: {v}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cfb4c3799a6f"
      },
      "source": [
        "#### Helper Function to Build Prompt\n",
        "\n",
        "Below function will be utilised to converting user input into actual prompt with few shots and a prefix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90f64ee30e90"
      },
      "outputs": [],
      "source": [
        "def build_prompt(user_prompt, example_selector):\n",
        "    prompt_template = f\"\"\"\\\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",
        "    \"\"\"\n",
        "    example_prompt = PromptTemplate(\n",
        "        input_variables=[\"question\", \"answer\"],\n",
        "        template=\"question: {question}\\nanswer: {answer}\",\n",
        "    )\n",
        "\n",
        "    prompt = FewShotPromptTemplate(\n",
        "        example_selector=example_selector,\n",
        "        example_prompt=example_prompt,\n",
        "        prefix=prompt_template,\n",
        "        suffix=\"question: {question}\\nanswer: \",\n",
        "        input_variables=[\"question\"],\n",
        "    )\n",
        "    final_prompt = prompt.format(question=user_prompt)\n",
        "    return final_prompt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jhxjQ6F0Fol0"
      },
      "source": [
        "## Generate SQL queries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JK8aYpx0Fse0"
      },
      "source": [
        "### Define helper function to generate SQL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nfqswuz2Fwqw"
      },
      "source": [
        "`generate_sql()`: This function is used to retrieve a SQL query from the Vertex AI LLM model using the prompt template we have built thus far.\n",
        "\n",
        "`execute_sql()`: This function is used to execute a SQL query against the live BigQuery dataset, and returning results as a dataframe.\n",
        "\n",
        "`build_prompt()`: This function is used to create the final prompt which includes common prefix and suffix for all the prompts\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": "O45Xw1uql0rF"
      },
      "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\"{PROJECT_ID}.{BQ_PROCESSED_DATASET}.{table_name}\"\n",
        "        )\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\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",
        "    example_selector,\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",
        "    model = VertexAI(\n",
        "        model_name=MODEL_ID,\n",
        "        temperature=temperature,\n",
        "        max_output_tokens=max_output_tokens,\n",
        "        top_k=top_k,\n",
        "        top_p=top_p,\n",
        "    )\n",
        "    final_prompt = build_prompt(prompt, example_selector)\n",
        "    print(final_prompt)\n",
        "    text = model.invoke(final_prompt)\n",
        "    print(\"Number of output tokens: \" + str(len(text)))\n",
        "    print(\"Response:\")\n",
        "    print(text)\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": "--NVnQyodhAC"
      },
      "source": [
        "### Example 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sR9hiB2Tdx53"
      },
      "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": "Wmk0c2H_lhcB"
      },
      "outputs": [],
      "source": [
        "user_prompt = \"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",
        "final_generated_prompt = build_prompt(user_prompt, example_selector)\n",
        "print(final_generated_prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kF0_OL5Odvdn"
      },
      "source": [
        "Let's test the generated query with the live dataset in BigQuery."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gEaH9nZbOloi"
      },
      "outputs": [],
      "source": [
        "output = generate_sql(example_selector, user_prompt)\n",
        "print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6uzuKxoud6Cf"
      },
      "source": [
        "Let's test the generated query against your BigQuery dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HeNR3q_xPfO2"
      },
      "outputs": [],
      "source": [
        "# Execute the query\n",
        "query_result = execute_sql(output)\n",
        "display(query_result.head(2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vnb2LuGjLjpN"
      },
      "source": [
        "### Example 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hW-fLol-LjpX"
      },
      "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": "JyE38cipLjpY"
      },
      "outputs": [],
      "source": [
        "user_prompt = \"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",
        "final_generated_prompt = build_prompt(user_prompt, example_selector)\n",
        "print(final_generated_prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AAG4I30NLjpY"
      },
      "source": [
        "Let's test the generated query with the live dataset in BigQuery."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h_LksGE3LjpY"
      },
      "outputs": [],
      "source": [
        "output = generate_sql(example_selector, user_prompt)\n",
        "print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9nqeQsWGLjpY"
      },
      "source": [
        "Let's test the generated query against your BigQuery dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJq19QC_LjpY"
      },
      "outputs": [],
      "source": [
        "# Execute the query\n",
        "query_result = execute_sql(output)\n",
        "display(query_result.head(2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QmI4izZRd-Id"
      },
      "source": [
        "## Evaluate model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3dkE8mCeFri"
      },
      "source": [
        "### Run model on evaluation dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmrMER40eI4O"
      },
      "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": "7UwUYtOEg-mf"
      },
      "outputs": [],
      "source": [
        "eval_df[\"Generated SQL Query\"] = eval_df[\"Question\"].apply(\n",
        "    lambda x: generate_sql(example_selector, x)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "udQqttGIGNmM"
      },
      "source": [
        "### Compare output result\n",
        "\n",
        "In the next code cell, we'll execute the original SQL query and then compare its output directly to the output of the generated SQL query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rffyi5XX1Giu"
      },
      "outputs": [],
      "source": [
        "def compare_dataframes(sql_query, generated_sql_query):\n",
        "    \"\"\"Compares two pandas DataFrames row-wise using columns from the second DataFrame.\n",
        "    Args:\n",
        "        SQL Query, Generated SQL Query\n",
        "    Returns:\n",
        "        True if output of both the SQL queries matches otherwise False\n",
        "    \"\"\"\n",
        "    df1 = execute_sql(sql_query)\n",
        "    df2 = execute_sql(generated_sql_query)\n",
        "\n",
        "    # If generated query returned an error instead of a dataframe with results:\n",
        "    if not isinstance(df2, pd.DataFrame):\n",
        "        return False\n",
        "\n",
        "    try:\n",
        "        df2 = df2[df1.columns]\n",
        "    except KeyError:\n",
        "        # Columns in results of ground truth query are missing\n",
        "        # from results returned by generated query\n",
        "        return False\n",
        "\n",
        "    comparison_result = df2.eq(df1)\n",
        "    matching_rows = comparison_result.all(axis=1)\n",
        "    matching_count = matching_rows.sum()\n",
        "    # return df1, df2\n",
        "    return True if matching_count == len(df1) else False\n",
        "\n",
        "\n",
        "eval_df[\"Data Match\"] = eval_df.apply(\n",
        "    lambda x: compare_dataframes(x[\"SQL Query\"], x[\"Generated SQL Query\"]), axis=1\n",
        ")\n",
        "# eval_df[\"sql_query_output\"],eval_df[\"generated_sql_query_output\"] = eval_df.apply(lambda x: compare_dataframes(x[\"SQL Query\"], x[\"Generated SQL Query\"]), axis=1)\n",
        "\n",
        "# Note: To save the output data to the final dataframe, make these changes: 1. Uncomment lines 26 and 30. 2. Comment out line 29."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ydZeqe3FGUQf"
      },
      "source": [
        "## Final Result\n",
        "\n",
        "In the next cell, we'll calculate the model's final score. This score represents the percentage of successful matches between the original and generated queries, as indicated in the 'Data Match' column."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S-a85rHhNCs2"
      },
      "source": [
        "### Score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3SnMTlFSHV17"
      },
      "outputs": [],
      "source": [
        "def get_prcntg_match(eval_df):\n",
        "    return round(eval_df[\"Data Match\"].sum() / len(eval_df) * 100)\n",
        "\n",
        "\n",
        "prcntg_match = get_prcntg_match(eval_df)\n",
        "print(f\"Final Score based on the percentage of data match: {prcntg_match}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2nQc4ePmIa-7"
      },
      "source": [
        "### Output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1eQxydrDEA-R"
      },
      "outputs": [],
      "source": [
        "pd.set_option(\"display.max_rows\", None)\n",
        "pd.set_option(\"display.max_columns\", None)\n",
        "pd.set_option(\"display.width\", None)\n",
        "pd.set_option(\"display.max_colwidth\", -1)\n",
        "display(eval_df)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sql_code_generation_langchain.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}