openai_poc.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data\n",
    "\n",
    "Loading and appending the Chicago Crime dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import duckdb\n",
    "import openai\n",
    "import chicago_data\n",
    "import datetime\n",
    "import time \n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = \"https://data.cityofchicago.org/resource/ijzp-q8t2\"\n",
    "start = datetime.datetime(2024,1,1,0,0,0)\n",
    "end = datetime.datetime(2024,4,24,0,0,0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "chicago_crime = chicago_data.backfill_chicago_data(endpoint = endpoint, \n",
    "                                       start = start, \n",
    "                                       end = end, \n",
    "                                       offset = 24 * 30,\n",
    "                                       limit = 100000)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "chicago_crime.to_csv(\"data/chicago_crime.csv\", index = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the path of the CSV files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['chicago_crime_2020-2024.csv', 'chicago_crime.csv']\n"
     ]
    }
   ],
   "source": [
    "path = \"./data\"\n",
    "\n",
    "files = [x for x in os.listdir(path = path) if \".csv\" in x]\n",
    "\n",
    "print(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "ParserError",
     "evalue": "Error tokenizing data. C error: Calling read(nbytes) on source failed. Try engine='python'.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mParserError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m chicago_crime \u001b[38;5;241m=\u001b[39m \u001b[43mpd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconcat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_csv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfiles\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m chicago_crime\u001b[38;5;241m.\u001b[39mhead()\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/core/reshape/concat.py:372\u001b[0m, in \u001b[0;36mconcat\u001b[0;34m(objs, axis, join, ignore_index, keys, levels, names, verify_integrity, sort, copy)\u001b[0m\n\u001b[1;32m    369\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m copy \u001b[38;5;129;01mand\u001b[39;00m using_copy_on_write():\n\u001b[1;32m    370\u001b[0m     copy \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m--> 372\u001b[0m op \u001b[38;5;241m=\u001b[39m \u001b[43m_Concatenator\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    373\u001b[0m \u001b[43m    \u001b[49m\u001b[43mobjs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    374\u001b[0m \u001b[43m    \u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43maxis\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    375\u001b[0m \u001b[43m    \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_index\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    376\u001b[0m \u001b[43m    \u001b[49m\u001b[43mjoin\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mjoin\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    377\u001b[0m \u001b[43m    \u001b[49m\u001b[43mkeys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    378\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlevels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlevels\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    379\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnames\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnames\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    380\u001b[0m \u001b[43m    \u001b[49m\u001b[43mverify_integrity\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mverify_integrity\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    381\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcopy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcopy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    382\u001b[0m \u001b[43m    \u001b[49m\u001b[43msort\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msort\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    383\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    385\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m op\u001b[38;5;241m.\u001b[39mget_result()\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/core/reshape/concat.py:426\u001b[0m, in \u001b[0;36m_Concatenator.__init__\u001b[0;34m(self, objs, axis, join, keys, levels, names, ignore_index, verify_integrity, copy, sort)\u001b[0m\n\u001b[1;32m    424\u001b[0m     objs \u001b[38;5;241m=\u001b[39m [objs[k] \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m keys]\n\u001b[1;32m    425\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 426\u001b[0m     objs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mlist\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mobjs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    428\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(objs) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m    429\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNo objects to concatenate\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[6], line 1\u001b[0m, in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[0;32m----> 1\u001b[0m chicago_crime \u001b[38;5;241m=\u001b[39m pd\u001b[38;5;241m.\u001b[39mconcat((\u001b[43mpd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_csv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m/\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m files), ignore_index\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m      3\u001b[0m chicago_crime\u001b[38;5;241m.\u001b[39mhead()\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/io/parsers/readers.py:912\u001b[0m, in \u001b[0;36mread_csv\u001b[0;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, date_format, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, encoding_errors, dialect, on_bad_lines, delim_whitespace, low_memory, memory_map, float_precision, storage_options, dtype_backend)\u001b[0m\n\u001b[1;32m    899\u001b[0m kwds_defaults \u001b[38;5;241m=\u001b[39m _refine_defaults_read(\n\u001b[1;32m    900\u001b[0m     dialect,\n\u001b[1;32m    901\u001b[0m     delimiter,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    908\u001b[0m     dtype_backend\u001b[38;5;241m=\u001b[39mdtype_backend,\n\u001b[1;32m    909\u001b[0m )\n\u001b[1;32m    910\u001b[0m kwds\u001b[38;5;241m.\u001b[39mupdate(kwds_defaults)\n\u001b[0;32m--> 912\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_read\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilepath_or_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/io/parsers/readers.py:583\u001b[0m, in \u001b[0;36m_read\u001b[0;34m(filepath_or_buffer, kwds)\u001b[0m\n\u001b[1;32m    580\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m parser\n\u001b[1;32m    582\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m parser:\n\u001b[0;32m--> 583\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mparser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnrows\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/io/parsers/readers.py:1704\u001b[0m, in \u001b[0;36mTextFileReader.read\u001b[0;34m(self, nrows)\u001b[0m\n\u001b[1;32m   1697\u001b[0m nrows \u001b[38;5;241m=\u001b[39m validate_integer(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnrows\u001b[39m\u001b[38;5;124m\"\u001b[39m, nrows)\n\u001b[1;32m   1698\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1699\u001b[0m     \u001b[38;5;66;03m# error: \"ParserBase\" has no attribute \"read\"\u001b[39;00m\n\u001b[1;32m   1700\u001b[0m     (\n\u001b[1;32m   1701\u001b[0m         index,\n\u001b[1;32m   1702\u001b[0m         columns,\n\u001b[1;32m   1703\u001b[0m         col_dict,\n\u001b[0;32m-> 1704\u001b[0m     ) \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[attr-defined]\u001b[39;49;00m\n\u001b[1;32m   1705\u001b[0m \u001b[43m        \u001b[49m\u001b[43mnrows\u001b[49m\n\u001b[1;32m   1706\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1707\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n\u001b[1;32m   1708\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclose()\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/io/parsers/c_parser_wrapper.py:234\u001b[0m, in \u001b[0;36mCParserWrapper.read\u001b[0;34m(self, nrows)\u001b[0m\n\u001b[1;32m    232\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    233\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlow_memory:\n\u001b[0;32m--> 234\u001b[0m         chunks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_reader\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_low_memory\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnrows\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    235\u001b[0m         \u001b[38;5;66;03m# destructive to chunks\u001b[39;00m\n\u001b[1;32m    236\u001b[0m         data \u001b[38;5;241m=\u001b[39m _concatenate_chunks(chunks)\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/_libs/parsers.pyx:812\u001b[0m, in \u001b[0;36mpandas._libs.parsers.TextReader.read_low_memory\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/_libs/parsers.pyx:873\u001b[0m, in \u001b[0;36mpandas._libs.parsers.TextReader._read_rows\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/_libs/parsers.pyx:848\u001b[0m, in \u001b[0;36mpandas._libs.parsers.TextReader._tokenize_rows\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/_libs/parsers.pyx:859\u001b[0m, in \u001b[0;36mpandas._libs.parsers.TextReader._check_tokenize_status\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/opt/conda/envs/lang2sql/lib/python3.10/site-packages/pandas/_libs/parsers.pyx:2025\u001b[0m, in \u001b[0;36mpandas._libs.parsers.raise_parser_error\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mParserError\u001b[0m: Error tokenizing data. C error: Calling read(nbytes) on source failed. Try engine='python'."
     ]
    }
   ],
   "source": [
    "chicago_crime = pd.concat((pd.read_csv(path +\"/\" + f) for f in files), ignore_index=True)\n",
    "\n",
    "chicago_crime.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chicago_crime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Source: https://platform.openai.com/examples/default-sql-translate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "\n",
    "Given the following SQL table, your job is to write queries given a user’s request. \\n\n",
    "\n",
    "CREATE TABLE {} ({}) \\n\n",
    "\n",
    "Write a SQL query that returns - {}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sql_prompt_generator(table_name, col_names, query):\n",
    "    prompt = prompt_template.format(table_name, col_names, query)\n",
    "    return prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table = \"chicago_crime\"\n",
    "col_names = str(list(chicago_crime.columns)).replace('[', '').replace(']', '')\n",
    "query = \"How many cases are from the year 2023?\"\n",
    "\n",
    "p = sql_prompt_generator(table_name = table, col_names = col_names, query = query)\n",
    "\n",
    "print(p)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai.api_key = os.getenv('OPENAI_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_api_models = pd.DataFrame(openai.Model.list()[\"data\"])\n",
    "\n",
    "openai_api_models.head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.Completion.create(engine = \"text-davinci-003\",\n",
    "                                     prompt= p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = response[\"choices\"][0][\"text\"]\n",
    "print(query)\n",
    "duckdb.sql(query).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.Completion.create(engine = \"text-davinci-003\",\n",
    "                                     prompt= p,\n",
    "                                     n = 2)\n",
    "print(response)\n",
    "query = response[\"choices\"][0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.Completion.create(engine = \"text-davinci-003\",\n",
    "                                     prompt= p,\n",
    "                                     n = 2,\n",
    "                                     temperature = 0)\n",
    "print(response)\n",
    "query = response[\"choices\"][0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.Completion.create(engine = \"text-davinci-003\",\n",
    "                                     prompt= p,\n",
    "                                     n = 2,\n",
    "                                     temperature = 2)\n",
    "print(response)\n",
    "query = response[\"choices\"][0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lang2sql(api_key, table_name, col_names, query, engine = \"text-davinci-003\", n = 1, temperature = 0):\n",
    "    prompt = sql_prompt_generator(table_name = table_name, col_names = col_names, query = query)\n",
    "    \n",
    "    openai.api_key = api_key\n",
    "\n",
    "    response = openai.Completion.create(engine = engine,\n",
    "                                        prompt= prompt,\n",
    "                                        n = n,\n",
    "                                        temperature = temperature)\n",
    "    return response\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_name = \"chicago_crime\"\n",
    "col_names = str(list(chicago_crime.columns)).replace('[', '').replace(']', '')\n",
    "api_key = os.getenv('OPENAI_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases are from the year 2023?\"\n",
    "\n",
    "response = lang2sql(api_key = api_key, table_name = table_name, col_names= col_names, query = query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)\n",
    "sql_query = response[\"choices\"][0][\"text\"]\n",
    "print(sql_query)\n",
    "\n",
    "duckdb.sql(sql_query).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases ended up with arrest?\"\n",
    "\n",
    "response = lang2sql(api_key = api_key, table_name = table_name, col_names= col_names, query = query)\n",
    "sql_query = response[\"choices\"][0][\"text\"]\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tbl_describe = duckdb.sql(\"DESCRIBE SELECT * FROM chicago_crime;\")\n",
    "tbl_describe\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "col_attr = tbl_describe.df()[[\"column_name\", \"column_type\"]]\n",
    "\n",
    "col_attr[\"column_joint\"] = col_attr[\"column_name\"] + \" \" +  col_attr[\"column_type\"]\n",
    "col_attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "col_names = str(list(col_attr[\"column_joint\"].values)).replace('[', '').replace(']', '').replace('\\'', '')\n",
    "query = \"How many cases ended up with arrest?\"\n",
    "prompt = sql_prompt_generator(table_name = table_name, col_names = col_names, query = query)\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "response = lang2sql(api_key = api_key, table_name = table_name, col_names= col_names, query = query)\n",
    "print(response)\n",
    "sql_query = response[\"choices\"][0][\"text\"]\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = \"\"\"\n",
    "Given the following SQL table, your job is to write queries given a user’s request. \n",
    "\n",
    "\n",
    "CREATE TABLE chicago_crime (ID BIGINT, Case Number VARCHAR, Date VARCHAR, Block VARCHAR, IUCR VARCHAR, Primary Type VARCHAR, Description VARCHAR, Location Description VARCHAR, Arrest BOOLEAN, Domestic BOOLEAN, Beat BIGINT, District BIGINT, Ward DOUBLE, Community Area BIGINT, FBI Code VARCHAR, X Coordinate DOUBLE, Y Coordinate DOUBLE, Year BIGINT, Updated On VARCHAR, Latitude DOUBLE, Longitude DOUBLE, Location VARCHAR) \n",
    "\"\"\"\n",
    "\n",
    "prompt = \"Write a SQL query that returns - How many cases ended up with arrest?\"\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-3.5-turbo\",\n",
    "  messages=[\n",
    "    {\n",
    "      \"role\": \"system\",\n",
    "      \"content\": system\n",
    "    },\n",
    "    {\n",
    "      \"role\": \"user\",\n",
    "      \"content\": prompt\n",
    "    }\n",
    "  ],\n",
    "  temperature=1,\n",
    "  max_tokens=256,\n",
    "  top_p=1,\n",
    "  frequency_penalty=0,\n",
    "  presence_penalty=0\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_query = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(sql_query)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(sql_query).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(\"DESCRIBE SELECT * FROM chicago_crime;\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_message(table_name, query):\n",
    "\n",
    "    class message:\n",
    "        def __init__(message, system, user, column_names, column_attr):\n",
    "            message.system = system\n",
    "            message.user = user\n",
    "            message.column_names = column_names\n",
    "            message.column_attr = column_attr\n",
    "\n",
    "    \n",
    "    system_template = \"\"\"\n",
    "\n",
    "    Given the following SQL table, your job is to write queries given a user’s request. \\n\n",
    "\n",
    "    CREATE TABLE {} ({}) \\n\n",
    "    \"\"\"\n",
    "\n",
    "    user_template = \"Write a SQL query that returns - {}\"\n",
    "    \n",
    "    tbl_describe = duckdb.sql(\"DESCRIBE SELECT * FROM \" + table_name +  \";\")\n",
    "    col_attr = tbl_describe.df()[[\"column_name\", \"column_type\"]]\n",
    "    col_attr[\"column_joint\"] = col_attr[\"column_name\"] + \" \" +  col_attr[\"column_type\"]\n",
    "    col_names = str(list(col_attr[\"column_joint\"].values)).replace('[', '').replace(']', '').replace('\\'', '')\n",
    "\n",
    "    system = system_template.format(table_name, col_names)\n",
    "    user = user_template.format(query)\n",
    "\n",
    "    m = message(system = system, user = user, column_names = col_attr[\"column_name\"], column_attr = col_attr[\"column_type\"])\n",
    "    return m\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases ended up with arrest?\"\n",
    "msg = create_message(table_name = \"chicago_crime\", query = query)\n",
    "\n",
    "print(msg.system)\n",
    "print(msg.user)\n",
    "print(msg.column_names)\n",
    "print(msg.column_attr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases ended up with arrest?\"\n",
    "prompt = create_message(table_name = \"chicago_crime\", query = query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  message = [\n",
    "    {\n",
    "      \"role\": \"system\",\n",
    "      \"content\": prompt.system\n",
    "    },\n",
    "    {\n",
    "      \"role\": \"user\",\n",
    "      \"content\": prompt.user\n",
    "    }\n",
    "    ]\n",
    "\n",
    "message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " response = openai.ChatCompletion.create(\n",
    "        model = \"gpt-3.5-turbo\",\n",
    "        messages = message,\n",
    "        temperature = 0,\n",
    "        max_tokens = 256,\n",
    "        frequency_penalty = 0,\n",
    "        presence_penalty = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "\n",
    "print(sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_quotes(query, col_names):\n",
    "    for i in col_names:\n",
    "        if i in query:\n",
    "            l = query.find(i)\n",
    "            if query[l-1] != \"'\" and query[l-1] != '\"': \n",
    "                query = str(query).replace(i, '\"' + i + '\"') \n",
    "\n",
    "    return(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q1 =\"SELECT COUNT(*) FROM chicago_crime WHERE Primary Type = 'ROBBERY';\"\n",
    "q2 =\"SELECT COUNT(*) FROM chicago_crime WHERE \\\"Primary Type\\\" = 'ROBBERY';\"\n",
    "q3 =\"SELECT COUNT(*) FROM chicago_crime WHERE 'Primary Type' = 'ROBBERY';\"\n",
    "\n",
    "print(add_quotes(query = q1, col_names= prompt.column_names))\n",
    "print(add_quotes(query = q2, col_names= prompt.column_names))\n",
    "print(add_quotes(query = q3, col_names= prompt.column_names))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_quotes(query = sql, col_names = prompt.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lang2sql(api_key, table_name, query, model = \"gpt-3.5-turbo\", temperature = 0, max_tokens = 256, frequency_penalty = 0,presence_penalty= 0):\n",
    "    class response:\n",
    "        def __init__(output, message, response, sql):\n",
    "            output.message = message\n",
    "            output.response = response\n",
    "            output.sql = sql\n",
    "\n",
    "    openai.api_key = api_key\n",
    "\n",
    "    m = create_message(table_name = table_name, query = query)\n",
    "\n",
    "    message = [\n",
    "    {\n",
    "      \"role\": \"system\",\n",
    "      \"content\": m.system\n",
    "    },\n",
    "    {\n",
    "      \"role\": \"user\",\n",
    "      \"content\": m.user\n",
    "    }\n",
    "    ]\n",
    "    \n",
    "    openai_response = openai.ChatCompletion.create(\n",
    "        model = model,\n",
    "        messages = message,\n",
    "        temperature = temperature,\n",
    "        max_tokens = max_tokens,\n",
    "        frequency_penalty = frequency_penalty,\n",
    "        presence_penalty = presence_penalty)\n",
    "    \n",
    "    sql_query = add_quotes(query = openai_response[\"choices\"][0][\"message\"][\"content\"], col_names = m.column_names)\n",
    "\n",
    "    output = response(message = m, response = openai_response, sql = sql_query)\n",
    "\n",
    "    return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "api_key = os.getenv('OPENAI_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases ended up with arrest?\"\n",
    "response = lang2sql(api_key = api_key, table_name = \"chicago_crime\", query = query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response.message)\n",
    "print(response.sql)\n",
    "print(response.response)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases ended up with arrest during 2022\"\n",
    "response = lang2sql(api_key = api_key, table_name = \"chicago_crime\", query = query)\n",
    "print(response.sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "query = \"How many cases ended up with arrest between 2022 and 2023?\"\n",
    "response = lang2sql(api_key = api_key, table_name = \"chicago_crime\", query = query)\n",
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"Summarize the cases by primary type\"\n",
    "response = lang2sql(api_key = api_key, table_name = \"chicago_crime\", query = query)\n",
    "\n",
    "print(response.sql)\n",
    "\n",
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"How many cases is the type of robbery?\"\n",
    "response = lang2sql(api_key = api_key, table_name = \"chicago_crime\", query = query)\n",
    "\n",
    "print(response.sql)\n",
    "\n",
    "duckdb.sql(response.sql).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duckdb.sql(response.sql).show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lang2sql",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}