5. Code Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a918a1c",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "    <div style=\"width: 150px; float: left;\"> <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_ball.png\" alt=\"Data For Science, Inc\" align=\"left\" border=\"0\" width=150px> </div>\n",
    "    <div style=\"float: left; margin-left: 10px;\"> <h1>Generative AI with OpenAI API</h1>\n",
    "<h1>Code Generation</h1>\n",
    "        <p>Bruno Gonçalves<br/>\n",
    "        <a href=\"http://www.data4sci.com/\">www.data4sci.com</a><br/>\n",
    "            @bgoncalves, @data4sci</p></div>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "598bbe29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from pprint import pprint\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import sqlite3\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "from ipywidgets import interact\n",
    "import openai\n",
    "from openai import OpenAI\n",
    "\n",
    "import os\n",
    "import gzip\n",
    "\n",
    "import tqdm as tq\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18228fa",
   "metadata": {},
   "source": [
    "We start by printing out the versions of the libraries we're using for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "203457b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python implementation: CPython\n",
      "Python version       : 3.11.7\n",
      "IPython version      : 8.12.3\n",
      "\n",
      "Compiler    : Clang 14.0.6 \n",
      "OS          : Darwin\n",
      "Release     : 23.4.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: 5217868ae783ab0315a9d79f2e1ac0fb25095003\n",
      "\n",
      "tqdm      : 4.66.2\n",
      "matplotlib: 3.8.0\n",
      "watermark : 2.4.3\n",
      "pandas    : 2.1.4\n",
      "openai    : 1.20.0\n",
      "numpy     : 1.26.4\n",
      "sqlite3   : 2.6.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e03745",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0d920890",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('d4sci.mplstyle')\n",
    "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8415402",
   "metadata": {},
   "source": [
    "# Text to Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cb551e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7f253a1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "        {\"role\": \"system\", \"content\": \"\"\"You are a grumpy but expert Python programmer \n",
    "        that interviewing for a job. Please be as concise with your answers as possible.\"\"\"},\n",
    "        {\"role\": \"user\", \"content\": \"\"\"Create a recursive Python function to compute \n",
    "        Fibonacci numbers. Don't provide any explanation, just the code\"\"\"},\n",
    "  ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e2ed698d",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=\"gpt-4\",\n",
    "    messages=messages,\n",
    "    temperature=0,\n",
    "    max_tokens=1024\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b6ef9d5",
   "metadata": {},
   "source": [
    "Which produces the expected result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1083e1a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def fibonacci(n):\n",
      "    if n <= 1:\n",
      "       return n\n",
      "    else:\n",
      "       return(fibonacci(n-1) + fibonacci(n-2))\n"
     ]
    }
   ],
   "source": [
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e07705aa",
   "metadata": {},
   "source": [
    "and works as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "84763859",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci(n):\n",
    "    if n <= 1:\n",
    "        return n\n",
    "    else:\n",
    "        return fibonacci(n-1) + fibonacci(n-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "915cd1fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "55"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fibonacci(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c07cb02",
   "metadata": {},
   "source": [
    "Let us define a utility function to make sequential queries easier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "827509d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(messages, prompt):\n",
    "    messages.append({\"role\":\"user\", \"content\":prompt})\n",
    "    \n",
    "    response = client.chat.completions.create(\n",
    "        # GPT-3.5-Turbo now incorporates the Codex functionality\n",
    "        model=\"gpt-4\",\n",
    "        messages=messages,\n",
    "        temperature=0,\n",
    "        max_tokens=1024\n",
    "    )\n",
    "    \n",
    "    messages.append(response.choices[0].message)\n",
    "    \n",
    "    return messages[-1].content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba151d34",
   "metadata": {},
   "source": [
    "# Adding comments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b9490d53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def fibonacci(n):\n",
      "    # Base case: if n is 0 or 1, return n\n",
      "    if n <= 1:\n",
      "       return n\n",
      "    # Recursive case: return the sum of the previous two Fibonacci numbers\n",
      "    else:\n",
      "       return(fibonacci(n-1) + fibonacci(n-2))\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"Can you add comments to this function?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e5ca8a67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recursion in this code is used to break down the problem of calculating a Fibonacci number into smaller, simpler problems. It allows the function to call itself with different arguments, computing the smaller Fibonacci numbers first, then adding them together to get the larger Fibonacci number.\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"What is the purpose of recursion in this piece of code?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f697b8",
   "metadata": {},
   "source": [
    "# Explaining Code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c2a4416",
   "metadata": {},
   "source": [
    "Let's use a relatively small python script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b10ad422",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_text = \"\".join(open(\"data/EpiModel.py\").readlines())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "26ccb74f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### −∗− mode : python ; −∗−\n",
      "# @file EpiModel.py\n",
      "# @author Bruno Goncalves\n",
      "######################################################\n",
      "\n",
      "import networkx as nx\n",
      "import numpy as np\n",
      "from numpy import linalg\n",
      "from numpy import random\n",
      "import scipy.integrate\n",
      "import pandas as pd\n",
      "import matplotlib.pyplot as plt\n",
      "\n",
      "from tqdm import tqdm\n",
      "tqdm.pandas()\n",
      "\n",
      "class EpiModel(object):\n",
      "    \"\"\"Simple Epidemic Model Implementation\n",
      "    \n",
      "        Provides a way to implement and numerically integrate \n",
      "    \"\"\"\n",
      "    def __init__(self, compartments=None):\n",
      "        self.transitions = nx.MultiDiGraph()\n",
      "        self.seasonality = None\n",
      "        \n",
      "        if compartments is not None:\n",
      "            self.transitions.add_nodes_from([comp for comp in compartments])\n",
      "    \n",
      "    def add_interaction(self, source, target, agent, rate):        \n",
      "        self.transitions.add_edge(source, target, agent=agent, rate=rate)        \n",
      "        \n",
      "    def add_spontaneous(self, source, target, rate):\n",
      "        self.transitions.add_edge(source, target, rate=rate)\n",
      "\n",
      "    def add_vaccination(self, source, target, rate, start):\n",
      "        self.transitions.add_edge(source, target, rate=rate, start=start)\n",
      "        \n",
      "    def _new_cases(self, population, time, pos):\n",
      "        \"\"\"Internal function used by integration routine\"\"\"\n",
      "        diff = np.zeros(len(pos))\n",
      "        N = np.sum(population)        \n",
      "        \n",
      "        for edge in self.transitions.edges(data=True):\n",
      "            source = edge[0]\n",
      "            target = edge[1]\n",
      "            trans = edge[2]\n",
      "            \n",
      "            rate = trans['rate']*population[pos[source]]\n",
      "            \n",
      "            if 'start' in trans and trans['start'] >= time:\n",
      "                continue\n",
      "\n",
      "            if 'agent' in trans:\n",
      "                agent = trans['agent']\n",
      "                rate *= population[pos[agent]]/N\n",
      "\n",
      "                if self.seasonality is not None:\n",
      "                    curr_t = int(time)%365\n",
      "                    season = float(self.seasonality[curr_t])\n",
      "                    rate *= season\n",
      "                \n",
      "            diff[pos[source]] -= rate\n",
      "            diff[pos[target]] += rate\n",
      "            \n",
      "        return diff\n",
      "    \n",
      "    def plot(self, title=None, normed=True, **kwargs):\n",
      "        \"\"\"Convenience function for plotting\"\"\"\n",
      "        try:\n",
      "            if normed:\n",
      "                N = self.values_.iloc[0].sum()\n",
      "                ax = (self.values_/N).plot(**kwargs)\n",
      "            else:\n",
      "                ax = self.values_.plot(**kwargs)\n",
      "                \n",
      "            ax.set_xlabel('Time')\n",
      "            ax.set_ylabel('Population')\n",
      "            \n",
      "            if title is not None:\n",
      "                ax.set_title(title)\n",
      "            \n",
      "            return ax\n",
      "        except:\n",
      "            raise NotInitialized('You must call integrate() first')\n",
      "    \n",
      "    def __getattr__(self, name):\n",
      "        \"\"\"Dynamic method to return the individual compartment values\"\"\"\n",
      "        if 'values_' in self.__dict__:\n",
      "            return self.values_[name]\n",
      "        else:\n",
      "            raise AttributeError(\"'EpiModel' object has no attribute '%s'\" % name)\n",
      "\n",
      "    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Stochastically simulate the epidemic model\"\"\"\n",
      "        pos = {comp: i for i, comp in enumerate(self.transitions.nodes())}\n",
      "        population=np.zeros(len(pos), dtype='int')\n",
      "\n",
      "        for comp in kwargs:\n",
      "            population[pos[comp]] = kwargs[comp]\n",
      "\n",
      "        values = []\n",
      "        values.append(population)\n",
      "\n",
      "        comps = list(self.transitions.nodes)\n",
      "        time = np.arange(t_min, t_min+timesteps, 1, dtype='int')\n",
      "\n",
      "        self.seasonality = seasonality\n",
      "\n",
      "        for t in time:\n",
      "            pop = values[-1]\n",
      "            new_pop = values[-1].copy()\n",
      "            N = np.sum(pop)\n",
      "\n",
      "\n",
      "            for comp in comps:\n",
      "                trans = list(self.transitions.edges(comp, data=True))             \n",
      "\n",
      "                prob = np.zeros(len(comps), dtype='float')\n",
      "\n",
      "                for _, node_j, data in trans:\n",
      "                    source = pos[comp]\n",
      "                    target = pos[node_j]\n",
      "\n",
      "                    rate = data['rate']\n",
      "\n",
      "                    if 'start' in data and data['start'] >= t:\n",
      "                        continue\n",
      "\n",
      "                    if 'agent' in data:\n",
      "                        agent = pos[data['agent']]\n",
      "                        rate *= pop[agent]/N\n",
      "\n",
      "                        if self.seasonality is not None:\n",
      "                            curr_t = int(t)%365\n",
      "                            season = float(self.seasonality[curr_t])\n",
      "                            rate *= season\n",
      "\n",
      "                    prob[target] = rate\n",
      "\n",
      "                prob[source] = 1-np.sum(prob)\n",
      "\n",
      "                delta = random.multinomial(pop[source], prob)\n",
      "                delta[source] = 0\n",
      "\n",
      "                changes = np.sum(delta)\n",
      "\n",
      "                if changes == 0:\n",
      "                    continue\n",
      "\n",
      "                new_pop[source] -= changes\n",
      "\n",
      "                for i in range(len(delta)):\n",
      "                    new_pop[i] += delta[i]\n",
      "\n",
      "            values.append(new_pop)\n",
      "\n",
      "        values = np.array(values)\n",
      "        self.values_ = pd.DataFrame(values[1:], columns=comps, index=time)\n",
      "    \n",
      "    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Numerically integrate the epidemic model\"\"\"\n",
      "        pos = {comp: i for i, comp in enumerate(self.transitions.nodes())}\n",
      "        population=np.zeros(len(pos))\n",
      "        \n",
      "        for comp in kwargs:\n",
      "            population[pos[comp]] = kwargs[comp]\n",
      "        \n",
      "        time = np.arange(t_min, t_min+timesteps, 1)\n",
      "\n",
      "        self.seasonality = seasonality\n",
      "        self.values_ = pd.DataFrame(scipy.integrate.odeint(self._new_cases, population, time, args=(pos,)), columns=pos.keys(), index=time)\n",
      "\n",
      "    def __repr__(self):\n",
      "        text = 'Epidemic Model with %u compartments and %u transitions:\\n\\n' % \\\n",
      "              (self.transitions.number_of_nodes(), \n",
      "               self.transitions.number_of_edges())\n",
      "        \n",
      "        for edge in self.transitions.edges(data=True):\n",
      "            source = edge[0]\n",
      "            target = edge[1]\n",
      "            trans = edge[2]\n",
      "            \n",
      "            rate = trans['rate']\n",
      "\n",
      "            if 'agent' in trans:\n",
      "                agent = trans['agent']\n",
      "                text += \"%s + %s = %s %f\\n\" % (source, agent, target, rate)\n",
      "            elif 'start' in trans:\n",
      "                start = trans['start']\n",
      "                text+=\"%s -> %s %f starting at %s days\\n\" % (source, target, rate, start)\n",
      "            else:\n",
      "                text+=\"%s -> %s %f\\n\" % (source, target, rate)\n",
      "        \n",
      "        R0 = self.R0()\n",
      "\n",
      "        if R0 is not None:\n",
      "            text += \"\\nR0=%1.2f\" % R0\n",
      "\n",
      "        return text\n",
      "\n",
      "    def _get_active(self):\n",
      "        active = set()\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                active.add(data['agent'])\n",
      "            else:\n",
      "                active.add(node_i)\n",
      "\n",
      "        return active\n",
      "\n",
      "    def _get_susceptible(self):\n",
      "        susceptible = set([node for node, deg in self.transitions.in_degree() if deg==0])\n",
      "\n",
      "        if len(susceptible) == 0:\n",
      "            for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "                if \"agent\" in data:\n",
      "                    susceptible.add(node_i)\n",
      "\n",
      "        return susceptible\n",
      "\n",
      "    def _get_infections(self):\n",
      "        inf = {}\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                agent = data['agent']\n",
      "\n",
      "                if agent not in inf:\n",
      "                    inf[agent] = {}\n",
      "\n",
      "                if node_i not in inf[agent]:\n",
      "                    inf[agent][node_i] = {}\n",
      "\n",
      "                inf[agent][node_i]['target'] = node_j\n",
      "                inf[agent][node_i]['rate'] = data['rate']\n",
      "\n",
      "        return inf\n",
      "\n",
      "\n",
      "    def R0(self):\n",
      "        infected = set()\n",
      "\n",
      "        susceptible = self._get_susceptible()\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                infected.add(data['agent'])\n",
      "                infected.add(node_j)\n",
      "\n",
      "\n",
      "        infected = sorted(infected)\n",
      "        N_infected = len(infected)\n",
      "\n",
      "        F = np.zeros((N_infected, N_infected), dtype='float')\n",
      "        V = np.zeros((N_infected, N_infected), dtype='float')\n",
      "\n",
      "        pos = dict(zip(infected, np.arange(N_infected)))\n",
      "\n",
      "        try:\n",
      "            for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "                rate = data['rate']\n",
      "\n",
      "                if \"agent\" in data:\n",
      "                    target = pos[node_j]\n",
      "                    agent = pos[data['agent']]\n",
      "\n",
      "                    if node_i in susceptible:\n",
      "                        F[target, agent] = rate\n",
      "                elif \"start\" in data:\n",
      "                    continue\n",
      "                else:\n",
      "                    source = pos[node_i]\n",
      "\n",
      "                    V[source, source] += rate\n",
      "\n",
      "                    if node_j in pos:\n",
      "                        target = pos[node_j]\n",
      "                        V[target, source] -= rate\n",
      "        \n",
      "            eig, v = linalg.eig(np.dot(F, linalg.inv(V)))\n",
      "\n",
      "            return eig.max()\n",
      "        except:\n",
      "            return None\n",
      "\n",
      "    def __getitem__(self, bla):\n",
      "        return self.values_[bla]\n",
      "\n",
      "if __name__ == '__main__':\n",
      "\n",
      "    beta = 0.2\n",
      "    mu = 0.1\n",
      "\n",
      "    SIR = EpiModel()\n",
      "    SIR.add_interaction('S', 'I', 'I', beta)\n",
      "    SIR.add_spontaneous('I', 'R', mu)\n",
      "    SIR.add_vaccination('S', 'V', 0.01, 75)\n",
      "    SIR.add_spontaneous('VI', 'VR', mu)\n",
      "    SIR.add_interaction('V', 'VI', 'I', beta*(1-.8))\n",
      "\n",
      "\n",
      "    print(SIR)\n",
      "\n",
      "    N = 100000\n",
      "    I0 = 10  \n",
      "\n",
      "    season = np.ones(365+1)\n",
      "    season[74:100] = 0.25\n",
      "\n",
      "    fig, ax = plt.subplots(1)\n",
      "\n",
      "    Nruns = 1000\n",
      "    values = []\n",
      "\n",
      "    #for i in tqdm(range(Nruns), total=Nruns):\n",
      "    SIR.integrate(365, S=.3*N-10, I=10, V=.7*N)\n",
      "\n",
      "    SIR[['I', 'VI', 'VR', 'R']].plot(ax=ax)\n",
      "    print(SIR.S.tail())\n",
      "    #ax.plot(SIR.I/N, lw=.1, c='b')\n",
      "\n",
      "    if SIR.I.max() > 10:\n",
      "        values.append(SIR.I)\n",
      "\n",
      "    values = pd.DataFrame(values)\n",
      "    (values.median(axis=0)/N).plot(ax=ax, c='r')\n",
      "\n",
      "    fig.savefig('SIR.png')\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(code_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "914d4a80",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This code defines a class `EpiModel` for simulating and analyzing epidemic models. The class allows for the creation of a model with multiple compartments (like susceptible, infected, and recovered in an SIR model), and transitions between these compartments. \n",
      "\n",
      "The transitions can be interactions between compartments (like susceptible and infected individuals interacting to produce more infected), spontaneous transitions (like infected individuals recovering), or vaccinations (moving individuals from susceptible to vaccinated). \n",
      "\n",
      "The class also includes methods for simulating the epidemic over time, either stochastically or by numerical integration, and for plotting the results. \n",
      "\n",
      "The `R0` method calculates the basic reproduction number of the epidemic, which is the expected number of secondary infections produced by a single infected individual in a completely susceptible population.\n",
      "\n",
      "The code at the end of the script creates an instance of the `EpiModel` class for an SIR model with vaccination, simulates it over a year, and plots the results.\n",
      "CPU times: user 8.99 ms, sys: 1.91 ms, total: 10.9 ms\n",
      "Wall time: 8.64 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Please explain what this piece of code does: ```%s```\" % code_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "27f6a518",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sure, here are the methods with added docstrings:\n",
      "\n",
      "```python\n",
      "class EpiModel(object):\n",
      "    \"\"\"Simple Epidemic Model Implementation\"\"\"\n",
      "    \n",
      "    def __init__(self, compartments=None):\n",
      "        \"\"\"Initialize the model with compartments if provided.\"\"\"\n",
      "        \n",
      "    def add_interaction(self, source, target, agent, rate):        \n",
      "        \"\"\"Add an interaction transition between compartments.\"\"\"\n",
      "        \n",
      "    def add_spontaneous(self, source, target, rate):\n",
      "        \"\"\"Add a spontaneous transition between compartments.\"\"\"\n",
      "        \n",
      "    def add_vaccination(self, source, target, rate, start):\n",
      "        \"\"\"Add a vaccination transition between compartments.\"\"\"\n",
      "        \n",
      "    def _new_cases(self, population, time, pos):\n",
      "        \"\"\"Calculate the number of new cases for each compartment.\"\"\"\n",
      "        \n",
      "    def plot(self, title=None, normed=True, **kwargs):\n",
      "        \"\"\"Plot the population in each compartment over time.\"\"\"\n",
      "        \n",
      "    def __getattr__(self, name):\n",
      "        \"\"\"Return the population in the compartment specified by name.\"\"\"\n",
      "        \n",
      "    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Simulate the epidemic stochastically over the specified number of timesteps.\"\"\"\n",
      "        \n",
      "    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Integrate the epidemic numerically over the specified number of timesteps.\"\"\"\n",
      "        \n",
      "    def __repr__(self):\n",
      "        \"\"\"Return a string representation of the model.\"\"\"\n",
      "        \n",
      "    def _get_active(self):\n",
      "        \"\"\"Get the active compartments in the model.\"\"\"\n",
      "        \n",
      "    def _get_susceptible(self):\n",
      "        \"\"\"Get the susceptible compartments in the model.\"\"\"\n",
      "        \n",
      "    def _get_infections(self):\n",
      "        \"\"\"Get the infection transitions in the model.\"\"\"\n",
      "        \n",
      "    def R0(self):\n",
      "        \"\"\"Calculate the basic reproduction number of the epidemic.\"\"\"\n",
      "        \n",
      "    def __getitem__(self, bla):\n",
      "        \"\"\"Return the population in the compartment specified by bla.\"\"\"\n",
      "```\n",
      "\n",
      "Please note that these are brief descriptions. For a more detailed explanation of what each method does, you would need to look at the code or provide a more detailed docstring.\n",
      "CPU times: user 13.7 ms, sys: 2.43 ms, total: 16.1 ms\n",
      "Wall time: 20 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Can you please add a doc string to each function and method?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3e196e3",
   "metadata": {},
   "source": [
    "# Interacting with a database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f680e5",
   "metadata": {},
   "source": [
    "Let us open a small test database. This file was downloaded from https://github.com/chineseballer06/Statistical-Analysis-of-Northwind-Database/blob/master/Northwind_small.sqlite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d817a24f",
   "metadata": {},
   "outputs": [],
   "source": [
    "con = sqlite3.connect(\"data/Northwind_small.sqlite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ebb6f516",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"\"\"You're a Database Administrator. \n",
    "    Please generate SQL queries to answer the following questions. \n",
    "    No comments are necessary.\"\"\"},\n",
    "    {\"role\": \"user\", \"content\": \"\"\"\n",
    "# Table Employee, columns = [Id, LastName, First Name]\n",
    "# Table Shipper, columns = [Id, CompanyName, Phone]\n",
    "# Table OrderDetail, columns = [OrderId, ProductId, Quantity]\n",
    "# Table EmployeeTerritory, columns = [Id, EmployeeId, TerritoryId]\n",
    "    \"\"\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "75d49715",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT Employee.FirstName, Employee.LastName, EmployeeTerritory.TerritoryId\n",
      "FROM Employee\n",
      "JOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId;\n"
     ]
    }
   ],
   "source": [
    "query_sql = chat(messages, \"Generate a table with employee first name, last name and territory id\")\n",
    "print(query_sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4628cef7",
   "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>FirstName</th>\n",
       "      <th>LastName</th>\n",
       "      <th>TerritoryId</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Nancy</td>\n",
       "      <td>Davolio</td>\n",
       "      <td>06897</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Nancy</td>\n",
       "      <td>Davolio</td>\n",
       "      <td>19713</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02139</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02184</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>40222</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>30346</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>31406</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>32859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>33607</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>20852</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>27403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>27511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>02903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>07960</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>08837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>10019</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>10038</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>11747</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>14450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>85014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>85251</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98004</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98052</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98104</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>60179</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>60601</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>80202</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>80909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>90405</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>94025</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>94105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95008</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95054</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95060</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>19428</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>44122</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>45839</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>53404</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>03049</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>03801</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>44</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48075</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48084</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48304</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>55113</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>55439</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   FirstName   LastName TerritoryId\n",
       "0      Nancy    Davolio       06897\n",
       "1      Nancy    Davolio       19713\n",
       "2     Andrew     Fuller       01581\n",
       "3     Andrew     Fuller       01730\n",
       "4     Andrew     Fuller       01833\n",
       "5     Andrew     Fuller       02116\n",
       "6     Andrew     Fuller       02139\n",
       "7     Andrew     Fuller       02184\n",
       "8     Andrew     Fuller       40222\n",
       "9      Janet  Leverling       30346\n",
       "10     Janet  Leverling       31406\n",
       "11     Janet  Leverling       32859\n",
       "12     Janet  Leverling       33607\n",
       "13  Margaret    Peacock       20852\n",
       "14  Margaret    Peacock       27403\n",
       "15  Margaret    Peacock       27511\n",
       "16    Steven   Buchanan       02903\n",
       "17    Steven   Buchanan       07960\n",
       "18    Steven   Buchanan       08837\n",
       "19    Steven   Buchanan       10019\n",
       "20    Steven   Buchanan       10038\n",
       "21    Steven   Buchanan       11747\n",
       "22    Steven   Buchanan       14450\n",
       "23   Michael     Suyama       85014\n",
       "24   Michael     Suyama       85251\n",
       "25   Michael     Suyama       98004\n",
       "26   Michael     Suyama       98052\n",
       "27   Michael     Suyama       98104\n",
       "28    Robert       King       60179\n",
       "29    Robert       King       60601\n",
       "30    Robert       King       80202\n",
       "31    Robert       King       80909\n",
       "32    Robert       King       90405\n",
       "33    Robert       King       94025\n",
       "34    Robert       King       94105\n",
       "35    Robert       King       95008\n",
       "36    Robert       King       95054\n",
       "37    Robert       King       95060\n",
       "38     Laura   Callahan       19428\n",
       "39     Laura   Callahan       44122\n",
       "40     Laura   Callahan       45839\n",
       "41     Laura   Callahan       53404\n",
       "42      Anne  Dodsworth       03049\n",
       "43      Anne  Dodsworth       03801\n",
       "44      Anne  Dodsworth       48075\n",
       "45      Anne  Dodsworth       48084\n",
       "46      Anne  Dodsworth       48304\n",
       "47      Anne  Dodsworth       55113\n",
       "48      Anne  Dodsworth       55439"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(query_sql, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cbaf3b33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT EmployeeTerritory.TerritoryId, COUNT(Employee.Id) as EmployeeCount\n",
      "FROM Employee\n",
      "JOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId\n",
      "GROUP BY EmployeeTerritory.TerritoryId;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"Compute how many employees work in each territory\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9631bd51",
   "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>TerritoryId</th>\n",
       "      <th>EmployeeCount</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>01581</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>01730</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>01833</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>02116</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>02139</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>02184</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>02903</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>03049</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>03801</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>06897</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>07960</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>08837</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>10019</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>10038</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>11747</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>14450</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>19428</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>19713</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>20852</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>27403</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>27511</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>30346</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>31406</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>32859</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>33607</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>40222</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>44122</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>45839</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>48075</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>48084</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>48304</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td>53404</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>55113</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>55439</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>60179</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>60601</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>80202</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>80909</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>85014</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>85251</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>90405</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>94025</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>94105</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>95008</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>44</th>\n",
       "      <td>95054</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>95060</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>98004</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>98052</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>98104</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   TerritoryId  EmployeeCount\n",
       "0        01581              1\n",
       "1        01730              1\n",
       "2        01833              1\n",
       "3        02116              1\n",
       "4        02139              1\n",
       "5        02184              1\n",
       "6        02903              1\n",
       "7        03049              1\n",
       "8        03801              1\n",
       "9        06897              1\n",
       "10       07960              1\n",
       "11       08837              1\n",
       "12       10019              1\n",
       "13       10038              1\n",
       "14       11747              1\n",
       "15       14450              1\n",
       "16       19428              1\n",
       "17       19713              1\n",
       "18       20852              1\n",
       "19       27403              1\n",
       "20       27511              1\n",
       "21       30346              1\n",
       "22       31406              1\n",
       "23       32859              1\n",
       "24       33607              1\n",
       "25       40222              1\n",
       "26       44122              1\n",
       "27       45839              1\n",
       "28       48075              1\n",
       "29       48084              1\n",
       "30       48304              1\n",
       "31       53404              1\n",
       "32       55113              1\n",
       "33       55439              1\n",
       "34       60179              1\n",
       "35       60601              1\n",
       "36       80202              1\n",
       "37       80909              1\n",
       "38       85014              1\n",
       "39       85251              1\n",
       "40       90405              1\n",
       "41       94025              1\n",
       "42       94105              1\n",
       "43       95008              1\n",
       "44       95054              1\n",
       "45       95060              1\n",
       "46       98004              1\n",
       "47       98052              1\n",
       "48       98104              1"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "ac90e35c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT COUNT(*) FROM Shipper;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"How many shippers do we work with?\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "beb8c49e",
   "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>COUNT(*)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   COUNT(*)\n",
       "0         3"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b35326",
   "metadata": {},
   "source": [
    "<center>\n",
    "     <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_full.png\" alt=\"Data For Science, Inc\" align=\"center\" border=\"0\" width=300px> \n",
    "</center>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}