5. Code Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4373962c",
   "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": "6e7b952f",
   "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": "ba391b25",
   "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": "b6fd90cd",
   "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     : 24.3.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: 3a7a9a8b6856eb5855cd2ac76a384e203382ab54\n",
      "\n",
      "watermark : 2.4.3\n",
      "openai    : 1.30.5\n",
      "matplotlib: 3.8.0\n",
      "sqlite3   : 2.6.0\n",
      "numpy     : 1.26.4\n",
      "json      : 2.0.9\n",
      "tqdm      : 4.66.4\n",
      "pandas    : 2.2.3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49239d5f",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bdfb3c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('d4sci.mplstyle')\n",
    "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0724b260",
   "metadata": {},
   "source": [
    "# Text to Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "27a08608",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c0d870ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "        {\n",
    "            \"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",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "            \"content\": \"\"\"Create a recursive Python function to compute \n",
    "                Fibonacci numbers. Don't provide any explanation, just the code\"\"\"\n",
    "        },\n",
    "  ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "094a7141",
   "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": "02871015",
   "metadata": {},
   "source": [
    "Which produces the expected result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "99000cde",
   "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": "27648ba2",
   "metadata": {},
   "source": [
    "and works as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c7b82c05",
   "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": 18,
   "id": "4b27ba8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "610"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fibonacci(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b44d275",
   "metadata": {},
   "source": [
    "Let us define a utility function to make sequential queries easier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fe28c154",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(messages, prompt):\n",
    "    messages.append({\"role\":\"user\", \"content\":prompt})\n",
    "    \n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\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": "cb4fcd3f",
   "metadata": {},
   "source": [
    "# Adding comments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6d6c5693",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```python\n",
      "def fibonacci(n):\n",
      "    # Base case\n",
      "    if n <= 1:\n",
      "        return n\n",
      "    else:\n",
      "        return fibonacci(n-1) + fibonacci(n-2)\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"Can you add comments to this function?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6244a3dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To calculate Fibonacci numbers by breaking down the problem into smaller subproblems.\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"What is the purpose of recursion in this piece of code?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "42df0a45",
   "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='```python\\ndef fibonacci(n):\\n    # Base case\\n    if n <= 1:\\n        return n\\n    else:\\n        return fibonacci(n-1) + fibonacci(n-2)\\n```', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user',\n",
       "  'content': 'What is the purpose of recursion in this piece of code?'},\n",
       " ChatCompletionMessage(content='To calculate Fibonacci numbers by breaking down the problem into smaller subproblems.', role='assistant', function_call=None, tool_calls=None, refusal=None)]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd15571",
   "metadata": {},
   "source": [
    "# Explain Existing Code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fedd4f0",
   "metadata": {},
   "source": [
    "Let's use a relatively small python script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f9d1c9b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_text = \"\".join(open(\"data/EpiModel.py\").readlines())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2ad0e5e3",
   "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": 25,
   "id": "42d75591",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This code defines a class `EpiModel` for simulating and numerically integrating a simple epidemic model. It includes methods for adding interactions, spontaneous transitions, and vaccinations, as well as for plotting and calculating the basic reproduction number (R0). The main script at the end demonstrates how to use the class to simulate an SIR model and plot the results.\n",
      "CPU times: user 13 ms, sys: 3.27 ms, total: 16.2 ms\n",
      "Wall time: 1.04 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Please explain what this piece of code does: ```%s```\" % code_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "15a62b20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```python\n",
      "def __init__(self, compartments=None):\n",
      "    \"\"\"Initialize the EpiModel object.\n",
      "    \n",
      "    Args:\n",
      "        compartments (list): List of compartment names.\n",
      "    \"\"\"\n",
      "\n",
      "def add_interaction(self, source, target, agent, rate):\n",
      "    \"\"\"Add an interaction between compartments.\n",
      "    \n",
      "    Args:\n",
      "        source (str): Source compartment name.\n",
      "        target (str): Target compartment name.\n",
      "        agent (str): Agent involved in the interaction.\n",
      "        rate (float): Rate of the interaction.\n",
      "    \"\"\"\n",
      "\n",
      "def add_spontaneous(self, source, target, rate):\n",
      "    \"\"\"Add a spontaneous transition between compartments.\n",
      "    \n",
      "    Args:\n",
      "        source (str): Source compartment name.\n",
      "        target (str): Target compartment name.\n",
      "        rate (float): Rate of the transition.\n",
      "    \"\"\"\n",
      "\n",
      "def add_vaccination(self, source, target, rate, start):\n",
      "    \"\"\"Add a vaccination transition between compartments.\n",
      "    \n",
      "    Args:\n",
      "        source (str): Source compartment name.\n",
      "        target (str): Target compartment name.\n",
      "        rate (float): Rate of the vaccination.\n",
      "        start (int): Start time of the vaccination.\n",
      "    \"\"\"\n",
      "\n",
      "def _new_cases(self, population, time, pos):\n",
      "    \"\"\"Internal function used by integration routine.\n",
      "    \n",
      "    Args:\n",
      "        population (numpy.array): Array of population counts.\n",
      "        time (int): Current time step.\n",
      "        pos (dict): Mapping of compartment names to indices.\n",
      "    \"\"\"\n",
      "\n",
      "# Add docstrings for the remaining methods in a similar format\n",
      "```\n",
      "CPU times: user 11.8 ms, sys: 2.63 ms, total: 14.5 ms\n",
      "Wall time: 2.87 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": "5d35344d",
   "metadata": {},
   "source": [
    "# Interacting with a database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29d8fe16",
   "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": 28,
   "id": "fe6c1156",
   "metadata": {},
   "outputs": [],
   "source": [
    "con = sqlite3.connect(\"data/Northwind_small.sqlite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "057bf8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"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\", \"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": 30,
   "id": "d248d3a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT e.FirstName, e.LastName, et.TerritoryId\n",
      "FROM Employee e\n",
      "JOIN EmployeeTerritory et ON e.Id = et.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": 31,
   "id": "673bbcce",
   "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": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(query_sql, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "67029d0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT et.TerritoryId, COUNT(e.Id) AS NumEmployees\n",
      "FROM Employee e\n",
      "JOIN EmployeeTerritory et ON e.Id = et.EmployeeId\n",
      "GROUP BY et.TerritoryId;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"Compute how many employees work in each territory\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "aea608f8",
   "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>NumEmployees</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  NumEmployees\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": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "2d929f55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT COUNT(*)\n",
      "FROM Shipper;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"How many shippers do we work with?\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8f37ba48",
   "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": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d71ea68c",
   "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 e.FirstName, e.LastName, et.TerritoryId\\nFROM Employee e\\nJOIN EmployeeTerritory et ON e.Id = et.EmployeeId;', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user',\n",
       "  'content': 'Compute how many employees work in each territory'},\n",
       " ChatCompletionMessage(content='SELECT TerritoryId, COUNT(EmployeeId) AS NumEmployees\\nFROM EmployeeTerritory\\nGROUP BY TerritoryId;', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user',\n",
       "  'content': 'Compute how many employees work in each territory'},\n",
       " ChatCompletionMessage(content='SELECT et.TerritoryId, COUNT(e.Id) AS NumEmployees\\nFROM Employee e\\nJOIN EmployeeTerritory et ON e.Id = et.EmployeeId\\nGROUP BY et.TerritoryId;', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user', 'content': 'How many shippers do we work with?'},\n",
       " ChatCompletionMessage(content='SELECT COUNT(*)\\nFROM Shipper;', role='assistant', function_call=None, tool_calls=None, refusal=None)]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "51069cfb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "14ef32ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=messages,\n",
    "        temperature=0,\n",
    "        max_tokens=1024\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "cfe914e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT COUNT(Id)\\nFROM Shipper;'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b998b2",
   "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
}