text2sql_astradb_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74043b68-324d-44c9-a191-9f13d37b7001",
   "metadata": {
    "id": "74043b68-324d-44c9-a191-9f13d37b7001"
   },
   "source": [
    "# Dynamic Few-Shot Prompting for SQL generation with Astra DB"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "607aeb3b-5245-4e8b-a710-d6d5026d602c",
   "metadata": {
    "id": "607aeb3b-5245-4e8b-a710-d6d5026d602c"
   },
   "source": [
    "Generating SQL to answer natural language questions is _hard_. There has been a large amount of research in this area, even before the rise of LLMs. Although every individual problem may require specific techniques to achieve the best performance, there are still some general principles that can apply across many domains.\n",
    "\n",
    "This notebook will demonstrate one of those generalizable techniques, \"Dynamic Few-Shot Prompting\", using Astra DB. This technique improves the accuracy of the generated SQL queries."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e2d3f3-9bc0-4bbf-835d-f6493cea7b09",
   "metadata": {},
   "source": [
    "In this notebook, you will:\n",
    "- download a dataset of questions and corresponding answering SQL statements\n",
    "- populate a relational (SQL) database with sample data to run the above queries\n",
    "- learn how to query an LLM to have it generate SQL given a natural-language question\n",
    "- ... then refine the technique with a few-shot improvement\n",
    "- ... then make the few-shot selection dynamic with the help of a vector database\n",
    "- quantify the effectiveness of the approaches\n",
    "\n",
    "At the end, you will have a clear picture of the main challenges and approach to generating SQL with Large Language Models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeea60cf-db7e-4152-b5eb-b134a9107aa5",
   "metadata": {},
   "source": [
    "To run the notebook, you will need:\n",
    "- OpenAI API key ([details](https://platform.openai.com/docs/quickstart/account-setup))\n",
    "- an Astra DB API Endpoint and associated Database Token ([details](https://docs.datastax.com/en/astra/astra-db-vector/get-started/quickstart.html#create-a-serverless-vector-database))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc2a69df-0fb0-43d3-b571-6bb26477c6c6",
   "metadata": {},
   "source": [
    "## About dynamic few-shot prompting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "843a7206-c4d2-4a64-9db9-13a5e6f997a4",
   "metadata": {},
   "source": [
    "It is a widely known fact that when prompting LLMs to generate an answer, listing some example question-answer pairs together with the actual question greatly improves the quality of the result. This technique is termed \"few-shot prompting\".\n",
    "\n",
    "Dynamic few-shot prompting is a variation on few-shot prompting where the _choice- of the examples is depending on the very question being asked. In other words, the selected question-answer pairs for the prompt are extracted from a larger collection of examples in such a way to be very similar to the actual question. This increases the effectiveness of few-shot prompting.\n",
    "\n",
    "If you generate vector embeddings for the example questions and store that information in an Astra DB collection, then you can use vector search to locate the most similar examples in milliseconds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffbdf67e-26b7-4713-8118-f5375488cc30",
   "metadata": {
    "id": "ffbdf67e-26b7-4713-8118-f5375488cc30"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a96785dd-6dcf-43c9-ba6e-6fe3b6ab9835",
   "metadata": {
    "id": "a96785dd-6dcf-43c9-ba6e-6fe3b6ab9835"
   },
   "source": [
    "### Requirements\n",
    "\n",
    "_Note: this demo uses the latest Data API python client, `astrapy 2.*`, which is in [pre-release](https://docs.datastax.com/en/astra-db-serverless/api-reference/client-versions.html#version-2-0-preview) at this time._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e052fbf7-d3bc-4c86-af86-2384f5d7bfcb",
   "metadata": {
    "id": "e052fbf7-d3bc-4c86-af86-2384f5d7bfcb"
   },
   "outputs": [],
   "source": [
    "!pip install -q \\\n",
    "  \"openai>=1.0,<2.0\" \\\n",
    "  \"astrapy==2.0.0-rc1\" \\\n",
    "  \"datasets==3.*\" \\\n",
    "  \"tenacity==9.*\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8f88ba3-136a-4a16-9032-9c1794e1ccf7",
   "metadata": {
    "id": "c8f88ba3-136a-4a16-9032-9c1794e1ccf7"
   },
   "source": [
    "### Import dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "51809f71-3869-4dc1-8c4e-56dd78abc7c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import re\n",
    "import sqlite3\n",
    "from functools import lru_cache, partial\n",
    "from getpass import getpass\n",
    "from typing import Any, Callable, Dict, List, Tuple\n",
    "\n",
    "import openai\n",
    "import pandas as pd\n",
    "from astrapy import DataAPIClient\n",
    "from astrapy.api_options import APIOptions, SerdesOptions\n",
    "from astrapy.info import CollectionDefinition\n",
    "from datasets import load_dataset\n",
    "from tenacity import retry, wait_exponential\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43ace97a-f446-496d-9127-ad5cb77b5291",
   "metadata": {},
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e68ce01-3193-4e35-b253-ad6cb8ce6f6c",
   "metadata": {},
   "source": [
    "This demo uses OpenAI for the LLM and the embeddings. The demo is set to use the GPT3.5 model, which is cheaper and - being less powerful than the most recent ones - better highlights the increase in accuracy of the improvements exemplified in this notebook. Feel free to edit these settings to your liking (including, with the necessary changes later on, switching to a different model vendor altogether). For example, you may want to use the \"gpt-4o-mini-2024-07-18\", or \"gpt-4o-2024-08-06\", OpenAI model.\n",
    "\n",
    "An embedding model is needed for later calculation of embedding vectors for the question examples.\n",
    "\n",
    "The SQL, relational database for this demo will be a local SQLite file for simplicity. You can, instead, choose to target Postgres, Amazon Aurora, Amazon Redshift or any other SQL-compatible database by rewriting a single function given in a later cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c04f400-0aba-40e1-b45d-7d5711a8f50b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM and embedding model settings\n",
    "LLM_MODEL_NAME = \"gpt-3.5-turbo\"\n",
    "\n",
    "EMBEDDING_MODEL = \"text-embedding-3-small\"\n",
    "EMBEDDING_DIMENSION = 1536\n",
    "\n",
    "# Collection for storing the examples for few-shot prompting\n",
    "ASTRA_DB_COLLECTION_NAME = \"text2sql_examples\"\n",
    "\n",
    "# This is needed for the SQLite database\n",
    "SQLITE_FILE_NAME = \"sample_database.db\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "744af48d-fa9c-4331-af76-36fccf8d7761",
   "metadata": {},
   "source": [
    "### initialize OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b44ee15d-943f-481e-a0a1-6e1e1d8b2ec7",
   "metadata": {
    "id": "b44ee15d-943f-481e-a0a1-6e1e1d8b2ec7"
   },
   "outputs": [],
   "source": [
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"OpenAI API Key: \")\n",
    "\n",
    "client = openai.OpenAI()\n",
    "\n",
    "# We equip the LLM and embedding calls with exponential backoff and retry\n",
    "# in case the service is throttled\n",
    "\n",
    "@retry(wait=wait_exponential(multiplier=1.2, min=4, max=30))\n",
    "def llm_completion(prompt: str) -> str:\n",
    "    response = client.chat.completions.create(\n",
    "        messages=[{\n",
    "            \"role\": \"user\",\n",
    "            \"content\": prompt,\n",
    "        }],\n",
    "        model=LLM_MODEL_NAME,\n",
    "    ).choices[0].message.content\n",
    "    return response\n",
    "\n",
    "\n",
    "@retry(wait=wait_exponential(multiplier=1.2, min=4, max=30))\n",
    "def compute_embedding_list(texts: List[str]) -> List[List[float]]:\n",
    "    return [\n",
    "        data.embedding\n",
    "        for data in client.embeddings.create(\n",
    "            input=texts,\n",
    "            model=EMBEDDING_MODEL,\n",
    "            timeout=10,\n",
    "        ).data\n",
    "    ]\n",
    "\n",
    "\n",
    "def compute_embedding(text: str) -> List[float]:\n",
    "    return compute_embedding_list([text])[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3f41179-13ea-4250-a0fb-6e330cef6db7",
   "metadata": {},
   "source": [
    "### Initialize Astra DB client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5e76c8dc-0e8f-40dd-9cd9-da203b193fc5",
   "metadata": {
    "id": "5e76c8dc-0e8f-40dd-9cd9-da203b193fc5",
    "outputId": "c948b8a7-27e4-47fb-af40-7fd743b14aa8"
   },
   "outputs": [],
   "source": [
    "ASTRA_DB_API_ENDPOINT = os.environ.get(\"ASTRA_DB_API_ENDPOINT\")\n",
    "if not ASTRA_DB_API_ENDPOINT:\n",
    "    ASTRA_DB_API_ENDPOINT = input(\"Astra DB API Endpoint: \")\n",
    "ASTRA_DB_APPLICATION_TOKEN = os.environ.get(\"ASTRA_DB_APPLICATION_TOKEN\")\n",
    "if not ASTRA_DB_APPLICATION_TOKEN:\n",
    "    ASTRA_DB_APPLICATION_TOKEN = getpass(\"Astra DB Token: \")\n",
    "\n",
    "data_api_client = DataAPIClient()\n",
    "astra_db = data_api_client.get_database(\n",
    "    ASTRA_DB_API_ENDPOINT,\n",
    "    token=ASTRA_DB_APPLICATION_TOKEN,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "529eedd2-6f55-429f-865a-b7f6b848d7de",
   "metadata": {},
   "source": [
    "### Initialize SQL execution\n",
    "\n",
    "**Note:** in order to target a different database supporting SQL, all you have to do is provide a different implementation of the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6175152d-5e28-494b-9f16-e1c91028fefe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_sql(\n",
    "    *statements: List[str], raise_on_error: bool = True\n",
    ") -> List[Tuple[Any, ...]]:\n",
    "    \"\"\"\n",
    "    Utility to execute DB SQL statements and return the result\n",
    "    of the final query.\n",
    "    \"\"\"\n",
    "    with sqlite3.connect(SQLITE_FILE_NAME) as connection:\n",
    "        cursor = connection.cursor()\n",
    "        try:\n",
    "            for statement in statements:\n",
    "                cursor.execute(statement)\n",
    "        except sqlite3.OperationalError as e:\n",
    "            # syntax errors or similar in running the query\n",
    "            if raise_on_error:\n",
    "                raise\n",
    "            return []\n",
    "        try:\n",
    "            result = cursor.fetchall()\n",
    "        except sqlite3.OperationalError as e:\n",
    "            # No result set (final statement was not a SELECT)\n",
    "            result = []\n",
    "\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f87a5931-557e-4eaf-abbf-7e23f1c9e5de",
   "metadata": {
    "id": "f87a5931-557e-4eaf-abbf-7e23f1c9e5de"
   },
   "source": [
    "## Load and prepare query data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a79da8-7a2c-475c-b60a-5ad2704ffba5",
   "metadata": {
    "id": "68a79da8-7a2c-475c-b60a-5ad2704ffba5"
   },
   "source": [
    "For this demo, we will use the [Spider](https://yale-lily.github.io/spider) dataset, which has been a standard to evaluate generated SQL performance for a few years now. This dataset consists of `question`, `query` pairs (**Example pairs**) to indicate the ideal query to be generated from a given natural language question.\n",
    "\n",
    "_Note: to apply this approach to a specific use case, the best strategy to collect data is to store the actual SQL generated by an application in some live environment. You can then either have a quick feedback interface for users to thumb-up/thumb-down these queries, or have internal human evaluators grade them offline. As more approved queries become available, you will have more and more example pairs to feed into the SQL-generating application, increasing its effectiveness._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "29b68b71-feef-4528-87dd-01345cfea7f0",
   "metadata": {
    "id": "29b68b71-feef-4528-87dd-01345cfea7f0"
   },
   "outputs": [],
   "source": [
    "spider = load_dataset(\"spider\", split=\"validation\")\n",
    "spider_df = spider.to_pandas()\n",
    "spider_schema = load_dataset(\"richardr1126/spider-schema\", split=\"train\")\n",
    "spider_schema_df = spider_schema.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2feb2a43-958f-4125-9218-3b8addf001d5",
   "metadata": {
    "id": "2feb2a43-958f-4125-9218-3b8addf001d5",
    "outputId": "35266e52-8299-4ccd-e195-3ecaba8996bc"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>db_id</th>\n",
       "      <th>query</th>\n",
       "      <th>question</th>\n",
       "      <th>query_toks</th>\n",
       "      <th>query_toks_no_value</th>\n",
       "      <th>question_toks</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>SELECT count(*) FROM singer</td>\n",
       "      <td>How many singers do we have?</td>\n",
       "      <td>[SELECT, count, (, *, ), FROM, singer]</td>\n",
       "      <td>[select, count, (, *, ), from, singer]</td>\n",
       "      <td>[How, many, singers, do, we, have, ?]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>SELECT count(*) FROM singer</td>\n",
       "      <td>What is the total number of singers?</td>\n",
       "      <td>[SELECT, count, (, *, ), FROM, singer]</td>\n",
       "      <td>[select, count, (, *, ), from, singer]</td>\n",
       "      <td>[What, is, the, total, number, of, singers, ?]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>SELECT name ,  country ,  age FROM singer ORDE...</td>\n",
       "      <td>Show name, country, age for all singers ordere...</td>\n",
       "      <td>[SELECT, name, ,, country, ,, age, FROM, singe...</td>\n",
       "      <td>[select, name, ,, country, ,, age, from, singe...</td>\n",
       "      <td>[Show, name, ,, country, ,, age, for, all, sin...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            db_id                                              query  \\\n",
       "0  concert_singer                        SELECT count(*) FROM singer   \n",
       "1  concert_singer                        SELECT count(*) FROM singer   \n",
       "2  concert_singer  SELECT name ,  country ,  age FROM singer ORDE...   \n",
       "\n",
       "                                            question  \\\n",
       "0                       How many singers do we have?   \n",
       "1               What is the total number of singers?   \n",
       "2  Show name, country, age for all singers ordere...   \n",
       "\n",
       "                                          query_toks  \\\n",
       "0             [SELECT, count, (, *, ), FROM, singer]   \n",
       "1             [SELECT, count, (, *, ), FROM, singer]   \n",
       "2  [SELECT, name, ,, country, ,, age, FROM, singe...   \n",
       "\n",
       "                                 query_toks_no_value  \\\n",
       "0             [select, count, (, *, ), from, singer]   \n",
       "1             [select, count, (, *, ), from, singer]   \n",
       "2  [select, name, ,, country, ,, age, from, singe...   \n",
       "\n",
       "                                       question_toks  \n",
       "0              [How, many, singers, do, we, have, ?]  \n",
       "1     [What, is, the, total, number, of, singers, ?]  \n",
       "2  [Show, name, ,, country, ,, age, for, all, sin...  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spider_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3038802b-b391-4941-b6e0-58b363834e1a",
   "metadata": {},
   "source": [
    "We now set a handful of rows aside, for testing the query generation. These questions come from two schemas in the Spider dataset. One schema contains tables about singers and where they’ve performed, and the other schema contains tables about pets and which students own those pets.\n",
    "\n",
    "The remaining rows will be used as candidate examples for the dynamic few-shot prompts later: importantly, to avoid data leakage and have a fair test we need to exclude from not only the test rows, but any additional row with the same query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecdbfa6e-0c1c-4d16-affe-008ed244e181",
   "metadata": {},
   "source": [
    "_Note: feel free to change the indices for the test set, but since the Spider dataset targets way more databases than the two we will actually populate with made-up data, keep the indices below ~80 to ensure only those two are targeted by the tests._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9cdd0af9-afb3-40b7-a2b5-4ff2571054e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 1034 rows loaded.\n",
      "Rows kept aside for testing: 6\n",
      "Rows to use as examples: 1023.\n",
      "  (5 further rows pruned from examples)\n"
     ]
    }
   ],
   "source": [
    "# picking some random questions for testing.\n",
    "test_indices = [25, 35, 45, 55, 65, 75]\n",
    "\n",
    "test_df = spider_df.loc[test_indices]\n",
    "test_queries = set(test_df[\"query\"])\n",
    "\n",
    "# Perform comparisons up to spacing and upper/lowercase differences:\n",
    "norm_test_queries = {qr.replace(\" \", \"\").lower() for qr in test_queries}\n",
    "\n",
    "def _is_test_query(qr: str) -> bool:\n",
    "    return qr.replace(\" \", \"\").lower() in norm_test_queries\n",
    "\n",
    "idx_to_remove = spider_df[spider_df[\"query\"].apply(_is_test_query)].index\n",
    "\n",
    "examples_df = spider_df.drop(idx_to_remove)\n",
    "\n",
    "print(\n",
    "    f\"Total {len(spider_df)} rows loaded.\\n\"\n",
    "    f\"Rows kept aside for testing: {len(test_queries)}\\n\"\n",
    "    f\"Rows to use as examples: {len(examples_df)}.\\n\"\n",
    "    f\"  ({len(idx_to_remove) - len(test_queries)} further rows pruned from examples)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168bffc1-a7d5-4b7d-92df-2ae9bb3c6a6c",
   "metadata": {},
   "source": [
    "This code prints the selected test questions and their Example SQL (the latter will be used only to evaluate the correctness of the SQL generation):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "98f2168b-e2e7-44b1-b7da-8151ee88094d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test question [0] (on DB concert_singer):\n",
      "    What is the name and capacity of the stadium with the most concerts after 2013 ?\n",
      "Corresponding Example SQL:\n",
      "    select t2.name ,  t2.capacity from concert as t1 join stadium as t2 on t1.stadium_id  =  t2.stadium_id where t1.year  >  2013 group by t2.stadium_id order by count(*) desc limit 1\n",
      "================================================================================\n",
      "Test question [1] (on DB concert_singer):\n",
      "    List singer names and number of concerts for each singer.\n",
      "Corresponding Example SQL:\n",
      "    SELECT T2.name ,  count(*) FROM singer_in_concert AS T1 JOIN singer AS T2 ON T1.singer_id  =  T2.singer_id GROUP BY T2.singer_id\n",
      "================================================================================\n",
      "Test question [2] (on DB pets_1):\n",
      "    Find the number of pets whose weight is heavier than 10.\n",
      "Corresponding Example SQL:\n",
      "    SELECT count(*) FROM pets WHERE weight  >  10\n",
      "================================================================================\n",
      "Test question [3] (on DB pets_1):\n",
      "    Find the number of distinct type of pets.\n",
      "Corresponding Example SQL:\n",
      "    SELECT count(DISTINCT pettype) FROM pets\n",
      "================================================================================\n",
      "Test question [4] (on DB pets_1):\n",
      "    Find the first name and age of students who have a dog but do not have a cat as a pet.\n",
      "Corresponding Example SQL:\n",
      "    SELECT T1.fname ,  T1.age FROM student AS T1 JOIN has_pet AS T2 ON T1.stuid  =  T2.stuid JOIN pets AS T3 ON T3.petid  =  T2.petid WHERE T3.pettype  =  'dog' AND T1.stuid NOT IN (SELECT T1.stuid FROM student AS T1 JOIN has_pet AS T2 ON T1.stuid  =  T2.stuid JOIN pets AS T3 ON T3.petid  =  T2.petid WHERE T3.pettype  =  'cat')\n",
      "================================================================================\n",
      "Test question [5] (on DB pets_1):\n",
      "    Find the first name and age of students who have a pet.\n",
      "Corresponding Example SQL:\n",
      "    SELECT DISTINCT T1.fname ,  T1.age FROM student AS T1 JOIN has_pet AS T2 ON T1.stuid  =  T2.stuid\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "test_queries = {\n",
    "    i: row\n",
    "    for i, (_, row) in enumerate(test_df.iterrows())\n",
    "}\n",
    "\n",
    "for i, row in test_queries.items():\n",
    "    if row[\"db_id\"] in {\"pets_1\", \"concert_singer\"}:\n",
    "        print(f\"Test question [{i}] (on DB {row['db_id']}):\\n    {row['question']}\")\n",
    "        print(f\"Corresponding Example SQL:\\n    {row['query']}\")\n",
    "        print(\"=\" * 80)\n",
    "    else:\n",
    "        raise ValueError(\n",
    "            f\"Test question detected targeting a db out of scope ({row['db_id']}). \"\n",
    "            \"Please lower the values in `test_indices` and try again.\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a00dee-0692-4f63-9248-0f0bd80b6279",
   "metadata": {
    "id": "44a00dee-0692-4f63-9248-0f0bd80b6279"
   },
   "source": [
    "## Populate database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da848ba-7452-4a15-b1b8-579e4ce7fa84",
   "metadata": {
    "id": "5da848ba-7452-4a15-b1b8-579e4ce7fa84"
   },
   "source": [
    "To insert example data in the database, the `execute_sql` function defined above will be used."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3d5050-26fa-4e9c-8b2e-863b3b876b48",
   "metadata": {},
   "source": [
    "Set up the `concert_singer` database. Create tables matching the Spider schema, and then insert sample data for each test question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fbfdb53c-dbb5-4809-a61e-f4f242d33574",
   "metadata": {
    "id": "fbfdb53c-dbb5-4809-a61e-f4f242d33574",
    "outputId": "8040964a-5a89-414c-bb0f-273dd348394d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set up concert_singer DB\n",
    "CREATE_TABLES_SQL = \"\"\"\n",
    "-- Creating the stadium table\n",
    "CREATE TABLE stadium (\n",
    "    Stadium_ID INT PRIMARY KEY,\n",
    "    Location TEXT,\n",
    "    Name TEXT,\n",
    "    Capacity INT,\n",
    "    Highest INT,\n",
    "    Lowest INT,\n",
    "    Average INT\n",
    ");\n",
    "\n",
    "-- Creating the singer table\n",
    "CREATE TABLE singer (\n",
    "    Singer_ID INT PRIMARY KEY,\n",
    "    Name TEXT,\n",
    "    Country TEXT,\n",
    "    Song_Name TEXT,\n",
    "    Song_release_year TEXT,\n",
    "    Age INT,\n",
    "    Is_male BOOLEAN\n",
    ");\n",
    "\n",
    "-- Creating the concert table\n",
    "CREATE TABLE concert (\n",
    "    concert_ID INT PRIMARY KEY,\n",
    "    concert_Name TEXT,\n",
    "    Theme TEXT,\n",
    "    Stadium_ID INT,\n",
    "    Year TEXT,\n",
    "    FOREIGN KEY (Stadium_ID) REFERENCES stadium(Stadium_ID)\n",
    ");\n",
    "\n",
    "-- Creating the singer_in_concert table\n",
    "CREATE TABLE singer_in_concert (\n",
    "    concert_ID INT,\n",
    "    Singer_ID INT,\n",
    "    PRIMARY KEY (concert_ID, Singer_ID),\n",
    "    FOREIGN KEY (concert_ID) REFERENCES concert(concert_ID),\n",
    "    FOREIGN KEY (Singer_ID) REFERENCES singer(Singer_ID)\n",
    ");\n",
    "\"\"\"\n",
    "\n",
    "POPULATE_DATA_SQL = \"\"\"\n",
    "-- Populating the stadium table\n",
    "INSERT INTO stadium (\n",
    "    Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average\n",
    ") VALUES\n",
    "(1, 'New York, USA', 'Liberty Stadium', 50000, 1000, 500, 750),\n",
    "(2, 'London, UK', 'Royal Arena', 60000, 1500, 600, 900),\n",
    "(3, 'Tokyo, Japan', 'Sunshine Dome', 55000, 1200, 550, 800),\n",
    "(4, 'Sydney, Australia', 'Ocean Field', 40000, 900, 400, 650),\n",
    "(5, 'Berlin, Germany', 'Eagle Grounds', 45000, 1100, 450, 700);\n",
    "\n",
    "-- Populating the singer table\n",
    "INSERT INTO singer (\n",
    "    Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male\n",
    ") VALUES\n",
    "(1, 'John Doe', 'USA', 'Freedom Song', '2018', 28, TRUE),\n",
    "(2, 'Emma Stone', 'UK', 'Rolling Hills', '2019', 25, FALSE),\n",
    "(3, 'Haruki Tanaka', 'Japan', 'Tokyo Lights', '2020', 30, TRUE),\n",
    "(4, 'Alice Johnson', 'Australia', 'Ocean Waves', '2021', 27, FALSE),\n",
    "(5, 'Max Müller', 'Germany', 'Berlin Nights', '2017', 32, TRUE);\n",
    "\n",
    "-- Populating the concert table\n",
    "INSERT INTO concert (concert_ID, concert_Name, Theme, Stadium_ID, Year) VALUES\n",
    "(1, 'Freedom Fest', 'Pop', 1, '2021'),\n",
    "(2, 'Rock Mania', 'Rock', 2, '2022'),\n",
    "(3, 'Electronic Waves', 'Electronic', 3, '2020'),\n",
    "(4, 'Jazz Evenings', 'Jazz', 3, '2019'),\n",
    "(5, 'Classical Mornings', 'Classical', 5, '2023');\n",
    "\n",
    "-- Populating the singer_in_concert table\n",
    "INSERT INTO singer_in_concert (concert_ID, Singer_ID) VALUES\n",
    "(1, 1),\n",
    "(1, 2),\n",
    "(2, 3),\n",
    "(3, 4),\n",
    "(4, 5),\n",
    "(5, 1),\n",
    "(2, 2),\n",
    "(3, 3),\n",
    "(4, 4),\n",
    "(5, 5);\n",
    "\"\"\"\n",
    "\n",
    "statements = [\n",
    "    statement.strip() for statement in (\n",
    "        CREATE_TABLES_SQL + POPULATE_DATA_SQL\n",
    "    ).split(\";\")\n",
    "    if len(statement.strip()) > 0\n",
    "]\n",
    "execute_sql(*statements)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc12c6d1-83fd-4e91-902c-05067f2b6f72",
   "metadata": {},
   "source": [
    "Set up the `pets_1` database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1c2b3c2d-d2c1-4da7-a473-f32c76a8036a",
   "metadata": {
    "id": "1b2a0be5-419a-4781-898d-bb87ef8f69de",
    "outputId": "1ecf9897-3a7e-472e-b251-b4ce88e1d6d6"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set up pets_1 DB\n",
    "CREATE_TABLES_SQL = \"\"\"\n",
    "-- Creating the Student table\n",
    "CREATE TABLE Student (\n",
    "    StuID INT PRIMARY KEY,\n",
    "    LName VARCHAR(255),\n",
    "    Fname VARCHAR(255),\n",
    "    Age INT,\n",
    "    Sex VARCHAR(10),\n",
    "    Major INT,\n",
    "    Advisor INT,\n",
    "    city_code VARCHAR(50)\n",
    ");\n",
    "\n",
    "-- Creating the Pets table\n",
    "CREATE TABLE Pets (\n",
    "    PetID INT PRIMARY KEY,\n",
    "    PetType VARCHAR(255),\n",
    "    pet_age INT,\n",
    "    weight INT\n",
    ");\n",
    "\n",
    "-- Creating the Has_Pet table\n",
    "CREATE TABLE Has_Pet (\n",
    "    StuID INT,\n",
    "    PetID INT,\n",
    "    FOREIGN KEY (StuID) REFERENCES Student(StuID),\n",
    "    FOREIGN KEY (PetID) REFERENCES Pets(PetID)\n",
    ");\n",
    "\"\"\"\n",
    "\n",
    "POPULATE_DATA_SQL = \"\"\"\n",
    "-- Populating the Student table\n",
    "INSERT INTO Student (\n",
    "    StuID, LName, Fname, Age, Sex, Major, Advisor, city_code\n",
    ") VALUES\n",
    "(101, 'Smith', 'John', 20, 'M', 501, 301, 'NYC'),\n",
    "(102, 'Johnson', 'Emma', 22, 'F', 502, 302, 'LAX'),\n",
    "(103, 'Williams', 'Michael', 21, 'M', 503, 303, 'CHI'),\n",
    "(104, 'Brown', 'Sarah', 23, 'F', 504, 304, 'HOU'),\n",
    "(105, 'Jones', 'David', 19, 'M', 505, 305, 'PHI');\n",
    "\n",
    "-- Populating the Pets table\n",
    "INSERT INTO Pets (PetID, PetType, pet_age, weight) VALUES\n",
    "(201, 'dog', 3, 20.5),\n",
    "(202, 'cat', 5, 10.2),\n",
    "(203, 'dog', 2, 8.1),\n",
    "(204, 'parrot', 4, 0.5),\n",
    "(205, 'hamster', 1, 0.7);\n",
    "\n",
    "-- Populating the Has_Pet table\n",
    "INSERT INTO Has_Pet (StuID, PetID) VALUES\n",
    "(101, 201),\n",
    "(101, 202),\n",
    "(105, 203),\n",
    "(103, 204),\n",
    "(104, 205),\n",
    "(105, 201);\n",
    "\"\"\"\n",
    "\n",
    "statements = [\n",
    "    statement.strip() for statement in (\n",
    "        CREATE_TABLES_SQL + POPULATE_DATA_SQL\n",
    "    ).split(\";\")\n",
    "    if len(statement.strip()) > 0\n",
    "]\n",
    "execute_sql(*statements)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1ec0d5f-3a34-4e03-808f-3fd0bb864172",
   "metadata": {},
   "source": [
    "Note that the example data, of the full Spider dataset, only covers the \"pets_1\" and \"concert_singer\" databases. For this reason, we had a check earlier to ensure the chosen test queries are all targeting those databases (lest the correctness check would trivially fail)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62b1a3ad-e8dc-4aa1-bd9e-31d2ee01805e",
   "metadata": {},
   "source": [
    "## Evaluator function for queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b4a2f0-a5d9-40d8-9929-5c047dafdca5",
   "metadata": {
    "id": "19b4a2f0-a5d9-40d8-9929-5c047dafdca5"
   },
   "source": [
    "Finally, we just need an evaluation function for the generated queries, to assess their correctness.\n",
    "\n",
    "This function runs the Example SQL counterpart alongside each of the generated queries,\n",
    "then compares the results to determine whether they return the same data.\n",
    "The evaluator must keep into account that results differing only in their ordering should be treated as returning the same data.\n",
    "\n",
    "The function returns a short report in the form of a table: what we want to maximize is the **accuracy**, the percentage of generated queries that do return the same data as the corresponding Example SQL .\n",
    "\n",
    "_Note: In a realistic application, the accuracy must be measured over a much larger amount of test questions than is shown in this demo._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7bde010d-5fb5-410b-8b66-a6376668c601",
   "metadata": {
    "id": "7bde010d-5fb5-410b-8b66-a6376668c601"
   },
   "outputs": [],
   "source": [
    "def eval_generated_queries(generated_queries: Dict[int, str]) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Evaluate the given queries against the test set,\n",
    "    and return a report of the performance on each row.\n",
    "    \"\"\"\n",
    "    report = []\n",
    "\n",
    "    for tq_index, tq_generated_sql in sorted(generated_queries.items()):\n",
    "        query_row = test_queries[tq_index]\n",
    "        example_sql = query_row[\"query\"]\n",
    "        example_results = execute_sql(example_sql)\n",
    "        \n",
    "        try:\n",
    "            gen_results = execute_sql(tq_generated_sql)\n",
    "            error = None\n",
    "\n",
    "            # Figure out correctness, not super straightforward\n",
    "            if len(example_results) == len(gen_results):\n",
    "                # (note: here and in the next sorting,\n",
    "                #  we ignore \"sorting collisions\" for simplicity...)\n",
    "                in_sorted_gen = [sorted(tpl, key=str) for tpl in gen_results]\n",
    "                in_sorted_example = [sorted(tpl, key=str) for tpl in example_results]\n",
    "                if \"ORDER BY\" in example_sql.upper():\n",
    "                    # ordering within tuples must match\n",
    "                    correct = in_sorted_gen == in_sorted_example\n",
    "                else:\n",
    "                    # sort sequence of tuples before comparison\\\n",
    "                    correct = sorted(in_sorted_gen) == sorted(in_sorted_example)\n",
    "            else:\n",
    "                correct = False\n",
    "\n",
    "        except Exception as e:\n",
    "            gen_results = None\n",
    "            error = e\n",
    "            correct = False\n",
    "\n",
    "        report.append({\n",
    "            \"DB ID\": query_row[\"db_id\"],\n",
    "            \"Question\": query_row[\"question\"],\n",
    "            \"Correct\": correct,\n",
    "            \"Error\": error,\n",
    "        })\n",
    "\n",
    "    return pd.DataFrame(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7e5429c-b76c-434d-ac8d-0e6ed433a611",
   "metadata": {
    "id": "d7e5429c-b76c-434d-ac8d-0e6ed433a611"
   },
   "source": [
    "## Generating SQL 1: Zero-Shot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb32bd7f-e958-49b3-8623-446bfad21529",
   "metadata": {
    "id": "fb32bd7f-e958-49b3-8623-446bfad21529"
   },
   "source": [
    "Now, let's go ahead and generate some SQL queries in a few different ways.\n",
    "\n",
    "First, we'll use an LLM with a prompt template from the [SQL-PaLM](https://arxiv.org/abs/2306.00739) paper as-is to see how it does."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cf9049a-6563-4e2b-a516-a9c7c42afde6",
   "metadata": {},
   "source": [
    "### SQL-generation tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "295c7cb4-5b4d-47bd-95fc-ff2767afd80d",
   "metadata": {
    "id": "295c7cb4-5b4d-47bd-95fc-ff2767afd80d"
   },
   "source": [
    "First, we have some utility functions to define.\n",
    "\n",
    "The main purpose of the tools below is to assemble a complete prompt for an LLM, urging it to produce SQL."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70ff71e4-9f6a-466b-bae0-dd517b90dfc0",
   "metadata": {},
   "source": [
    "#### Schema format converter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ccb22c-1572-4c51-ad5e-b42101965032",
   "metadata": {},
   "source": [
    "The logic below translates the Spider schema format to match the prompt template of the SQL-PaLM paper, which has proven to be effective."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "34de348c-1f1f-47db-8dd1-05ca397003b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _format_schema(\n",
    "    db_id: str, spider_schema_row: str\n",
    ") -> Tuple[str, str, str, str]:\n",
    "    \"\"\"\n",
    "    Converts the existing schema format of spider_schema_df to\n",
    "    (schema, columns, primary keys, foreign keys) as used by SQL-PaLM.\n",
    "    \"\"\"\n",
    "    schema_str = f\"| {db_id} \"\n",
    "    column_str = \"\"\n",
    "\n",
    "    spider_schema_vt_str = spider_schema_row[\"Schema (values (type))\"]\n",
    "\n",
    "    for table_name in re.findall(r\"([^,: ]*) :\", spider_schema_vt_str):\n",
    "        schema_str += f\"| {table_name} : \"\n",
    "\n",
    "        start_ndx = spider_schema_vt_str.find(table_name + \" :\")\n",
    "        end_ndx = spider_schema_vt_str.find(\n",
    "            \":\", start_ndx + len(table_name) + 4\n",
    "        )\n",
    "        if end_ndx == -1:\n",
    "            end_ndx = len(spider_schema_vt_str)\n",
    "        current_substr = spider_schema_vt_str[start_ndx:end_ndx]\n",
    "        for col_name, col_type in re.findall(\n",
    "            r\" ([^:,]*) \\(([^,]*)\\)\", current_substr\n",
    "        ):\n",
    "            schema_str += f\"{col_name} , \"\n",
    "            column_str += f\"{table_name} : {col_name} ({col_type}) | \"\n",
    "\n",
    "        schema_str = schema_str[:-2]\n",
    "    column_str = column_str[:-2]\n",
    "\n",
    "    return (\n",
    "        schema_str + \";\",\n",
    "        column_str + \";\",\n",
    "        spider_schema_row[\"Primary Keys\"],\n",
    "        spider_schema_row[\"Foreign Keys\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec7ff334-7071-4d0d-bc4b-dda8ebe50d45",
   "metadata": {},
   "source": [
    "For illustration purposes, this is the effect of this conversion on the pets database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bf8e93bd-eeae-456b-aae6-f6e931475780",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BEFORE CONVERSION:\n",
      "{\n",
      "  \"db_id\": \"pets_1\",\n",
      "  \"Schema (values (type))\": \"Student : StuID (number) , LName (text) , Fname (text) , Age (number) , Sex (text) , Major (number) , Advisor (number) , city_code (text) | Has_Pet : StuID (number) , PetID (number) | Pets : PetID (number) , PetType (text) , pet_age (number) , weight (number)\",\n",
      "  \"Primary Keys\": \"Student : StuID | Pets : PetID\",\n",
      "  \"Foreign Keys\": \"Has_Pet : StuID equals Student : StuID | Has_Pet : PetID equals Pets : PetID\"\n",
      "}\n",
      "\n",
      "======\n",
      "AFTER CONVERSION:\n",
      "- SCHEMA:\n",
      "      | pets_1 | Student : StuID , LName , Fname , Age , Sex , Major , Advisor , city_code | Has_Pet : StuID , PetID | Pets : PetID , PetType , pet_age , weight ;\n",
      "- COLS:\n",
      "      Student : StuID (number) | Student : LName (text) | Student : Fname (text) | Student : Age (number) | Student : Sex (text) | Student : Major (number) | Student : Advisor (number) | Student : city_code (text) | Has_Pet : StuID (number) | Has_Pet : PetID (number) | Pets : PetID (number) | Pets : PetType (text) | Pets : pet_age (number) | Pets : weight (number) ;\n",
      "- PR_KEYS:\n",
      "      Student : StuID | Pets : PetID\n",
      "- F_KEYS:\n",
      "      Has_Pet : StuID equals Student : StuID | Has_Pet : PetID equals Pets : PetID\n"
     ]
    }
   ],
   "source": [
    "spider_schema_row0 = spider_schema_df[spider_schema_df[\"db_id\"]==\"pets_1\"].iloc[0]\n",
    "spider_schema0 = spider_schema_row0\n",
    "palm_schema0, palm_cols0, palm_pk0, palm_fk0 = _format_schema(\"pets_1\", spider_schema0)\n",
    "\n",
    "print(f\"BEFORE CONVERSION:\\n{json.dumps(spider_schema0.to_dict(), indent=2)}\")\n",
    "print(f\"\\n======\\nAFTER CONVERSION:\\n- SCHEMA:\\n      {palm_schema0}\")\n",
    "print(f\"- COLS:\\n      {palm_cols0}\")\n",
    "\n",
    "print(f\"- PR_KEYS:\\n      {palm_pk0}\")\n",
    "print(f\"- F_KEYS:\\n      {palm_fk0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b327a70-1c35-49b0-b997-8e18ceaf19c0",
   "metadata": {},
   "source": [
    "#### Assemble prompt for SQL generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dab167a7-36e8-4c97-bf99-3f26da26c46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "SQL_PROMPT_TEMPLATE = \"\"\"Convert text to SQL:\n",
    "\n",
    "[Schema : (values)]: {schema_str}\n",
    "\n",
    "[Column names (type)]: {column_str}\n",
    "\n",
    "[Primary Keys]: {pk_str}\n",
    "\n",
    "[Foreign Keys]: {fk_str}\n",
    "\n",
    "[Q]: {question}\n",
    "\n",
    "[SQL]: \"\"\"\n",
    "\n",
    "QUESTION_PREFIX_STR = (\n",
    "    \"Given the following schema information, generate valid SQL \"\n",
    "    \"to answer the provided query. Enclose the query in markdown \"\n",
    "    \"code-block syntax.\\n\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5d90bb54-bc50-47b2-a9bb-27cefc273cd7",
   "metadata": {
    "id": "5d90bb54-bc50-47b2-a9bb-27cefc273cd7"
   },
   "outputs": [],
   "source": [
    "@lru_cache\n",
    "def _get_spider_schema_by_db_id(db_id: str) -> pd.Series:\n",
    "    return spider_schema_df[spider_schema_df[\"db_id\"] == db_id].iloc[0]\n",
    "\n",
    "\n",
    "def _format_sql_prompt(db_id: str, question: str) -> str:\n",
    "    \"\"\"\n",
    "    Returns a formatted section of the prompt describing the DB Schema.\n",
    "    This core logic is factored for later re-use in the few-shot approach.\n",
    "    \"\"\"\n",
    "    spider_schema_row = _get_spider_schema_by_db_id(db_id)\n",
    "    schema_str, column_str, pk_str, fk_str = _format_schema(\n",
    "        db_id, spider_schema_row\n",
    "    )\n",
    "    sql_prompt_str = SQL_PROMPT_TEMPLATE.format(\n",
    "        schema_str=schema_str,\n",
    "        column_str=column_str,\n",
    "        pk_str=pk_str,\n",
    "        fk_str=fk_str,\n",
    "        question=question,\n",
    "    )\n",
    "    return sql_prompt_str\n",
    "\n",
    "\n",
    "def _format_question_prompt(db_id: str, question: str) -> str:\n",
    "    core_prompt = _format_sql_prompt(db_id, question)\n",
    "    return QUESTION_PREFIX_STR + core_prompt\n",
    "\n",
    "\n",
    "def _clean_sql_response(response: str) -> str:\n",
    "    _response2 = response.replace(\"```sql\", \"```\")\n",
    "    _start = _response2.find(\"```\")\n",
    "    if _start < 0:\n",
    "        # assume the whole response is the SQL, no markdown stuff\n",
    "        # (gpt-3.5 tends to do this)\n",
    "        return \"\\n\".join(l.strip() for l in _response2.split(\"\\n\") if l.strip())\n",
    "    rest = _response2[_start+3:]\n",
    "    _stop = rest.find(\"```\")\n",
    "    if _stop < 0:\n",
    "        raise ValueError(\"Invalid answer from LLM.\")\n",
    "    _sql = rest[:_stop].strip()\n",
    "    return _sql\n",
    "\n",
    "\n",
    "def zero_shot_prompt(db_id: str, question: str) -> str:\n",
    "    \"\"\"Creates a Zero-Shot prompt for the given question & db_id\"\"\"\n",
    "    return _format_question_prompt(db_id, question)\n",
    "\n",
    "\n",
    "def generate_sql(\n",
    "    question: str,\n",
    "    db_id: str,\n",
    "    prompt_fn: Callable[[str, str], str],\n",
    "    debug_prompt: bool = False,\n",
    ") -> str:\n",
    "    \"\"\"Use the LLM to generate a SQL response for a given Spider question.\"\"\"\n",
    "    prompt = prompt_fn(db_id=db_id, question=question)\n",
    "\n",
    "    if debug_prompt:\n",
    "        print(f\"LLM prompt for SQL generation:\\n======\\n{prompt}\\n======\")\n",
    "\n",
    "    response = llm_completion(prompt)\n",
    "    return _clean_sql_response(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fbc9748-cf18-4fa5-a793-f2903e53791b",
   "metadata": {
    "id": "6fbc9748-cf18-4fa5-a793-f2903e53791b"
   },
   "source": [
    "Next, we can try with the questions we pulled out from the examples above.\n",
    "\n",
    "Check out the full prompt to the LLM used to generate the query:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f56e52b2-6e4e-4437-a703-53fe6eb169ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_queries_zeroshot = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1f1242da-5819-4454-9292-63ebec3b4929",
   "metadata": {
    "id": "1f1242da-5819-4454-9292-63ebec3b4929",
    "outputId": "d6b8fcb0-bc35-47a5-e6a7-4fe16532d770"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM prompt for SQL generation:\n",
      "======\n",
      "Given the following schema information, generate valid SQL to answer the provided query. Enclose the query in markdown code-block syntax.\n",
      "Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID ;\n",
      "\n",
      "[Column names (type)]: stadium : Stadium_ID (number) | stadium : Location (text) | stadium : Name (text) | stadium : Capacity (number) | stadium : Highest (number) | stadium : Lowest (number) | stadium : Average (number) | singer : Singer_ID (number) | singer : Name (text) | singer : Country (text) | singer : Song_Name (text) | singer : Song_release_year (text) | singer : Age (number) | singer : Is_male (others) | concert : concert_ID (number) | concert : concert_Name (text) | concert : Theme (text) | concert : Stadium_ID (text) | concert : Year (text) | singer_in_concert : concert_ID (number) | singer_in_concert : Singer_ID (text) ;\n",
      "\n",
      "[Primary Keys]: stadium : Stadium_ID | singer : Singer_ID | concert : concert_ID | singer_in_concert : concert_ID\n",
      "\n",
      "[Foreign Keys]: concert : Stadium_ID equals stadium : Stadium_ID | singer_in_concert : Singer_ID equals singer : Singer_ID | singer_in_concert : concert_ID equals concert : concert_ID\n",
      "\n",
      "[Q]: What is the name and capacity of the stadium with the most concerts after 2013 ?\n",
      "\n",
      "[SQL]: \n",
      "======\n",
      "\n",
      "===\n",
      "Generated SQL:\n",
      "SELECT s.Name, s.Capacity\n",
      "FROM stadium s\n",
      "WHERE s.Stadium_ID IN (\n",
      "    SELECT c.Stadium_ID\n",
      "    FROM concert c\n",
      "    WHERE c.Year > '2013'\n",
      "    GROUP BY c.Stadium_ID\n",
      "    ORDER BY COUNT(c.concert_ID) DESC\n",
      "    LIMIT 1\n",
      ");\n"
     ]
    }
   ],
   "source": [
    "question0 = test_queries[0][\"question\"]\n",
    "db_id0 = test_queries[0][\"db_id\"]\n",
    "# Debug mode to inspect the complete prompt:\n",
    "gen_sql0 = generate_sql(\n",
    "    question=question0,\n",
    "    db_id=db_id0,\n",
    "    prompt_fn=zero_shot_prompt,\n",
    "    debug_prompt=True,\n",
    ")\n",
    "\n",
    "print(\"\\n===\\nGenerated SQL:\")\n",
    "print(gen_sql0)\n",
    "\n",
    "generated_queries_zeroshot[0] = gen_sql0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76651e92-8184-474f-95ec-e37871eceb7c",
   "metadata": {},
   "source": [
    "The other test questions can be similarly generated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4d140c20-0d3e-4216-8e4f-05bd1d49239d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "====\n",
      "Added for Q[1]='List singer names and number of concerts for each singer.' -> SQL:\n",
      "----\n",
      "SELECT singer.Name, COUNT(concert_ID) AS Number_of_Concerts\n",
      "FROM singer\n",
      "JOIN singer_in_concert ON singer.Singer_ID = singer_in_concert.Singer_ID\n",
      "GROUP BY singer.Name;\n",
      "\n",
      "====\n",
      "Added for Q[2]='Find the number of pets whose weight is heavier than 10.' -> SQL:\n",
      "----\n",
      "SELECT COUNT(*)\n",
      "FROM Pets\n",
      "WHERE weight > 10;\n",
      "\n",
      "====\n",
      "Added for Q[3]='Find the number of distinct type of pets.' -> SQL:\n",
      "----\n",
      "```\n",
      "SELECT COUNT(DISTINCT PetType) AS num_distinct_pet_types\n",
      "FROM Pets;\n",
      "```\n",
      "\n",
      "====\n",
      "Added for Q[4]='Find the first name and age of students who have a dog but do not have a cat as a pet.' -> SQL:\n",
      "----\n",
      "SELECT Fname, Age\n",
      "FROM Student\n",
      "WHERE StuID IN (\n",
      "    SELECT StuID\n",
      "    FROM Has_Pet\n",
      "    JOIN Pets ON Has_Pet.PetID = Pets.PetID\n",
      "    WHERE PetType = 'dog'\n",
      ")\n",
      "AND StuID NOT IN (\n",
      "    SELECT StuID\n",
      "    FROM Has_Pet\n",
      "    JOIN Pets ON Has_Pet.PetID = Pets.PetID\n",
      "    WHERE PetType = 'cat'\n",
      ");\n",
      "\n",
      "====\n",
      "Added for Q[5]='Find the first name and age of students who have a pet.' -> SQL:\n",
      "----\n",
      "SELECT Fname, Age\n",
      "FROM Student\n",
      "WHERE StuID IN (\n",
      "    SELECT StuID\n",
      "    FROM Has_Pet\n",
      ");\n"
     ]
    }
   ],
   "source": [
    "for test_i, test_row in test_queries.items():\n",
    "    if test_i not in generated_queries_zeroshot:\n",
    "        generated_queries_zeroshot[test_i] =  generate_sql(\n",
    "            question=test_row[\"question\"],\n",
    "            db_id=test_row[\"db_id\"],\n",
    "            prompt_fn=zero_shot_prompt,\n",
    "        )\n",
    "        print(\n",
    "            f\"\\n====\\nAdded for Q[{test_i}]=\"\n",
    "            f\"'{test_row['question']}' -> SQL:\\n----\"\n",
    "        )\n",
    "        print(generated_queries_zeroshot[test_i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e44223e0-85a9-4625-a115-32b8db9c1f95",
   "metadata": {},
   "source": [
    "Running the evaluation on these queries yields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "10deb49a-b8d0-4b9a-99d8-eff10bbfbb25",
   "metadata": {
    "id": "10deb49a-b8d0-4b9a-99d8-eff10bbfbb25",
    "outputId": "bc6acc06-7000-481f-af17-d1d41f4aa981"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>DB ID</th>\n",
       "      <th>Question</th>\n",
       "      <th>Correct</th>\n",
       "      <th>Error</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>What is the name and capacity of the stadium w...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>List singer names and number of concerts for e...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of pets whose weight is heavie...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of distinct type of pets.</td>\n",
       "      <td>False</td>\n",
       "      <td>near \"SELECT COUNT(DISTINCT PetType) AS n...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            DB ID                                           Question  Correct  \\\n",
       "0  concert_singer  What is the name and capacity of the stadium w...     True   \n",
       "1  concert_singer  List singer names and number of concerts for e...     True   \n",
       "2          pets_1  Find the number of pets whose weight is heavie...     True   \n",
       "3          pets_1          Find the number of distinct type of pets.    False   \n",
       "4          pets_1  Find the first name and age of students who ha...    False   \n",
       "5          pets_1  Find the first name and age of students who ha...     True   \n",
       "\n",
       "                                               Error  \n",
       "0                                               None  \n",
       "1                                               None  \n",
       "2                                               None  \n",
       "3  near \"```\\nSELECT COUNT(DISTINCT PetType) AS n...  \n",
       "4                                               None  \n",
       "5                                               None  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "report = eval_generated_queries(generated_queries_zeroshot)\n",
    "display(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66300234-b579-4548-8eb5-c18f414c3a48",
   "metadata": {
    "id": "66300234-b579-4548-8eb5-c18f414c3a48"
   },
   "source": [
    "So, **most questions get a working, correct answer** already.\n",
    "\n",
    "The next step is to see if the performance can be improved with a few-shot approach, i.e. by passing some examples to the SQL generation prompt, in order to better \"guide\" the LLM to the answer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7969b0d9-7267-40c1-bbbb-05b7953d4312",
   "metadata": {
    "id": "7969b0d9-7267-40c1-bbbb-05b7953d4312"
   },
   "source": [
    "## Generating SQL 2: Few-Shot (static)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d49c0640-db19-4d5c-9db7-2008998c0a45",
   "metadata": {
    "id": "d49c0640-db19-4d5c-9db7-2008998c0a45"
   },
   "source": [
    "Few-Shot prompting has been a technique for a few years, since LLMs were first conceived. The theory behind few-shot prompting (also known as \"In-Context Learning\") is that we are showing the model examples of our desired input/output pairs, and guiding the model to make decisions more in line with what we expect for the task at hand. Since the original GPT-3 paper, [Language Models are Few-Shot Learners](https://arxiv.org/abs/2005.14165), few-shot methods have been shown to greatly improve performance over zero-shot methods (aka no examples in the prompt). Because of this, let's see how our model does if we give it a few examples of correct output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ce2960f1-a4b3-4bcf-9392-1cace2f4a5ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXAMPLE_PREFIX_STR = \"Here is an example: \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8eaa926a-3f5f-4253-be8b-4b2bbbf94762",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _format_example_prompt(db_id: str, question: str, example_sql: str) -> str:\n",
    "    core_prompt = _format_sql_prompt(db_id, question)\n",
    "    return EXAMPLE_PREFIX_STR + core_prompt + example_sql + \"\\n\\n\"\n",
    "\n",
    "\n",
    "def few_shot_prompt(\n",
    "    db_id: str, question: str, example_indices: List[int]\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    Create a few shot prompt using the given indices to construct the example.\n",
    "    \"\"\"\n",
    "    prompt = \"\"\n",
    "    for e_index in example_indices:\n",
    "        e_row = examples_df.loc[e_index]\n",
    "        prompt += _format_example_prompt(\n",
    "            e_row[\"db_id\"], e_row[\"question\"], e_row[\"query\"]\n",
    "        )\n",
    "    return prompt + _format_question_prompt(db_id, question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "27675c80-926c-4aa4-84e6-1641b88a3c1b",
   "metadata": {
    "id": "27675c80-926c-4aa4-84e6-1641b88a3c1b"
   },
   "outputs": [],
   "source": [
    "static_few_shot_prompt_fn = partial(few_shot_prompt, example_indices=[100, 200])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcfe2f65-fa94-4404-a382-4702b516b117",
   "metadata": {
    "id": "dcfe2f65-fa94-4404-a382-4702b516b117"
   },
   "source": [
    "Note: one of the biggest tradeoffs with few-shot learning is the context length. We end up sacrificing context-length (and thus increasing costs, latency, and decreasing the space left for the generation) for each example shown to the model. For examples where the input context is very long (like showing schemas for SQL queries), this can be a difficult tradeoff to make. In our scenario, we will settle for just showing 2 examples to the model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f6bda4a-98b2-428a-87b9-3c5defdfabd4",
   "metadata": {},
   "source": [
    "Now, again a query is generated for each of the test question. For the first one, take a look at the full prompt with its examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b143a80f-726f-4669-bed0-9a98a2c5412a",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_queries_fewshot = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "eed50d04-daf0-487a-963d-2bd064b84667",
   "metadata": {
    "id": "eed50d04-daf0-487a-963d-2bd064b84667",
    "outputId": "dba1e25a-ca31-46ad-d980-32c41878d08b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM prompt for SQL generation:\n",
      "======\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | car_1 | continents : ContId , Continent | countries : CountryId , CountryName , Continent | car_makers : Id , Maker , FullName , Country | model_list : ModelId , Maker , Model | car_names : MakeId , Model , Make | cars_data : Id , MPG , Cylinders , Edispl , Horsepower , Weight , Accelerate , Year ;\n",
      "\n",
      "[Column names (type)]: continents : ContId (number) | continents : Continent (text) | countries : CountryId (number) | countries : CountryName (text) | countries : Continent (number) | car_makers : Id (number) | car_makers : Maker (text) | car_makers : FullName (text) | car_makers : Country (text) | model_list : ModelId (number) | model_list : Maker (number) | model_list : Model (text) | car_names : MakeId (number) | car_names : Model (text) | car_names : Make (text) | cars_data : Id (number) | cars_data : MPG (text) | cars_data : Cylinders (number) | cars_data : Edispl (number) | cars_data : Horsepower (text) | cars_data : Weight (number) | cars_data : Accelerate (number) | cars_data : Year (number) ;\n",
      "\n",
      "[Primary Keys]: continents : ContId | countries : CountryId | car_makers : Id | model_list : ModelId | car_names : MakeId | cars_data : Id\n",
      "\n",
      "[Foreign Keys]: countries : Continent equals continents : ContId | car_makers : Country equals countries : CountryId | model_list : Maker equals car_makers : Id | car_names : Model equals model_list : Model | cars_data : Id equals car_names : MakeId\n",
      "\n",
      "[Q]: What is the name of the different car makers who produced a car in 1970?\n",
      "\n",
      "[SQL]: SELECT DISTINCT T1.Maker FROM CAR_MAKERS AS T1 JOIN MODEL_LIST AS T2 ON T1.Id  =  T2.Maker JOIN CAR_NAMES AS T3 ON T2.model  =  T3.model JOIN CARS_DATA AS T4 ON T3.MakeId  =  T4.id WHERE T4.year  =  '1970';\n",
      "\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | flight_2 | airlines : uid , Airline , Abbreviation , Country | airports : City , AirportCode , AirportName , Country , CountryAbbrev | flights : Airline , FlightNo , SourceAirport , DestAirport ;\n",
      "\n",
      "[Column names (type)]: airlines : uid (number) | airlines : Airline (text) | airlines : Abbreviation (text) | airlines : Country (text) | airports : City (text) | airports : AirportCode (text) | airports : AirportName (text) | airports : Country (text) | airports : CountryAbbrev (text) | flights : Airline (number) | flights : FlightNo (number) | flights : SourceAirport (text) | flights : DestAirport (text) ;\n",
      "\n",
      "[Primary Keys]: airlines : uid | airports : AirportCode | flights : Airline\n",
      "\n",
      "[Foreign Keys]: flights : DestAirport equals airports : AirportCode | flights : SourceAirport equals airports : AirportCode\n",
      "\n",
      "[Q]: Return the name of the airport with code 'AKO'.\n",
      "\n",
      "[SQL]: SELECT AirportName FROM AIRPORTS WHERE AirportCode  =  \"AKO\"\n",
      "\n",
      "Given the following schema information, generate valid SQL to answer the provided query. Enclose the query in markdown code-block syntax.\n",
      "Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID ;\n",
      "\n",
      "[Column names (type)]: stadium : Stadium_ID (number) | stadium : Location (text) | stadium : Name (text) | stadium : Capacity (number) | stadium : Highest (number) | stadium : Lowest (number) | stadium : Average (number) | singer : Singer_ID (number) | singer : Name (text) | singer : Country (text) | singer : Song_Name (text) | singer : Song_release_year (text) | singer : Age (number) | singer : Is_male (others) | concert : concert_ID (number) | concert : concert_Name (text) | concert : Theme (text) | concert : Stadium_ID (text) | concert : Year (text) | singer_in_concert : concert_ID (number) | singer_in_concert : Singer_ID (text) ;\n",
      "\n",
      "[Primary Keys]: stadium : Stadium_ID | singer : Singer_ID | concert : concert_ID | singer_in_concert : concert_ID\n",
      "\n",
      "[Foreign Keys]: concert : Stadium_ID equals stadium : Stadium_ID | singer_in_concert : Singer_ID equals singer : Singer_ID | singer_in_concert : concert_ID equals concert : concert_ID\n",
      "\n",
      "[Q]: What is the name and capacity of the stadium with the most concerts after 2013 ?\n",
      "\n",
      "[SQL]: \n",
      "======\n",
      "\n",
      "===\n",
      "Generated SQL:\n",
      "SELECT Name, Capacity\n",
      "FROM stadium\n",
      "WHERE Stadium_ID IN (SELECT Stadium_ID\n",
      "FROM concert\n",
      "WHERE Year > '2013'\n",
      "GROUP BY Stadium_ID\n",
      "ORDER BY COUNT(concert_ID) DESC\n",
      "LIMIT 1);\n"
     ]
    }
   ],
   "source": [
    "question0 = test_queries[0][\"question\"]\n",
    "db_id0 = test_queries[0][\"db_id\"]\n",
    "# Debug mode to inspect the complete prompt:\n",
    "gen_sql0 = generate_sql(\n",
    "    question=question0,\n",
    "    db_id=db_id0,\n",
    "    prompt_fn=static_few_shot_prompt_fn,\n",
    "    debug_prompt=True,\n",
    ")\n",
    "\n",
    "print(\"\\n===\\nGenerated SQL:\")\n",
    "print(gen_sql0)\n",
    "\n",
    "generated_queries_fewshot[0] = gen_sql0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a33a2476-aae3-4a25-a7a2-c744fb877f32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "====\n",
      "Added for Q[1]='List singer names and number of concerts for each singer.' -> SQL:\n",
      "----\n",
      "SELECT singer.Name, COUNT(singer_in_concert.Singer_ID) AS Num_Concerts\n",
      "FROM singer\n",
      "JOIN singer_in_concert ON singer.Singer_ID = singer_in_concert.Singer_ID\n",
      "GROUP BY singer.Name;\n",
      "\n",
      "====\n",
      "Added for Q[2]='Find the number of pets whose weight is heavier than 10.' -> SQL:\n",
      "----\n",
      "SELECT COUNT(*) FROM PETS WHERE weight > 10;\n",
      "\n",
      "====\n",
      "Added for Q[3]='Find the number of distinct type of pets.' -> SQL:\n",
      "----\n",
      "SELECT COUNT(DISTINCT PetType) AS NumDistinctPets FROM PETS;\n",
      "\n",
      "====\n",
      "Added for Q[4]='Find the first name and age of students who have a dog but do not have a cat as a pet.' -> SQL:\n",
      "----\n",
      "SELECT Fname, Age\n",
      "FROM Student\n",
      "WHERE StuID IN\n",
      "(SELECT StuID\n",
      "FROM Has_Pet\n",
      "INNER JOIN Pets ON Has_Pet.PetID = Pets.PetID\n",
      "WHERE PetType = 'dog'\n",
      "AND StuID NOT IN\n",
      "(SELECT StuID\n",
      "FROM Has_Pet\n",
      "INNER JOIN Pets ON Has_Pet.PetID = Pets.PetID\n",
      "WHERE PetType = 'cat'))\n",
      "\n",
      "====\n",
      "Added for Q[5]='Find the first name and age of students who have a pet.' -> SQL:\n",
      "----\n",
      "SELECT Fname, Age\n",
      "FROM Student\n",
      "WHERE StuID IN (SELECT StuID FROM Has_Pet);\n"
     ]
    }
   ],
   "source": [
    "for test_i, test_row in test_queries.items():\n",
    "    if test_i not in generated_queries_fewshot:\n",
    "        generated_queries_fewshot[test_i] =  generate_sql(\n",
    "            question=test_row[\"question\"],\n",
    "            db_id=test_row[\"db_id\"],\n",
    "            prompt_fn=static_few_shot_prompt_fn,\n",
    "        )\n",
    "        print(\n",
    "            f\"\\n====\\nAdded for Q[{test_i}]=\"\n",
    "            f\"'{test_row['question']}' -> SQL:\\n----\"\n",
    "        )\n",
    "        print(generated_queries_fewshot[test_i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecb458a9-9148-4551-82bb-26069e10b923",
   "metadata": {},
   "source": [
    "Similarly as for the zero-shot approach, let's run the query evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c2b9038f-62a7-451b-88ee-71309677319a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>DB ID</th>\n",
       "      <th>Question</th>\n",
       "      <th>Correct</th>\n",
       "      <th>Error</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>What is the name and capacity of the stadium w...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>List singer names and number of concerts for e...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of pets whose weight is heavie...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of distinct type of pets.</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            DB ID                                           Question  Correct  \\\n",
       "0  concert_singer  What is the name and capacity of the stadium w...     True   \n",
       "1  concert_singer  List singer names and number of concerts for e...     True   \n",
       "2          pets_1  Find the number of pets whose weight is heavie...     True   \n",
       "3          pets_1          Find the number of distinct type of pets.     True   \n",
       "4          pets_1  Find the first name and age of students who ha...    False   \n",
       "5          pets_1  Find the first name and age of students who ha...     True   \n",
       "\n",
       "  Error  \n",
       "0  None  \n",
       "1  None  \n",
       "2  None  \n",
       "3  None  \n",
       "4  None  \n",
       "5  None  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "report = eval_generated_queries(generated_queries_fewshot)\n",
    "display(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cf5b64-0aa3-4358-9379-efc2788c3656",
   "metadata": {},
   "source": [
    "Few-shot generation **usually improves the situation a little**.\n",
    "\n",
    "Because the examples stuffed in the prompt -- always the same regardless of the question -- may have little relation to the structure of the actual problem at hand, we cannot expect this approach to yield big gains on the long run.\n",
    "\n",
    "The \"dynamic\" few-shot strategy, instead, will make a difference."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c21e36b-15e3-471e-93b1-906873e476a2",
   "metadata": {
    "id": "1c21e36b-15e3-471e-93b1-906873e476a2"
   },
   "source": [
    "## Generating SQL 3: Dynamic Few-Shot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a0f81a-c730-4220-a41f-2a9b3d9cd023",
   "metadata": {
    "id": "c6a0f81a-c730-4220-a41f-2a9b3d9cd023"
   },
   "source": [
    "The methods seen so far, zero-shot prompting and few-shot prompting with fixed examples, are a good starting point for a text-to-SQL implementation. However, the accuracy can be boosted with the introduction of dynamic few-shot prompting with Astra DB Serverless.\n",
    "\n",
    "To implement dynamic few-shot prompting in this case, we need to pick the most similar known question/query pairs from the examples dataset and insert them as additional context into the SQL generation prompt.\n",
    "\n",
    "We will compute **vector embeddings** for the example questions and store them, together with the question/query pair itself, in an Astra DB collection: in this way, vector search enables locating the best tailored examples to enrich the SQL-generation prompt."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f1ed973-ada7-4830-a219-3ad3123de8fe",
   "metadata": {},
   "source": [
    "### Create an Astra DB collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0c94bd4e-3055-4de2-b149-cd5a2b1dc6fa",
   "metadata": {
    "id": "0c94bd4e-3055-4de2-b149-cd5a2b1dc6fa",
    "outputId": "c0e2db06-2f5c-4ae9-b45d-a18bb929b9bc"
   },
   "outputs": [],
   "source": [
    "astra_db_collection = astra_db.create_collection(\n",
    "    ASTRA_DB_COLLECTION_NAME,\n",
    "    definition=(\n",
    "        CollectionDefinition.builder()\n",
    "        .set_vector_dimension(1536)\n",
    "        .build()\n",
    "    ),\n",
    "    # this allows `ndarray` objects in documents to insert:\n",
    "    spawn_api_options=APIOptions(\n",
    "        serdes_options=SerdesOptions(\n",
    "            unroll_iterables_to_lists=True,\n",
    "        ),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5328bdc-3f61-406c-877d-5fdb2159486b",
   "metadata": {},
   "source": [
    "### Calculate an embedding vector for each example question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8cbb6eba-1b71-4ee8-b09f-84e178041a11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32fb69abf7a24102879a024a34f197b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/26 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Batching these requests speeds the calculation up by a factor of about 13\n",
    "EMBEDDING_BATCH_SIZE = 40\n",
    "\n",
    "tqdm.pandas()\n",
    "\n",
    "examples_embeddings = (\n",
    "    examples_df\n",
    "    .groupby(examples_df.index // EMBEDDING_BATCH_SIZE)[\"question\"]\n",
    "    .progress_transform(lambda questions: compute_embedding_list(questions.tolist()))\n",
    ")\n",
    "examples_df[\"question_embedding\"] = examples_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5782cc1e-2315-43b9-bb69-4179bb2f261e",
   "metadata": {},
   "source": [
    "### Write the documents (with embeddings) into the collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "e68a0c48-fc1e-4377-930c-f1e87b0e6f35",
   "metadata": {
    "colab": {
     "referenced_widgets": [
      "0ec19321799344bebc7be71f312ab48a"
     ]
    },
    "id": "e68a0c48-fc1e-4377-930c-f1e87b0e6f35",
    "outputId": "8ea01b7f-3693-4e6a-d259-6f6c30f43042"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CollectionInsertManyResult(inserted_ids=[0, 1, 2, 3, 4 ... (1023 total)], raw_results=...)\n"
     ]
    }
   ],
   "source": [
    "# The documents' _id in the collection will be the indices in the DataFrame:\n",
    "examples_df[\"_id\"] = examples_df.index\n",
    "# The embeddings must be placed in the special \"$vector\" key:\n",
    "examples_documents = examples_df.rename(\n",
    "    columns={\"question_embedding\": \"$vector\"}\n",
    ").to_dict(orient=\"records\")\n",
    "\n",
    "result = astra_db_collection.insert_many(\n",
    "    examples_documents,\n",
    "    # these are ~1000 big documents, so let's be overzealous with timeouts:\n",
    "    timeout_ms=45000,          # for the whole insertion (which is chunked) ...\n",
    "    request_timeout_ms=20000,  # and this is for each HTTP request (one chunk).\n",
    ")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03491892-e130-417a-8d46-02ba078c64fd",
   "metadata": {},
   "source": [
    "### Construct the dynamic few-shot prompting function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebd29db8-a254-4a13-8aa1-74ada5ce87f7",
   "metadata": {},
   "source": [
    "The new prompt function first runs a vector search on the Astra DB collection, then uses the results to prepare a few-shot prompt tailored to the question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "4a69e0d7-116a-486c-ba9b-ec0ccd0b900d",
   "metadata": {
    "id": "4a69e0d7-116a-486c-ba9b-ec0ccd0b900d"
   },
   "outputs": [],
   "source": [
    "def dynamic_few_shot_prompt_fn(db_id: str, question: str) -> str:\n",
    "    embedding = compute_embedding(question)\n",
    "    documents = astra_db_collection.find(sort={\"$vector\": embedding}, limit=2)\n",
    "    closest_q_ids = [\n",
    "        doc[\"_id\"] for doc in documents\n",
    "    ]\n",
    "    return few_shot_prompt(db_id, question, example_indices=closest_q_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab97ee8e-c7c5-4f69-97de-7c39ad90739b",
   "metadata": {},
   "source": [
    "Let's put this new generation technique to use. In the next two cells, the whole prompt is also printed: you can verify the examples differ in the two invocations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5e9bc523-c6be-48ba-b68a-4554010630bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_queries_dynamic = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "beb5f7a9-40e5-4e4e-827d-f6c213a6ffc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM prompt for SQL generation:\n",
      "======\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID ;\n",
      "\n",
      "[Column names (type)]: stadium : Stadium_ID (number) | stadium : Location (text) | stadium : Name (text) | stadium : Capacity (number) | stadium : Highest (number) | stadium : Lowest (number) | stadium : Average (number) | singer : Singer_ID (number) | singer : Name (text) | singer : Country (text) | singer : Song_Name (text) | singer : Song_release_year (text) | singer : Age (number) | singer : Is_male (others) | concert : concert_ID (number) | concert : concert_Name (text) | concert : Theme (text) | concert : Stadium_ID (text) | concert : Year (text) | singer_in_concert : concert_ID (number) | singer_in_concert : Singer_ID (text) ;\n",
      "\n",
      "[Primary Keys]: stadium : Stadium_ID | singer : Singer_ID | concert : concert_ID | singer_in_concert : concert_ID\n",
      "\n",
      "[Foreign Keys]: concert : Stadium_ID equals stadium : Stadium_ID | singer_in_concert : Singer_ID equals singer : Singer_ID | singer_in_concert : concert_ID equals concert : concert_ID\n",
      "\n",
      "[Q]: What are the number of concerts that occurred in the stadium with the largest capacity ?\n",
      "\n",
      "[SQL]: select count(*) from concert where stadium_id = (select stadium_id from stadium order by capacity desc limit 1)\n",
      "\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID ;\n",
      "\n",
      "[Column names (type)]: stadium : Stadium_ID (number) | stadium : Location (text) | stadium : Name (text) | stadium : Capacity (number) | stadium : Highest (number) | stadium : Lowest (number) | stadium : Average (number) | singer : Singer_ID (number) | singer : Name (text) | singer : Country (text) | singer : Song_Name (text) | singer : Song_release_year (text) | singer : Age (number) | singer : Is_male (others) | concert : concert_ID (number) | concert : concert_Name (text) | concert : Theme (text) | concert : Stadium_ID (text) | concert : Year (text) | singer_in_concert : concert_ID (number) | singer_in_concert : Singer_ID (text) ;\n",
      "\n",
      "[Primary Keys]: stadium : Stadium_ID | singer : Singer_ID | concert : concert_ID | singer_in_concert : concert_ID\n",
      "\n",
      "[Foreign Keys]: concert : Stadium_ID equals stadium : Stadium_ID | singer_in_concert : Singer_ID equals singer : Singer_ID | singer_in_concert : concert_ID equals concert : concert_ID\n",
      "\n",
      "[Q]: Show the stadium name and capacity with most number of concerts in year 2014 or after.\n",
      "\n",
      "[SQL]: SELECT T2.name ,  T2.capacity FROM concert AS T1 JOIN stadium AS T2 ON T1.stadium_id  =  T2.stadium_id WHERE T1.year  >=  2014 GROUP BY T2.stadium_id ORDER BY count(*) DESC LIMIT 1\n",
      "\n",
      "Given the following schema information, generate valid SQL to answer the provided query. Enclose the query in markdown code-block syntax.\n",
      "Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID ;\n",
      "\n",
      "[Column names (type)]: stadium : Stadium_ID (number) | stadium : Location (text) | stadium : Name (text) | stadium : Capacity (number) | stadium : Highest (number) | stadium : Lowest (number) | stadium : Average (number) | singer : Singer_ID (number) | singer : Name (text) | singer : Country (text) | singer : Song_Name (text) | singer : Song_release_year (text) | singer : Age (number) | singer : Is_male (others) | concert : concert_ID (number) | concert : concert_Name (text) | concert : Theme (text) | concert : Stadium_ID (text) | concert : Year (text) | singer_in_concert : concert_ID (number) | singer_in_concert : Singer_ID (text) ;\n",
      "\n",
      "[Primary Keys]: stadium : Stadium_ID | singer : Singer_ID | concert : concert_ID | singer_in_concert : concert_ID\n",
      "\n",
      "[Foreign Keys]: concert : Stadium_ID equals stadium : Stadium_ID | singer_in_concert : Singer_ID equals singer : Singer_ID | singer_in_concert : concert_ID equals concert : concert_ID\n",
      "\n",
      "[Q]: What is the name and capacity of the stadium with the most concerts after 2013 ?\n",
      "\n",
      "[SQL]: \n",
      "======\n",
      "\n",
      "===\n",
      "Generated SQL:\n",
      "SELECT s.Name, s.Capacity\n",
      "FROM stadium s\n",
      "WHERE s.Stadium_ID = (\n",
      "SELECT c.Stadium_ID\n",
      "FROM concert c\n",
      "WHERE c.Year > 2013\n",
      "GROUP BY c.Stadium_ID\n",
      "ORDER BY COUNT(*) DESC\n",
      "LIMIT 1\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "question0 = test_queries[0][\"question\"]\n",
    "db_id0 = test_queries[0][\"db_id\"]\n",
    "# Debug mode to inspect the complete prompt:\n",
    "gen_sql0 = generate_sql(\n",
    "    question=question0,\n",
    "    db_id=db_id0,\n",
    "    prompt_fn=dynamic_few_shot_prompt_fn,\n",
    "    debug_prompt=True,\n",
    ")\n",
    "\n",
    "print(\"\\n===\\nGenerated SQL:\")\n",
    "print(gen_sql0)\n",
    "\n",
    "generated_queries_dynamic[0] = gen_sql0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "00610fa1-f77f-4222-8724-30d53a89651e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM prompt for SQL generation:\n",
      "======\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | pets_1 | Student : StuID , LName , Fname , Age , Sex , Major , Advisor , city_code | Has_Pet : StuID , PetID | Pets : PetID , PetType , pet_age , weight ;\n",
      "\n",
      "[Column names (type)]: Student : StuID (number) | Student : LName (text) | Student : Fname (text) | Student : Age (number) | Student : Sex (text) | Student : Major (number) | Student : Advisor (number) | Student : city_code (text) | Has_Pet : StuID (number) | Has_Pet : PetID (number) | Pets : PetID (number) | Pets : PetType (text) | Pets : pet_age (number) | Pets : weight (number) ;\n",
      "\n",
      "[Primary Keys]: Student : StuID | Pets : PetID\n",
      "\n",
      "[Foreign Keys]: Has_Pet : StuID equals Student : StuID | Has_Pet : PetID equals Pets : PetID\n",
      "\n",
      "[Q]: Find the first name of students who have cat or dog pet.\n",
      "\n",
      "[SQL]: SELECT DISTINCT T1.Fname FROM student AS T1 JOIN has_pet AS T2 ON T1.stuid  =  T2.stuid JOIN pets AS T3 ON T3.petid  =  T2.petid WHERE T3.pettype  =  'cat' OR T3.pettype  =  'dog'\n",
      "\n",
      "Here is an example: Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | pets_1 | Student : StuID , LName , Fname , Age , Sex , Major , Advisor , city_code | Has_Pet : StuID , PetID | Pets : PetID , PetType , pet_age , weight ;\n",
      "\n",
      "[Column names (type)]: Student : StuID (number) | Student : LName (text) | Student : Fname (text) | Student : Age (number) | Student : Sex (text) | Student : Major (number) | Student : Advisor (number) | Student : city_code (text) | Has_Pet : StuID (number) | Has_Pet : PetID (number) | Pets : PetID (number) | Pets : PetType (text) | Pets : pet_age (number) | Pets : weight (number) ;\n",
      "\n",
      "[Primary Keys]: Student : StuID | Pets : PetID\n",
      "\n",
      "[Foreign Keys]: Has_Pet : StuID equals Student : StuID | Has_Pet : PetID equals Pets : PetID\n",
      "\n",
      "[Q]: Find the first name of students who have both cat and dog pets .\n",
      "\n",
      "[SQL]: select t1.fname from student as t1 join has_pet as t2 on t1.stuid  =  t2.stuid join pets as t3 on t3.petid  =  t2.petid where t3.pettype  =  'cat' intersect select t1.fname from student as t1 join has_pet as t2 on t1.stuid  =  t2.stuid join pets as t3 on t3.petid  =  t2.petid where t3.pettype  =  'dog'\n",
      "\n",
      "Given the following schema information, generate valid SQL to answer the provided query. Enclose the query in markdown code-block syntax.\n",
      "Convert text to SQL:\n",
      "\n",
      "[Schema : (values)]: | pets_1 | Student : StuID , LName , Fname , Age , Sex , Major , Advisor , city_code | Has_Pet : StuID , PetID | Pets : PetID , PetType , pet_age , weight ;\n",
      "\n",
      "[Column names (type)]: Student : StuID (number) | Student : LName (text) | Student : Fname (text) | Student : Age (number) | Student : Sex (text) | Student : Major (number) | Student : Advisor (number) | Student : city_code (text) | Has_Pet : StuID (number) | Has_Pet : PetID (number) | Pets : PetID (number) | Pets : PetType (text) | Pets : pet_age (number) | Pets : weight (number) ;\n",
      "\n",
      "[Primary Keys]: Student : StuID | Pets : PetID\n",
      "\n",
      "[Foreign Keys]: Has_Pet : StuID equals Student : StuID | Has_Pet : PetID equals Pets : PetID\n",
      "\n",
      "[Q]: Find the first name and age of students who have a dog but do not have a cat as a pet.\n",
      "\n",
      "[SQL]: \n",
      "======\n",
      "\n",
      "===\n",
      "Generated SQL:\n",
      "SELECT T1.Fname, T1.Age\n",
      "FROM Student AS T1\n",
      "JOIN Has_Pet AS T2 ON T1.StuID = T2.StuID\n",
      "JOIN Pets AS T3 ON T2.PetID = T3.PetID\n",
      "WHERE T3.PetType = 'dog'\n",
      "AND T1.StuID NOT IN (\n",
      "SELECT T1.StuID\n",
      "FROM Student AS T1\n",
      "JOIN Has_Pet AS T2 ON T1.StuID = T2.StuID\n",
      "JOIN Pets AS T3 ON T2.PetID = T3.PetID\n",
      "WHERE T3.PetType = 'cat'\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "question4 = test_queries[4][\"question\"]\n",
    "db_id4 = test_queries[4][\"db_id\"]\n",
    "# Debug mode to inspect the complete prompt:\n",
    "gen_sql4 = generate_sql(\n",
    "    question=question4,\n",
    "    db_id=db_id4,\n",
    "    prompt_fn=dynamic_few_shot_prompt_fn,\n",
    "    debug_prompt=True,\n",
    ")\n",
    "\n",
    "print(\"\\n===\\nGenerated SQL:\")\n",
    "print(gen_sql4)\n",
    "\n",
    "generated_queries_dynamic[4] = gen_sql4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cb5a3e4-1228-4d95-b766-8769221e8505",
   "metadata": {},
   "source": [
    "The other queries are similarly computed (with prompt debug turned off):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "31a530ac-996c-4f2d-868b-174a606c1c68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "====\n",
      "Added for Q[1]='List singer names and number of concerts for each singer.' -> SQL:\n",
      "----\n",
      "SELECT s.Name, COUNT(sc.concert_ID) AS num_concerts\n",
      "FROM singer s\n",
      "JOIN singer_in_concert sc ON s.Singer_ID = sc.Singer_ID\n",
      "GROUP BY s.Name;\n",
      "\n",
      "====\n",
      "Added for Q[2]='Find the number of pets whose weight is heavier than 10.' -> SQL:\n",
      "----\n",
      "SELECT COUNT(*)\n",
      "FROM PETS\n",
      "WHERE weight > 10;\n",
      "\n",
      "====\n",
      "Added for Q[3]='Find the number of distinct type of pets.' -> SQL:\n",
      "----\n",
      "SELECT COUNT(DISTINCT PetType) AS num_distinct_pet_types FROM PETS;\n",
      "\n",
      "====\n",
      "Added for Q[5]='Find the first name and age of students who have a pet.' -> SQL:\n",
      "----\n",
      "SELECT Fname, Age\n",
      "FROM STUDENT\n",
      "WHERE StuID IN (SELECT StuID FROM HAS_PET)\n"
     ]
    }
   ],
   "source": [
    "for test_i, test_row in test_queries.items():\n",
    "    if test_i not in generated_queries_dynamic:\n",
    "        generated_queries_dynamic[test_i] =  generate_sql(\n",
    "            question=test_row[\"question\"],\n",
    "            db_id=test_row[\"db_id\"],\n",
    "            prompt_fn=dynamic_few_shot_prompt_fn,\n",
    "        )\n",
    "        print(\n",
    "            f\"\\n====\\nAdded for Q[{test_i}]=\"\n",
    "            f\"'{test_row['question']}' -> SQL:\\n----\"\n",
    "        )\n",
    "        print(generated_queries_dynamic[test_i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e11cac6b-45c8-4ee0-9fc1-a78ccbcc1f69",
   "metadata": {},
   "source": [
    "Time to run the evaluation step for all these queries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "8e779a23-cab3-4773-ad1e-38c2c7c7a80b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>DB ID</th>\n",
       "      <th>Question</th>\n",
       "      <th>Correct</th>\n",
       "      <th>Error</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>What is the name and capacity of the stadium w...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>concert_singer</td>\n",
       "      <td>List singer names and number of concerts for e...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "      <th>2</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of pets whose weight is heavie...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the number of distinct type of pets.</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pets_1</td>\n",
       "      <td>Find the first name and age of students who ha...</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            DB ID                                           Question  Correct  \\\n",
       "0  concert_singer  What is the name and capacity of the stadium w...     True   \n",
       "1  concert_singer  List singer names and number of concerts for e...     True   \n",
       "2          pets_1  Find the number of pets whose weight is heavie...     True   \n",
       "3          pets_1          Find the number of distinct type of pets.     True   \n",
       "4          pets_1  Find the first name and age of students who ha...     True   \n",
       "5          pets_1  Find the first name and age of students who ha...     True   \n",
       "\n",
       "  Error  \n",
       "0  None  \n",
       "1  None  \n",
       "2  None  \n",
       "3  None  \n",
       "4  None  \n",
       "5  None  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "report = eval_generated_queries(generated_queries_dynamic)\n",
    "display(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd98ced5-a3f3-451d-b558-edb535497581",
   "metadata": {},
   "source": [
    "Adoption of dynamic few-shot prompting has proven effective to get all questions right. **Success!**\n",
    "\n",
    "_Note: keep in mind that, as is usual when working with LLMs which are notably non-deterministic,\n",
    "results may differ between runs: you could still see the occasional incorrectly-generated SQL when running this notebook._"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6d900e6-f921-4b07-a207-95caed748624",
   "metadata": {
    "id": "2ea044b5-7319-4b2e-b083-11d0549f678d"
   },
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cada3d4-ea50-4edf-ae21-6c575d976cb6",
   "metadata": {},
   "source": [
    "Text-to-sql is a very critical skill for an LLM, enabling a variety of use cases and GenAI-based automation opportunities.\n",
    "\n",
    "AI-based solution can often be deployed at various levels of sophistication: in the case at hand, adopting a dynamic approach to few-shot\n",
    "prompting makes SQL generation nearly 100% accurate.\n",
    "\n",
    "For the retrieval of the most appropriate question/query examples, Astra DB is a perfect fit, thanks to its capabilities as a vector database.\n",
    "\n",
    "If you want to learn more on vector databases, GenAI, or you simply want to experiment with Astra DB, visit the [Astra DB docs](https://docs.datastax.com/en/astra/astra-db-vector/get-started/quickstart.html#create-a-serverless-vector-database)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "021f96ab-069e-4b80-b740-c14d2c5f813a",
   "metadata": {},
   "source": [
    "## Additional references"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef8c547a-f854-4331-ab62-881738107dab",
   "metadata": {},
   "source": [
    "- [Spider homepage](https://yale-lily.github.io/spider)\n",
    "- [SQL-PaLM](https://arxiv.org/abs/2306.00739) (arXiv)\n",
    "- [Language Models are Few-Shot Learners](https://arxiv.org/abs/2005.14165) (arXiv)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}