5. PandasAI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "    <div style=\"width: 150px; float: left;\"> <img src=\"data/D4Sci_logo_ball.png\" alt=\"Data For Science, Inc\" align=\"left\" border=\"0\"> </div>\n",
    "    <div style=\"float: left; margin-left: 10px;\"> <h1>LLMs for Data Science</h1>\n",
    "<h1>PandasAI</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,
   "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 matplotlib.pyplot as plt \n",
    "import sqlite3\n",
    "\n",
    "import pandasai\n",
    "from pandasai import SmartDataframe, SmartDatalake, Agent\n",
    "from pandasai.llm import OpenAI\n",
    "from pandasai.connectors import SqliteConnector\n",
    "\n",
    "import watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by print out the versions of the libraries we're using for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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: 03802c3bf87993c3670c2fd8bf86e59d3d60bdfd\n",
      "\n",
      "numpy     : 1.26.4\n",
      "matplotlib: 3.8.0\n",
      "sqlite3   : 2.6.0\n",
      "pandas    : 2.2.3\n",
      "pandasai  : 2.2.14\n",
      "watermark : 2.4.3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.style.use('./d4sci.mplstyle')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SmartDataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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>country</th>\n",
       "      <th>sales</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>United States</td>\n",
       "      <td>5000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>United Kingdom</td>\n",
       "      <td>3200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>France</td>\n",
       "      <td>2900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Germany</td>\n",
       "      <td>4100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Italy</td>\n",
       "      <td>2300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Spain</td>\n",
       "      <td>2100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Canada</td>\n",
       "      <td>2500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Australia</td>\n",
       "      <td>2600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Japan</td>\n",
       "      <td>4500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>China</td>\n",
       "      <td>7000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          country  sales\n",
       "0   United States   5000\n",
       "1  United Kingdom   3200\n",
       "2          France   2900\n",
       "3         Germany   4100\n",
       "4           Italy   2300\n",
       "5           Spain   2100\n",
       "6          Canada   2500\n",
       "7       Australia   2600\n",
       "8           Japan   4500\n",
       "9           China   7000"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sales_by_country = pd.DataFrame({\n",
    "    \"country\": [\"United States\", \"United Kingdom\", \"France\", \"Germany\", \"Italy\", \"Spain\", \"Canada\", \"Australia\", \"Japan\", \"China\"],\n",
    "    \"sales\": [5000, 3200, 2900, 4100, 2300, 2100, 2500, 2600, 4500, 7000]\n",
    "})\n",
    "sales_by_country"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unless we specify otherwise, it will default to __BambooLLM__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The top 5 countries by sales are: China, United States, Japan, Germany, United Kingdom'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = SmartDataframe(sales_by_country)\n",
    "df.chat('Which are the top 5 countries by sales?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>country</th>\n",
       "      <th>sales</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>China</td>\n",
       "      <td>7000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>United States</td>\n",
       "      <td>5000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Japan</td>\n",
       "      <td>4500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Germany</td>\n",
       "      <td>4100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>United Kingdom</td>\n",
       "      <td>3200</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          country  sales\n",
       "9           China   7000\n",
       "0   United States   5000\n",
       "8           Japan   4500\n",
       "3         Germany   4100\n",
       "1  United Kingdom   3200"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sales_by_country.sort_values('sales', ascending=False).head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create an OpenAI instance (with our own api key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And provide it to the SmartDataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = SmartDataframe(sales_by_country, config={\"llm\": openai})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can interact with by asking specific questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The top 5 countries by sales are: China, United States, Japan, Germany, United Kingdom'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.chat('Which are the top 5 countries by sales?')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Questions can be almost arbitrarily complex, as long as the LLM is able to understand them :D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'type': 'dataframe', 'value':          country  sales\n",
      "0  United States   5000\n",
      "9          China   7000}\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>country</th>\n",
       "      <th>sales</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>United States</td>\n",
       "      <td>5000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>China</td>\n",
       "      <td>7000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         country  sales\n",
       "0  United States   5000\n",
       "9          China   7000"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.chat('which countries have 5000 sales or more?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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>country</th>\n",
       "      <th>sales</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>United States</td>\n",
       "      <td>5000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>China</td>\n",
       "      <td>7000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         country  sales\n",
       "0  United States   5000\n",
       "9          China   7000"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sales_by_country[sales_by_country.sales >= 5000]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SmartDatalake"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To interact with multiple dataframes we must use a SmartDatalake. We start by defining a couple of small data dictionaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "employees_data = {\n",
    "    'EmployeeID': [1, 2, 3, 4, 5],\n",
    "    'Name': ['John', 'Emma', 'Liam', 'Olivia', 'William'],\n",
    "    'Department': ['HR', 'Sales', 'IT', 'Marketing', 'Finance']\n",
    "}\n",
    "\n",
    "salaries_data = {\n",
    "    'EmployeeID': [1, 2, 3, 4, 5],\n",
    "    'Salary': [5000, 6000, 4500, 7000, 5500]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That we load into individual pandas Dataframe objects (these could also be SmartDataframes instead)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "employees_df = pd.DataFrame(employees_data)\n",
    "salaries_df = pd.DataFrame(salaries_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can now pass to our Smartdatalake constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "lake = SmartDatalake([employees_df, salaries_df], config={\"llm\": openai})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And interact with just as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'type': 'string', 'value': 'The highest paid employee is Olivia with a salary of 7000.'}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The highest paid employee is Olivia with a salary of 7000.'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lake.chat(\"Who gets paid the most?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please not that to answer this question, PandasAI had a to perform a join between the two data frames "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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>EmployeeID</th>\n",
       "      <th>Name</th>\n",
       "      <th>Department</th>\n",
       "      <th>Salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>Olivia</td>\n",
       "      <td>Marketing</td>\n",
       "      <td>7000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Emma</td>\n",
       "      <td>Sales</td>\n",
       "      <td>6000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>William</td>\n",
       "      <td>Finance</td>\n",
       "      <td>5500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>John</td>\n",
       "      <td>HR</td>\n",
       "      <td>5000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Liam</td>\n",
       "      <td>IT</td>\n",
       "      <td>4500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   EmployeeID     Name Department  Salary\n",
       "3           4   Olivia  Marketing    7000\n",
       "1           2     Emma      Sales    6000\n",
       "4           5  William    Finance    5500\n",
       "0           1     John         HR    5000\n",
       "2           3     Liam         IT    4500"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joint_df = employees_df.merge(salaries_df, on='EmployeeID')\n",
    "joint_df.sort_values('Salary', ascending=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then select the Name corresponding to the largest Salary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Olivia'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joint_df.iloc[joint_df['Salary'].idxmax()].Name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SqliteConnector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We instantiate a SqliteConnector with all the necessary information about our table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "employee_connector = SqliteConnector(config={\n",
    "    \"database\" : \"data/Northwind_small.sqlite\",\n",
    "    \"table\" : \"Employee\",\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can use as just any other data source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "employee = SmartDataframe(employee_connector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "employee.chat('How many records are there ?')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which we can easily confirm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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>9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   COUNT(*)\n",
       "0         9"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "con = sqlite3.connect(\"data/Northwind_small.sqlite\")\n",
    "pd.read_sql_query(\"SELECT COUNT(*) FROM employee;\", con)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now create a SmartDatalake with a few other tables. We have to create a connector for each of them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "territory_connector = SqliteConnector(config={\n",
    "    \"database\" : \"data/Northwind_small.sqlite\",\n",
    "    \"table\" : \"EmployeeTerritory\",\n",
    "})\n",
    "\n",
    "shipper_connector = SqliteConnector(config={\n",
    "    \"database\" : \"data/Northwind_small.sqlite\",\n",
    "    \"table\" : \"Shipper\",\n",
    "})\n",
    "\n",
    "detail_connector = SqliteConnector(config={\n",
    "    \"database\" : \"data/Northwind_small.sqlite\",\n",
    "    \"table\" : \"OrderDetail\",\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "northwind = Agent([\n",
    "    employee_connector, \n",
    "    territory_connector, \n",
    "    shipper_connector, \n",
    "    detail_connector\n",
    "], \n",
    "config={\"llm\": openai})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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>6897</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>1581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>1730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>1833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>2116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>2139</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>2184</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>2903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>7960</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>8837</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>3049</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>3801</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         6897\n",
       "1      Nancy    Davolio        19713\n",
       "2     Andrew     Fuller         1581\n",
       "3     Andrew     Fuller         1730\n",
       "4     Andrew     Fuller         1833\n",
       "5     Andrew     Fuller         2116\n",
       "6     Andrew     Fuller         2139\n",
       "7     Andrew     Fuller         2184\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         2903\n",
       "17    Steven   Buchanan         7960\n",
       "18    Steven   Buchanan         8837\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         3049\n",
       "43      Anne  Dodsworth         3801\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": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "northwind.chat(\"Generate a table with employee first name, last name and territory id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'type': 'dataframe', 'value':     TerritoryId  EmployeeCount\n",
      "0          6897              1\n",
      "1         98004              1\n",
      "2         98104              1\n",
      "3         60179              1\n",
      "4         60601              1\n",
      "5         80202              1\n",
      "6         80909              1\n",
      "7         90405              1\n",
      "8         94025              1\n",
      "9         94105              1\n",
      "10        95008              1\n",
      "11        95054              1\n",
      "12        95060              1\n",
      "13        19428              1\n",
      "14        44122              1\n",
      "15        45839              1\n",
      "16        53404              1\n",
      "17         3049              1\n",
      "18         3801              1\n",
      "19        48075              1\n",
      "20        48084              1\n",
      "21        48304              1\n",
      "22        55113              1\n",
      "23        98052              1\n",
      "24        85251              1\n",
      "25        19713              1\n",
      "26        85014              1\n",
      "27         1581              1\n",
      "28         1730              1\n",
      "29         1833              1\n",
      "30         2116              1\n",
      "31         2139              1\n",
      "32         2184              1\n",
      "33        40222              1\n",
      "34        30346              1\n",
      "35        31406              1\n",
      "36        32859              1\n",
      "37        33607              1\n",
      "38        20852              1\n",
      "39        27403              1\n",
      "40        27511              1\n",
      "41         2903              1\n",
      "42         7960              1\n",
      "43         8837              1\n",
      "44        10019              1\n",
      "45        10038              1\n",
      "46        11747              1\n",
      "47        14450              1\n",
      "48        55439              1}\n"
     ]
    }
   ],
   "source": [
    "northwind.chat(\"Compute how many employees work in each territory\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which again required the LLM to figure out the not so trivial JOIN to perform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql_query(\"\"\"\n",
    "SELECT \n",
    "    TerritoryID, \n",
    "    COUNT(DISTINCT EmployeeID) AS EmployeeCount\n",
    "FROM EmployeeTerritory \n",
    "JOIN Employee ON Employee.ID=EmployeeID \n",
    "GROUP BY 1\n",
    "\"\"\", con)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can confirm the way the agent is thinking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To create the code, I started by thinking about the information we needed: the names of employees and the territories they belong to. First, I combined two sets of information—one that had employee details and another that listed which territories those employees are associated with.\n",
      "\n",
      "Once I had this combined information, I focused on counting how many employees were in each territory. I organized the data so that it clearly showed each territory alongside the number of employees working there. Finally, I prepared the results in a way that makes it easy to understand and use.\n",
      "\n",
      "In summary, the process involved gathering the right information, combining it, counting the employees per territory, and then presenting the results clearly.\n"
     ]
    }
   ],
   "source": [
    "print(northwind.explain())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "code = northwind.generate_code(\"Compute how many employees work in each territory\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('employees_df = dfs[0]\\n'\n",
      " 'employee_territory_df = dfs[1]\\n'\n",
      " \"merged_df = pd.merge(employee_territory_df, employees_df[['Id', 'FirstName', \"\n",
      " \"'LastName']], left_on='EmployeeId', right_on='Id')\\n\"\n",
      " \"territory_counts = merged_df['TerritoryId'].value_counts().reset_index()\\n\"\n",
      " \"territory_counts.columns = ['TerritoryId', 'EmployeeCount']\\n\"\n",
      " \"result = {'type': 'dataframe', 'value': territory_counts}\\n\"\n",
      " 'print(result)')\n"
     ]
    }
   ],
   "source": [
    "pprint(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "territory = SmartDataframe(territory_connector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "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>EmployeeId</th>\n",
       "      <th>NumTerritories</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>6</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>7</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>8</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>9</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   EmployeeId  NumTerritories\n",
       "0           1               2\n",
       "1           2               7\n",
       "2           3               4\n",
       "3           4               3\n",
       "4           5               7\n",
       "5           6               5\n",
       "6           7              10\n",
       "7           8               4\n",
       "8           9               7"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "territory.chat('Compute the number of territories by employeeid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting Generated Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PandasAI allows us to see what code it used to produce the previous output by looking into the __last_code_executed__ field. See [here](https://github.com/Sinaptik-AI/pandas-ai/discussions/500) for details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "employee_df = dfs[0]\n",
      "employee_territory_df = dfs[1]\n",
      "merged_df = pd.merge(employee_df[['Id', 'FirstName', 'LastName']], employee_territory_df, left_on='Id', right_on='EmployeeId')\n",
      "territory_counts = merged_df['TerritoryId'].value_counts().reset_index()\n",
      "territory_counts.columns = ['TerritoryId', 'EmployeeCount']\n",
      "result = {'type': 'dataframe', 'value': territory_counts}\n",
      "print(result)\n"
     ]
    }
   ],
   "source": [
    "print(northwind.last_code_executed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "     <img src=\"data/D4Sci_logo_full.png\" alt=\"Data For Science, Inc\" align=\"center\" border=\"0\" width=300px> \n",
    "</div>"
   ]
  }
 ],
 "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"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "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": 4
}