04-Pattern-IV/00-self-correct.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infer schema directly & self-correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vertexai.language_models import CodeGenerationModel\n",
    "from google.cloud import bigquery\n",
    "import logging \n",
    "import os "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Setup logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "logger.addHandler(logging.StreamHandler())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup essentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVICE_ACCOUNT_CREDENTIALS = './../credentials/vai-key.json'\n",
    "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = SERVICE_ACCOUNT_CREDENTIALS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = 'arun-genai-bb'\n",
    "CODE_GEN_MODEL_NAME = 'code-bison@latest'\n",
    "TEMPERATURE = 0 # default value = 0\n",
    "MAX_OUTPUT_TOKENS = 2048  # length of the output response | overridding the default value which is 128\n",
    "# TOP_P = 0.95  # default value\n",
    "# TOP_K = 40  # default value\n",
    "LOCATION = 'us-central1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = 'flight_reservations'\n",
    "TABLES = ['customers', 'flights', 'reservations', 'transactions', 'loyality_points']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "code_gen_model = CodeGenerationModel.from_pretrained(CODE_GEN_MODEL_NAME)\n",
    "bq_client = bigquery.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "    SELECT *\n",
      "    FROM `arun-genai-bb.flight_reservations.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS`\n",
      "    WHERE table_name in (\"customers\",\"flights\",\"reservations\",\"transactions\",\"loyality_points\")\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query = f\"\"\"\n",
    "    SELECT *\n",
    "    FROM `{PROJECT_ID}.{DATASET}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS`\n",
    "    WHERE table_name in ({','.join([f'\"{table}\"' for table in TABLES])})\n",
    "\"\"\"\n",
    "logger.info(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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>table_catalog</th>\n",
       "      <th>table_schema</th>\n",
       "      <th>table_name</th>\n",
       "      <th>column_name</th>\n",
       "      <th>field_path</th>\n",
       "      <th>data_type</th>\n",
       "      <th>description</th>\n",
       "      <th>collation_name</th>\n",
       "      <th>rounding_mode</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>transaction_id</td>\n",
       "      <td>transaction_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>amount</td>\n",
       "      <td>amount</td>\n",
       "      <td>FLOAT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>transaction_datetime</td>\n",
       "      <td>transaction_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>reservation_datetime</td>\n",
       "      <td>reservation_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>status</td>\n",
       "      <td>status</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>origin</td>\n",
       "      <td>origin</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>destination</td>\n",
       "      <td>destination</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>departure_datetime</td>\n",
       "      <td>departure_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>arrival_datetime</td>\n",
       "      <td>arrival_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>carrier</td>\n",
       "      <td>carrier</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>price</td>\n",
       "      <td>price</td>\n",
       "      <td>FLOAT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>first_name</td>\n",
       "      <td>first_name</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>last_name</td>\n",
       "      <td>last_name</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>email</td>\n",
       "      <td>email</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>date_of_birth</td>\n",
       "      <td>date_of_birth</td>\n",
       "      <td>DATE</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>created_at</td>\n",
       "      <td>created_at</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    table_catalog         table_schema    table_name           column_name  \\\n",
       "0   arun-genai-bb  flight_reservations  transactions        transaction_id   \n",
       "1   arun-genai-bb  flight_reservations  transactions        reservation_id   \n",
       "2   arun-genai-bb  flight_reservations  transactions                amount   \n",
       "3   arun-genai-bb  flight_reservations  transactions  transaction_datetime   \n",
       "4   arun-genai-bb  flight_reservations  reservations        reservation_id   \n",
       "5   arun-genai-bb  flight_reservations  reservations           customer_id   \n",
       "6   arun-genai-bb  flight_reservations  reservations             flight_id   \n",
       "7   arun-genai-bb  flight_reservations  reservations  reservation_datetime   \n",
       "8   arun-genai-bb  flight_reservations  reservations                status   \n",
       "9   arun-genai-bb  flight_reservations       flights             flight_id   \n",
       "10  arun-genai-bb  flight_reservations       flights                origin   \n",
       "11  arun-genai-bb  flight_reservations       flights           destination   \n",
       "12  arun-genai-bb  flight_reservations       flights    departure_datetime   \n",
       "13  arun-genai-bb  flight_reservations       flights      arrival_datetime   \n",
       "14  arun-genai-bb  flight_reservations       flights               carrier   \n",
       "15  arun-genai-bb  flight_reservations       flights                 price   \n",
       "16  arun-genai-bb  flight_reservations     customers           customer_id   \n",
       "17  arun-genai-bb  flight_reservations     customers            first_name   \n",
       "18  arun-genai-bb  flight_reservations     customers             last_name   \n",
       "19  arun-genai-bb  flight_reservations     customers                 email   \n",
       "20  arun-genai-bb  flight_reservations     customers         date_of_birth   \n",
       "21  arun-genai-bb  flight_reservations     customers            created_at   \n",
       "\n",
       "              field_path data_type description collation_name rounding_mode  \n",
       "0         transaction_id     INT64        None           NULL          None  \n",
       "1         reservation_id     INT64        None           NULL          None  \n",
       "2                 amount   FLOAT64        None           NULL          None  \n",
       "3   transaction_datetime  DATETIME        None           NULL          None  \n",
       "4         reservation_id     INT64        None           NULL          None  \n",
       "5            customer_id     INT64        None           NULL          None  \n",
       "6              flight_id     INT64        None           NULL          None  \n",
       "7   reservation_datetime  DATETIME        None           NULL          None  \n",
       "8                 status    STRING        None           NULL          None  \n",
       "9              flight_id     INT64        None           NULL          None  \n",
       "10                origin    STRING        None           NULL          None  \n",
       "11           destination    STRING        None           NULL          None  \n",
       "12    departure_datetime  DATETIME        None           NULL          None  \n",
       "13      arrival_datetime  DATETIME        None           NULL          None  \n",
       "14               carrier    STRING        None           NULL          None  \n",
       "15                 price   FLOAT64        None           NULL          None  \n",
       "16           customer_id     INT64        None           NULL          None  \n",
       "17            first_name    STRING        None           NULL          None  \n",
       "18             last_name    STRING        None           NULL          None  \n",
       "19                 email    STRING        None           NULL          None  \n",
       "20         date_of_birth      DATE        None           NULL          None  \n",
       "21            created_at  DATETIME        None           NULL          None  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema_columns = bq_client.query(query=query).to_dataframe()\n",
    "schema_columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n"
     ]
    }
   ],
   "source": [
    "schema_columns = schema_columns.to_markdown(index=False)\n",
    "logger.info(schema_columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Utility function to convert text into SQL and automatically correct the SQL query if any execution errors occur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_and_execute_sql(prompt, max_tries=8):\n",
    "    \"\"\"\n",
    "    Generate an SQL query using the code_gen_model and execute it using bq_client.\n",
    "    \n",
    "    Args:\n",
    "    - prompt (str): Prompt to provide to the model for generating SQL.\n",
    "    - max_tries (int): Maximum number of attempts to generate and execute SQL.\n",
    "    \n",
    "    Returns:\n",
    "    - dict: A dictionary containing the successful dataframe or error messages and prompt evolution.\n",
    "    \"\"\"\n",
    "    \n",
    "    tries = 0\n",
    "    error_messages = []\n",
    "    prompts = [prompt]\n",
    "    df = None\n",
    "    \n",
    "    while tries < max_tries:\n",
    "        logger.info(f'ATTEMPT: {tries+1}')\n",
    "        try:\n",
    "            # Predict SQL using the model\n",
    "            response = code_gen_model.predict(prompt, temperature=TEMPERATURE, max_output_tokens=MAX_OUTPUT_TOKENS)\n",
    "            generated_sql_query = response.text\n",
    "            generated_sql_query = '\\n'.join(generated_sql_query.split('\\n')[1:-1])\n",
    "            logger.info('-' * 50)\n",
    "            logger.info(generated_sql_query)\n",
    "            logger.info('-' * 50)\n",
    "            # Execute SQL using BigQuery client\n",
    "            df = bq_client.query(generated_sql_query).to_dataframe()\n",
    "            logger.info('SUCCEEDED')\n",
    "            return {\"dataframe\": df, \"prompts\": prompts, \"errors\": error_messages}\n",
    "        except Exception as e:\n",
    "            logger.error('FAILED')\n",
    "            # Catch the error, store the message, and try again\n",
    "            msg = str(e)\n",
    "            error_messages.append(msg)\n",
    "            # Evolve the prompt by appending the error message and asking the model to correct it\n",
    "            prompt = f\"\"\"{prompt}\n",
    "Encountered an error: {msg}. \n",
    "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
    "Follow the instructions mentioned above to remediate the error. \n",
    "\n",
    "Modify the below SQL query to resolve the issue:\n",
    "{generated_sql_query}\n",
    "\n",
    "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\"\"\"\n",
    "            prompts.append(prompt)\n",
    "            tries += 1\n",
    "        logger.info('=' * 100)\n",
    "\n",
    "    return {\"dataframe\": df, \"prompts\": prompts, \"errors\": error_messages}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Text-to-SQL scenarios"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct the SEED prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "{}\n",
      "\n",
      "SCHEMA:\n",
      "{}\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "seed_prompt = \"\"\"\n",
    "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
    "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
    "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
    "Additionally, the output column names should specify units where applicable.\\n\n",
    "QUESTION:\n",
    "{}\\n\n",
    "SCHEMA:\n",
    "{}\\n\n",
    "IMPORTANT: \n",
    "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
    "--\n",
    "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
    "\"\"\"\n",
    "logger.info(seed_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 1: Retrieve Active Reservations for a Specific Date Range"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this scenario, you want to find all active reservations within a specific date range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Provide a list of all reservations from October 10th to October 15th, 2023\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Provide a list of all reservations from October 10th to October 15th, 2023\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)\n",
    "logger.info(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  r.reservation_id,\n",
      "  r.customer_id,\n",
      "  r.flight_id,\n",
      "  r.reservation_datetime AS reservation_datetime_utc,\n",
      "  f.origin,\n",
      "  f.destination,\n",
      "  f.departure_datetime AS departure_datetime_utc,\n",
      "  f.arrival_datetime AS arrival_datetime_utc,\n",
      "  f.carrier,\n",
      "  f.price\n",
      "FROM arun-genai-bb.flight_reservations.reservations AS r\n",
      "JOIN arun-genai-bb.flight_reservations.flights AS f\n",
      "ON CAST(r.flight_id AS STRING) = f.flight_id\n",
      "WHERE r.reservation_datetime >= '2023-10-10 00:00:00'\n",
      "AND r.reservation_datetime < '2023-10-16 00:00:00';\n",
      "--------------------------------------------------\n",
      "FAILED\n",
      "====================================================================================================\n",
      "ATTEMPT: 2\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  r.reservation_id,\n",
      "  r.customer_id,\n",
      "  r.flight_id,\n",
      "  r.reservation_datetime AS reservation_datetime_utc,\n",
      "  f.origin,\n",
      "  f.destination,\n",
      "  f.departure_datetime AS departure_datetime_utc,\n",
      "  f.arrival_datetime AS arrival_datetime_utc,\n",
      "  f.carrier,\n",
      "  f.price\n",
      "FROM arun-genai-bb.flight_reservations.reservations AS r\n",
      "JOIN arun-genai-bb.flight_reservations.flights AS f\n",
      "ON CAST(r.flight_id AS STRING) = CAST(f.flight_id AS STRING)\n",
      "WHERE r.reservation_datetime >= '2023-10-10 00:00:00'\n",
      "AND r.reservation_datetime < '2023-10-16 00:00:00';\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 115 ms, sys: 19.5 ms, total: 135 ms\n",
      "Wall time: 9.64 s\n"
     ]
    },
    {
     "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>reservation_id</th>\n",
       "      <th>customer_id</th>\n",
       "      <th>flight_id</th>\n",
       "      <th>reservation_datetime_utc</th>\n",
       "      <th>origin</th>\n",
       "      <th>destination</th>\n",
       "      <th>departure_datetime_utc</th>\n",
       "      <th>arrival_datetime_utc</th>\n",
       "      <th>carrier</th>\n",
       "      <th>price</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>6</td>\n",
       "      <td>6</td>\n",
       "      <td>2023-10-10 10:00:00</td>\n",
       "      <td>SEA</td>\n",
       "      <td>JFK</td>\n",
       "      <td>2023-11-25 06:00:00</td>\n",
       "      <td>2023-11-25 14:30:00</td>\n",
       "      <td>United</td>\n",
       "      <td>550.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7</td>\n",
       "      <td>6</td>\n",
       "      <td>7</td>\n",
       "      <td>2023-10-12 11:30:00</td>\n",
       "      <td>JFK</td>\n",
       "      <td>MIA</td>\n",
       "      <td>2023-11-27 20:00:00</td>\n",
       "      <td>2023-11-27 23:30:00</td>\n",
       "      <td>American</td>\n",
       "      <td>380.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>8</td>\n",
       "      <td>8</td>\n",
       "      <td>8</td>\n",
       "      <td>2023-10-15 13:20:00</td>\n",
       "      <td>MIA</td>\n",
       "      <td>JFK</td>\n",
       "      <td>2023-11-30 10:00:00</td>\n",
       "      <td>2023-11-30 13:30:00</td>\n",
       "      <td>American</td>\n",
       "      <td>380.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   reservation_id  customer_id  flight_id reservation_datetime_utc origin  \\\n",
       "0               6            6          6      2023-10-10 10:00:00    SEA   \n",
       "1               7            6          7      2023-10-12 11:30:00    JFK   \n",
       "2               8            8          8      2023-10-15 13:20:00    MIA   \n",
       "\n",
       "  destination departure_datetime_utc arrival_datetime_utc   carrier  price  \n",
       "0         JFK    2023-11-25 06:00:00  2023-11-25 14:30:00    United  550.0  \n",
       "1         MIA    2023-11-27 20:00:00  2023-11-27 23:30:00  American  380.0  \n",
       "2         JFK    2023-11-30 10:00:00  2023-11-30 13:30:00  American  380.0  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 2: Identify customers who made reservations in the past N days."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Identify all customers who have made flight reservations within the last 7 days.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       c.email,\n",
      "       t.transaction_datetime AS reservation_date\n",
      "FROM flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "JOIN flight_reservations.transactions t\n",
      "ON r.reservation_id = t.reservation_id\n",
      "WHERE t.transaction_datetime >= DATE_SUB(CURRENT_DATETIME(), INTERVAL 7 DAY)\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 50.3 ms, sys: 9.21 ms, total: 59.5 ms\n",
      "Wall time: 4.86 s\n"
     ]
    },
    {
     "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>customer_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>email</th>\n",
       "      <th>reservation_date</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>Diana</td>\n",
       "      <td>Prince</td>\n",
       "      <td>diana.p@example.com</td>\n",
       "      <td>2023-10-12 11:31:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>8</td>\n",
       "      <td>Fiona</td>\n",
       "      <td>Shrek</td>\n",
       "      <td>fiona.s@example.com</td>\n",
       "      <td>2023-10-15 13:21:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>10</td>\n",
       "      <td>Hannah</td>\n",
       "      <td>Montana</td>\n",
       "      <td>hannah.m@example.com</td>\n",
       "      <td>2023-10-22 15:46:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>11</td>\n",
       "      <td>Ian</td>\n",
       "      <td>Somerhalder</td>\n",
       "      <td>ian.s@example.com</td>\n",
       "      <td>2023-10-25 12:31:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>11</td>\n",
       "      <td>Ian</td>\n",
       "      <td>Somerhalder</td>\n",
       "      <td>ian.s@example.com</td>\n",
       "      <td>2023-10-28 17:11:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>13</td>\n",
       "      <td>Kate</td>\n",
       "      <td>Winslet</td>\n",
       "      <td>kate.w@example.com</td>\n",
       "      <td>2023-11-02 08:21:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>15</td>\n",
       "      <td>Mary</td>\n",
       "      <td>Jane</td>\n",
       "      <td>mary.j@example.com</td>\n",
       "      <td>2023-11-04 10:46:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>16</td>\n",
       "      <td>Nick</td>\n",
       "      <td>Fury</td>\n",
       "      <td>nick.f@example.com</td>\n",
       "      <td>2023-11-08 15:31:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>17</td>\n",
       "      <td>Olivia</td>\n",
       "      <td>Newton</td>\n",
       "      <td>olivia.n@example.com</td>\n",
       "      <td>2023-11-11 10:16:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>18</td>\n",
       "      <td>Peter</td>\n",
       "      <td>Parker</td>\n",
       "      <td>peter.p@example.com</td>\n",
       "      <td>2023-11-15 12:51:00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>20</td>\n",
       "      <td>Ryan</td>\n",
       "      <td>Reynolds</td>\n",
       "      <td>ryan.r@example.com</td>\n",
       "      <td>2023-11-22 09:11:00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    customer_id first_name    last_name                 email  \\\n",
       "0             6      Diana       Prince   diana.p@example.com   \n",
       "1             8      Fiona        Shrek   fiona.s@example.com   \n",
       "2            10     Hannah      Montana  hannah.m@example.com   \n",
       "3            11        Ian  Somerhalder     ian.s@example.com   \n",
       "4            11        Ian  Somerhalder     ian.s@example.com   \n",
       "5            13       Kate      Winslet    kate.w@example.com   \n",
       "6            15       Mary         Jane    mary.j@example.com   \n",
       "7            16       Nick         Fury    nick.f@example.com   \n",
       "8            17     Olivia       Newton  olivia.n@example.com   \n",
       "9            18      Peter       Parker   peter.p@example.com   \n",
       "10           20       Ryan     Reynolds    ryan.r@example.com   \n",
       "\n",
       "      reservation_date  \n",
       "0  2023-10-12 11:31:00  \n",
       "1  2023-10-15 13:21:00  \n",
       "2  2023-10-22 15:46:00  \n",
       "3  2023-10-25 12:31:00  \n",
       "4  2023-10-28 17:11:00  \n",
       "5  2023-11-02 08:21:00  \n",
       "6  2023-11-04 10:46:00  \n",
       "7  2023-11-08 15:31:00  \n",
       "8  2023-11-11 10:16:00  \n",
       "9  2023-11-15 12:51:00  \n",
       "10 2023-11-22 09:11:00  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 3: Calculate Monthly Revenue\n",
    "Calculate the total revenue generated from transactions for a given month and year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Calculate the total revenue generated from transactions in October 2023, specifically from all reservations with a Confirmed status.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  SUM(transactions.amount) AS total_revenue_usd\n",
      "FROM\n",
      "  arun-genai-bb.flight_reservations.transactions AS transactions\n",
      "JOIN\n",
      "  arun-genai-bb.flight_reservations.reservations AS reservations\n",
      "ON\n",
      "  transactions.reservation_id = reservations.reservation_id\n",
      "WHERE\n",
      "  reservations.status = 'Confirmed'\n",
      "  AND transactions.transaction_datetime BETWEEN '2023-10-01 00:00:00' AND '2023-10-31 23:59:59';\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 60.3 ms, sys: 10.4 ms, total: 70.8 ms\n",
      "Wall time: 4.44 s\n"
     ]
    },
    {
     "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>total_revenue_usd</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3860.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   total_revenue_usd\n",
       "0             3860.0"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 4: Popular Flight Times\n",
    "Identify the most popular departure hours or days for a given day or month or year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Determine the departure months with the highest frequency for the year 2023.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT DATE_TRUNC(departure_datetime, MONTH) AS departure_month,\n",
      "       COUNT(*) AS num_departures\n",
      "FROM flight_reservations.flights\n",
      "WHERE departure_datetime BETWEEN '2023-01-01' AND '2023-12-31'\n",
      "GROUP BY departure_month\n",
      "ORDER BY num_departures DESC\n",
      "LIMIT 10;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 63.7 ms, sys: 10.9 ms, total: 74.6 ms\n",
      "Wall time: 4.69 s\n"
     ]
    },
    {
     "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>departure_month</th>\n",
       "      <th>num_departures</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2023-11-01</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2023-12-01</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  departure_month  num_departures\n",
       "0      2023-11-01               8\n",
       "1      2023-12-01               7"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 5: Customer Age Group\n",
    "Group customers by age brackets and count the number in each bracket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Group customers into five distinct age brackets and count the number of customers in each bracket.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  CASE\n",
      "    WHEN date_of_birth IS NULL THEN 'Unknown'\n",
      "    WHEN date_of_birth >= DATE_SUB(CURRENT_DATE(), INTERVAL 65 YEAR) THEN '65+'\n",
      "    WHEN date_of_birth >= DATE_SUB(CURRENT_DATE(), INTERVAL 55 YEAR) THEN '55-64'\n",
      "    WHEN date_of_birth >= DATE_SUB(CURRENT_DATE(), INTERVAL 45 YEAR) THEN '45-54'\n",
      "    WHEN date_of_birth >= DATE_SUB(CURRENT_DATE(), INTERVAL 35 YEAR) THEN '35-44'\n",
      "    ELSE '0-34'\n",
      "  END AS age_bracket,\n",
      "  COUNT(customer_id) AS num_customers\n",
      "FROM flight_reservations.customers\n",
      "GROUP BY age_bracket;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 60.9 ms, sys: 10.7 ms, total: 71.6 ms\n",
      "Wall time: 5.62 s\n"
     ]
    },
    {
     "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>age_bracket</th>\n",
       "      <th>num_customers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>65+</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0-34</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  age_bracket  num_customers\n",
       "0         65+             17\n",
       "1        0-34              3"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 6: Age Calculation\n",
    "Calculate the age of customers based on their date of birth and filter those who are above X years old."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 >= 18\n",
      "ORDER BY age\n",
      "--------------------------------------------------\n",
      "FAILED\n",
      "====================================================================================================\n",
      "ATTEMPT: 2\n",
      "--------------------------------------------------\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) >= 18\n",
      "--------------------------------------------------\n",
      "FAILED\n",
      "====================================================================================================\n",
      "ATTEMPT: 3\n",
      "--------------------------------------------------\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) >= 18\n",
      "ORDER BY age;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 86.8 ms, sys: 21.3 ms, total: 108 ms\n",
      "Wall time: 16 s\n"
     ]
    },
    {
     "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>customer_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "      <th>reservation_id</th>\n",
       "      <th>status</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10</td>\n",
       "      <td>Hannah</td>\n",
       "      <td>Montana</td>\n",
       "      <td>25</td>\n",
       "      <td>10</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>8</td>\n",
       "      <td>Fiona</td>\n",
       "      <td>Shrek</td>\n",
       "      <td>31</td>\n",
       "      <td>8</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Alice</td>\n",
       "      <td>Johnson</td>\n",
       "      <td>33</td>\n",
       "      <td>3</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>6</td>\n",
       "      <td>Diana</td>\n",
       "      <td>Prince</td>\n",
       "      <td>35</td>\n",
       "      <td>5</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>6</td>\n",
       "      <td>Diana</td>\n",
       "      <td>Prince</td>\n",
       "      <td>35</td>\n",
       "      <td>6</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>6</td>\n",
       "      <td>Diana</td>\n",
       "      <td>Prince</td>\n",
       "      <td>35</td>\n",
       "      <td>7</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>Jane</td>\n",
       "      <td>Doe</td>\n",
       "      <td>36</td>\n",
       "      <td>2</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>John</td>\n",
       "      <td>Doe</td>\n",
       "      <td>38</td>\n",
       "      <td>1</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>11</td>\n",
       "      <td>Ian</td>\n",
       "      <td>Somerhalder</td>\n",
       "      <td>45</td>\n",
       "      <td>11</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>11</td>\n",
       "      <td>Ian</td>\n",
       "      <td>Somerhalder</td>\n",
       "      <td>45</td>\n",
       "      <td>12</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   customer_id first_name    last_name  age  reservation_id     status\n",
       "0           10     Hannah      Montana   25              10  Confirmed\n",
       "1            8      Fiona        Shrek   31               8  Confirmed\n",
       "2            3      Alice      Johnson   33               3  Confirmed\n",
       "3            6      Diana       Prince   35               5  Confirmed\n",
       "4            6      Diana       Prince   35               6  Confirmed\n",
       "5            6      Diana       Prince   35               7  Confirmed\n",
       "6            2       Jane          Doe   36               2  Confirmed\n",
       "7            1       John          Doe   38               1  Confirmed\n",
       "8           11        Ian  Somerhalder   45              11  Confirmed\n",
       "9           11        Ian  Somerhalder   45              12  Confirmed"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect the evolution of the seed prompt and how SQL query was fixed automatically by the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "==================== ATTEMPT 1 ====================\n",
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n",
      "==================== ATTEMPT 2 ====================\n",
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n",
      "Encountered an error: 400 Function not found: DATEDIFF; Did you mean date_diff? at [18:5]\n",
      "\n",
      "Location: US\n",
      "Job ID: 410a5db4-0782-4e1b-b4a2-a8d8c31c021f\n",
      ". \n",
      "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
      "Follow the instructions mentioned above to remediate the error. \n",
      "\n",
      "Modify the below SQL query to resolve the issue:\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 >= 18\n",
      "ORDER BY age\n",
      "\n",
      "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\n",
      "==================== ATTEMPT 3 ====================\n",
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n",
      "Encountered an error: 400 Function not found: DATEDIFF; Did you mean date_diff? at [18:5]\n",
      "\n",
      "Location: US\n",
      "Job ID: 410a5db4-0782-4e1b-b4a2-a8d8c31c021f\n",
      ". \n",
      "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
      "Follow the instructions mentioned above to remediate the error. \n",
      "\n",
      "Modify the below SQL query to resolve the issue:\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATEDIFF(current_month.end_of_month, c.date_of_birth) / 365 >= 18\n",
      "ORDER BY age\n",
      "\n",
      "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\n",
      "Encountered an error: 400 Syntax error: Expected end of input but got \")\" at [3:1]\n",
      "\n",
      "Location: US\n",
      "Job ID: 66a0a6c7-51fa-40e0-901b-1fc55c269ae0\n",
      ". \n",
      "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
      "Follow the instructions mentioned above to remediate the error. \n",
      "\n",
      "Modify the below SQL query to resolve the issue:\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_of_month,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_of_month\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month\n",
      "CROSS JOIN flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.status = 'Confirmed'\n",
      "AND r.reservation_datetime BETWEEN current_month.start_of_month AND current_month.end_of_month\n",
      "AND DATE_DIFF(current_month.end_of_month, c.date_of_birth, YEAR) >= 18\n",
      "\n",
      "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\n"
     ]
    }
   ],
   "source": [
    "for i, prompt in enumerate(response['prompts']):\n",
    "    logger.info(f'==================== ATTEMPT {i+1} ====================')\n",
    "    logger.info(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".bq-sql-agent",
   "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.9.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}