5. Code Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e0240d1c",
   "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": "6e8f47ee",
   "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": "85ba323f",
   "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": "54272170",
   "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.5.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: 2f98fe4b5aabf6921b7e3b8fba71f74cb0ebead8\n",
      "\n",
      "json      : 2.0.9\n",
      "watermark : 2.4.3\n",
      "tqdm      : 4.66.4\n",
      "pandas    : 2.1.4\n",
      "sqlite3   : 2.6.0\n",
      "matplotlib: 3.8.0\n",
      "numpy     : 1.26.4\n",
      "openai    : 1.30.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c16249cb",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "17819e90",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('d4sci.mplstyle')\n",
    "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a62c637",
   "metadata": {},
   "source": [
    "# Text to Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5c9a949c",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8af69e2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "        {\"role\": \"system\", \"content\": \"\"\"You are a grumpy but expert Python software engineer \n",
    "        thats 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": 10,
   "id": "2b4265bb",
   "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": "7c9426be",
   "metadata": {},
   "source": [
    "Which produces the expected result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d7321188",
   "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": "b3951f88",
   "metadata": {},
   "source": [
    "and works as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "818d2467",
   "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": 14,
   "id": "bcee4c40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fibonacci(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4670c71",
   "metadata": {},
   "source": [
    "Let us define a utility function to make sequential queries easier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5fcb17a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(messages, prompt):\n",
    "    messages.append({\"role\":\"user\", \"content\":prompt})\n",
    "    \n",
    "    response = client.chat.completions.create(\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": "00e4e334",
   "metadata": {},
   "source": [
    "# Adding comments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3ae7b281",
   "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": 17,
   "id": "50cb02c1",
   "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, repeatedly breaking down the problem until it reaches the base case.\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"What is the purpose of recursion in this piece of code?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d3325006",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'system',\n",
       "  'content': 'You are a grumpy but expert Python software engineer \\n        thats interviewing for a job. Please be as concise with your answers as possible.'},\n",
       " {'role': 'user',\n",
       "  'content': \"Create a recursive Python function to compute \\n        Fibonacci numbers. Don't provide any explanation, just the code\"},\n",
       " {'role': 'user', 'content': 'Can you add comments to this function?'},\n",
       " ChatCompletionMessage(content='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))', role='assistant', function_call=None, tool_calls=None),\n",
       " {'role': 'user',\n",
       "  'content': 'What is the purpose of recursion in this piece of code?'},\n",
       " ChatCompletionMessage(content='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, repeatedly breaking down the problem until it reaches the base case.', role='assistant', function_call=None, tool_calls=None)]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67095e65",
   "metadata": {},
   "source": [
    "# Explaining Code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27c2a8f1",
   "metadata": {},
   "source": [
    "Let's use a relatively small python script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1f452cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_text = \"\".join(open(\"data/EpiModel.py\").readlines())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "69b41124",
   "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": 21,
   "id": "d9cc593b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This code defines a class `EpiModel` for simulating and analyzing epidemic models. The class allows you to define different compartments (like Susceptible, Infected, Recovered), add interactions between them, and simulate the spread of a disease over time. \n",
      "\n",
      "The `simulate` method stochastically simulates the epidemic model over a given number of timesteps, while the `integrate` method numerically integrates the model over time. \n",
      "\n",
      "The `plot` method provides a way to visualize the results of the simulation, and the `R0` method calculates the basic reproduction number of the disease, which is a key parameter in epidemiology.\n",
      "\n",
      "The code at the end of the script creates an instance of the `EpiModel` class, defines a Susceptible-Infected-Recovered (SIR) model with vaccination, and simulates it over a year. The results are then plotted and saved to a file.\n",
      "CPU times: user 12.4 ms, sys: 2.81 ms, total: 15.2 ms\n",
      "Wall time: 12.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Please explain what this piece of code does: ```%s```\" % code_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a48b1612",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sure, here are the docstrings for each method in the `EpiModel` class:\n",
      "\n",
      "```python\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",
      "        \"\"\"\n",
      "        Initialize the EpiModel object.\n",
      "\n",
      "        Args:\n",
      "            compartments (list): A list of compartments for the model.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def add_interaction(self, source, target, agent, rate):        \n",
      "        \"\"\"\n",
      "        Add an interaction between compartments.\n",
      "\n",
      "        Args:\n",
      "            source (str): The source compartment.\n",
      "            target (str): The target compartment.\n",
      "            agent (str): The agent that causes the transition.\n",
      "            rate (float): The rate of the transition.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def add_spontaneous(self, source, target, rate):\n",
      "        \"\"\"\n",
      "        Add a spontaneous transition.\n",
      "\n",
      "        Args:\n",
      "            source (str): The source compartment.\n",
      "            target (str): The target compartment.\n",
      "            rate (float): The rate of the transition.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def add_vaccination(self, source, target, rate, start):\n",
      "        \"\"\"\n",
      "        Add a vaccination transition.\n",
      "\n",
      "        Args:\n",
      "            source (str): The source compartment.\n",
      "            target (str): The target compartment.\n",
      "            rate (float): The rate of the transition.\n",
      "            start (int): The start time of the vaccination.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def _new_cases(self, population, time, pos):\n",
      "        \"\"\"\n",
      "        Calculate the number of new cases.\n",
      "\n",
      "        Args:\n",
      "            population (np.array): The current population in each compartment.\n",
      "            time (int): The current time.\n",
      "            pos (dict): A dictionary mapping compartment names to indices.\n",
      "\n",
      "        Returns:\n",
      "            np.array: The number of new cases in each compartment.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def plot(self, title=None, normed=True, **kwargs):\n",
      "        \"\"\"\n",
      "        Plot the results of the simulation.\n",
      "\n",
      "        Args:\n",
      "            title (str): The title of the plot.\n",
      "            normed (bool): Whether to normalize the results.\n",
      "            **kwargs: Additional keyword arguments for the plot function.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"\n",
      "        Stochastically simulate the epidemic model.\n",
      "\n",
      "        Args:\n",
      "            timesteps (int): The number of timesteps to simulate.\n",
      "            t_min (int): The start time of the simulation.\n",
      "            seasonality (np.array): The seasonality of the disease.\n",
      "            **kwargs: Initial population in each compartment.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"\n",
      "        Numerically integrate the epidemic model.\n",
      "\n",
      "        Args:\n",
      "            timesteps (int): The number of timesteps to integrate.\n",
      "            t_min (int): The start time of the integration.\n",
      "            seasonality (np.array): The seasonality of the disease.\n",
      "            **kwargs: Initial population in each compartment.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "\n",
      "    def R0(self):\n",
      "        \"\"\"\n",
      "        Calculate the basic reproduction number of the disease.\n",
      "\n",
      "        Returns:\n",
      "            float: The basic reproduction number.\n",
      "        \"\"\"\n",
      "        # ...\n",
      "```\n",
      "\n",
      "Note that I've omitted the docstrings for the private methods and the `__getattr__` and `__getitem__` methods, as these are typically not meant to be used directly by the user.\n",
      "CPU times: user 15.1 ms, sys: 3.09 ms, total: 18.2 ms\n",
      "Wall time: 36.8 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Can you please add a doc string to each function and method? Please include information about each argument of the function\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07dcd290",
   "metadata": {},
   "source": [
    "# Interacting with a database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2a61149",
   "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": 23,
   "id": "296572a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "con = sqlite3.connect(\"data/Northwind_small.sqlite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f2b0b8df",
   "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": 27,
   "id": "9744f222",
   "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": 28,
   "id": "55e352c6",
   "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": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(query_sql, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "1aa2f05f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT EmployeeTerritory.TerritoryId, COUNT(EmployeeTerritory.EmployeeId) as EmployeeCount\n",
      "FROM EmployeeTerritory\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": 31,
   "id": "593a3709",
   "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": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "227d9dc8",
   "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": 33,
   "id": "04ed4302",
   "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": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d14d897c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'system',\n",
       "  'content': \"You're a Database Administrator. \\n        Please generate SQL queries to answer the following questions. \\n        No comments are necessary.\"},\n",
       " {'role': 'user',\n",
       "  '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",
       " {'role': 'user',\n",
       "  'content': 'Generate a table with employee first name, last name and territory id'},\n",
       " ChatCompletionMessage(content='SELECT Employee.FirstName, Employee.LastName, EmployeeTerritory.TerritoryId\\nFROM Employee\\nJOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId;', role='assistant', function_call=None, tool_calls=None),\n",
       " {'role': 'user',\n",
       "  'content': 'Compute how many employees work in each territory'},\n",
       " ChatCompletionMessage(content='SELECT EmployeeTerritory.TerritoryId, COUNT(Employee.Id) as EmployeeCount\\nFROM Employee\\nJOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId\\nGROUP BY EmployeeTerritory.TerritoryId;', role='assistant', function_call=None, tool_calls=None),\n",
       " {'role': 'user',\n",
       "  'content': 'Compute how many employees work in each territory'},\n",
       " ChatCompletionMessage(content='SELECT EmployeeTerritory.TerritoryId, COUNT(EmployeeTerritory.EmployeeId) as EmployeeCount\\nFROM EmployeeTerritory\\nGROUP BY EmployeeTerritory.TerritoryId;', role='assistant', function_call=None, tool_calls=None),\n",
       " {'role': 'user', 'content': 'How many shippers do we work with?'},\n",
       " ChatCompletionMessage(content='SELECT COUNT(*) FROM Shipper;', role='assistant', function_call=None, tool_calls=None)]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d59f320",
   "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
}