LangChain with SQL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ed6aab1",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# SQL Chain example\n",
    "\n",
    "This example demonstrates the use of the `SQLDatabaseChain` for answering questions over a database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f66479",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Under the hood, LangChain uses SQLAlchemy to connect to SQL databases. The `SQLDatabaseChain` can therefore be used with any SQL dialect supported by SQLAlchemy, such as MS SQL, MySQL, MariaDB, PostgreSQL, Oracle SQL, and SQLite. Please refer to the SQLAlchemy documentation for more information about requirements for connecting to your database. For example, a connection to MySQL requires an appropriate connector such as PyMySQL. A URI for a MySQL connection might look like: `mysql+pymysql://user:pass@some_mysql_db_address/db_name`\n",
    "\n",
    "This demonstration uses SQLite and the example Chinook database.\n",
    "To set it up, follow the instructions on https://database.guide/2-sample-databases-sqlite/, placing the `.db` file in a notebooks folder at the root of this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting pymysql\n",
      "  Downloading PyMySQL-1.0.3-py3-none-any.whl (43 kB)\n",
      "\u001b[K     |████████████████████████████████| 43 kB 581 kB/s eta 0:00:01\n",
      "\u001b[?25hInstalling collected packages: pymysql\n",
      "Successfully installed pymysql-1.0.3\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install pymysql"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "231fa5c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting mysql-connector-python\n",
      "  Downloading mysql_connector_python-8.0.33-py2.py3-none-any.whl (390 kB)\n",
      "\u001b[K     |████████████████████████████████| 390 kB 3.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: protobuf<=3.20.3,>=3.11.0 in /Users/raydenx/opt/anaconda3/lib/python3.9/site-packages (from mysql-connector-python) (3.20.3)\n",
      "Installing collected packages: mysql-connector-python\n",
      "Successfully installed mysql-connector-python-8.0.33\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install mysql-connector-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d0e27d88",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from langchain import OpenAI, SQLDatabase, SQLDatabaseChain\n",
    "import pandas as pd\n",
    "import os \n",
    "import pymysql\n",
    "# from sqlalchemy import create_engine\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7bcf7828",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-04-26 22:34:26,789 INFO sqlalchemy.engine.Engine SHOW VARIABLES LIKE 'sql_mode'\n",
      "2023-04-26 22:34:26,790 INFO sqlalchemy.engine.Engine [raw sql] {}\n",
      "2023-04-26 22:34:26,807 INFO sqlalchemy.engine.Engine SHOW VARIABLES LIKE 'lower_case_table_names'\n",
      "2023-04-26 22:34:26,808 INFO sqlalchemy.engine.Engine [generated in 0.00078s] {}\n",
      "2023-04-26 22:34:26,812 INFO sqlalchemy.engine.Engine SELECT DATABASE()\n",
      "2023-04-26 22:34:26,813 INFO sqlalchemy.engine.Engine [raw sql] {}\n",
      "2023-04-26 22:34:26,819 INFO sqlalchemy.engine.Engine \n",
      "SELECT * FROM invoicing.clients\n",
      "\n",
      "2023-04-26 22:34:26,820 INFO sqlalchemy.engine.Engine [raw sql] {}\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>client_id</th>\n",
       "      <th>name</th>\n",
       "      <th>address</th>\n",
       "      <th>city</th>\n",
       "      <th>state</th>\n",
       "      <th>phone</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Vinte</td>\n",
       "      <td>3 Nevada Parkway</td>\n",
       "      <td>Syracuse</td>\n",
       "      <td>NY</td>\n",
       "      <td>315-252-7305</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Myworks</td>\n",
       "      <td>34267 Glendale Parkway</td>\n",
       "      <td>Huntington</td>\n",
       "      <td>WV</td>\n",
       "      <td>304-659-1170</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Yadel</td>\n",
       "      <td>096 Pawling Parkway</td>\n",
       "      <td>San Francisco</td>\n",
       "      <td>CA</td>\n",
       "      <td>415-144-6037</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>Kwideo</td>\n",
       "      <td>81674 Westerfield Circle</td>\n",
       "      <td>Waco</td>\n",
       "      <td>TX</td>\n",
       "      <td>254-750-0784</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>Topiclounge</td>\n",
       "      <td>0863 Farmco Road</td>\n",
       "      <td>Portland</td>\n",
       "      <td>OR</td>\n",
       "      <td>971-888-9129</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   client_id         name                   address           city state  \\\n",
       "0          1        Vinte          3 Nevada Parkway       Syracuse    NY   \n",
       "1          2      Myworks    34267 Glendale Parkway     Huntington    WV   \n",
       "2          3        Yadel       096 Pawling Parkway  San Francisco    CA   \n",
       "3          4       Kwideo  81674 Westerfield Circle           Waco    TX   \n",
       "4          5  Topiclounge          0863 Farmco Road       Portland    OR   \n",
       "\n",
       "          phone  \n",
       "0  315-252-7305  \n",
       "1  304-659-1170  \n",
       "2  415-144-6037  \n",
       "3  254-750-0784  \n",
       "4  971-888-9129  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "# 创建连接引擎\n",
    "engine = create_engine(\"mysql+pymysql://root:xrd93123875@127.0.0.1:3306\",echo=True)\n",
    "# 查询表数据\n",
    "query = '''\n",
    "SELECT * FROM invoicing.clients\n",
    "'''\n",
    "df = pd.read_sql_query(query, engine)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "307debbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy\n",
    "# sqlalchemy.orm.configure_mappers()\n",
    "db = SQLDatabase.from_uri(\"mysql+pymysql://address/classicmodels\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9404b076",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "872c882b",
   "metadata": {},
   "source": [
    "### py文件导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d42fb478",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "\n",
      "             write an SQL query to find how many customers live in USA.\n",
      "             \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m \n",
      "SELECT COUNT(*) AS 'Number of Customers in USA' \n",
      "FROM customers \n",
      "WHERE country = 'USA';\n",
      "\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(36,)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m 36 customers live in USA.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' 36 customers live in USA.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "OPENAI_API_KEY = 'open-ai-key'\n",
    "# db = SQLDatabase.from_uri(\"mysql+pymysql://root:xrd93123875@localhost:3306/store\")\n",
    "db = SQLDatabase.from_uri(\"mysql+pymysql://address/classicmodels\")\n",
    "llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)\n",
    "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True) \n",
    "db_chain.run('''\n",
    "             write an SQL query to find how many customers live in USA.\n",
    "             ''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1e692e",
   "metadata": {},
   "source": [
    "**NOTE:** For data-sensitive projects, you can specify `return_direct=True` in the `SQLDatabaseChain` initialization to directly return the output of the SQL query without any additional formatting. This prevents the LLM from seeing any contents within the database. Note, however, the LLM still has access to the database scheme (i.e. dialect, table and key names) by default."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "51392c04",
   "metadata": {},
   "source": [
    "# English Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dfdc90ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "\n",
      "write SQL to  Query the hierarchical relationship between employees and their subordinates in the classicmodels.employees table,sort by hierarchical, please use recursive CTE, and return the column name with data in the output, use table format.\n",
      "\n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m \n",
      "WITH RECURSIVE cte (employeeNumber, lastName, firstName, reportsTo, jobTitle, level) AS (\n",
      "    SELECT employeeNumber, lastName, firstName, reportsTo, jobTitle, 0 AS level\n",
      "    FROM employees\n",
      "    WHERE reportsTo IS NULL\n",
      "    UNION ALL\n",
      "    SELECT e.employeeNumber, e.lastName, e.firstName, e.reportsTo, e.jobTitle, cte.level + 1\n",
      "    FROM employees e\n",
      "    INNER JOIN cte\n",
      "    ON e.reportsTo = cte.employeeNumber\n",
      ")\n",
      "SELECT employeeNumber, lastName, firstName, reportsTo, jobTitle, level\n",
      "FROM cte\n",
      "ORDER BY level;\n",
      "\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(1002, 'Murphy', 'Diane', None, 'President', 0), (1056, 'Patterson', 'Mary', 1002, 'VP Sales', 1), (1076, 'Firrelli', 'Jeff', 1002, 'VP Marketing', 1), (1102, 'Bondur', 'Gerard', 1056, 'Sale Manager (EMEA)', 2), (1143, 'Bow', 'Anthony', 1056, 'Sales Manager (NA)', 2), (1621, 'Nishi', 'Mami', 1056, 'Sales Rep', 2), (1088, 'Patterson', 'William', 1056, 'Sales Manager (APAC)', 2), (1401, 'Castillo', 'Pamela', 1102, 'Sales Rep', 3), (1286, 'Tseng', 'Foon Yue', 1143, 'Sales Rep', 3), (1501, 'Bott', 'Larry', 1102, 'Sales Rep', 3), (1323, 'Vanauf', 'George', 1143, 'Sales Rep', 3), (1504, 'Jones', 'Barry', 1102, 'Sales Rep', 3), (1625, 'Kato', 'Yoshimi', 1621, 'Sales Rep', 3), (1611, 'Fixter', 'Andy', 1088, 'Sales Rep', 3), (1702, 'Gerard', 'Martin', 1102, 'Sales Rep', 3), (1612, 'Marsh', 'Peter', 1088, 'Sales Rep', 3), (1165, 'Jennings', 'Leslie', 1143, 'Sales Rep', 3), (1619, 'King', 'Tom', 1088, 'Sales Rep', 3), (1166, 'Thompson', 'Leslie', 1143, 'Sales Rep', 3), (1337, 'Bondur', 'Loui', 1102, 'Sales Rep', 3), (1188, 'Firrelli', 'Julie', 1143, 'Sales Rep', 3), (1370, 'Hernandez', 'Gerard', 1102, 'Sales Rep', 3), (1216, 'Patterson', 'Steve', 1143, 'Sales Rep', 3)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m The hierarchical relationship between employees and their subordinates in the classicmodels.employees table is as follows: \n",
      "\n",
      "employeeNumber\tlastName\tfirstName\treportsTo\tjobTitle\tlevel\n",
      "1002\tMurphy\tDiane\tNone\tPresident\t0\n",
      "1056\tPatterson\tMary\t1002\tVP Sales\t1\n",
      "1076\tFirrelli\tJeff\t1002\tVP Marketing\t1\n",
      "1102\tBondur\tGerard\t1056\tSale Manager (EMEA)\t2\n",
      "1143\tBow\tAnthony\t1056\tSales Manager (NA)\t2\n",
      "1621\tNishi\tMami\t1056\tSales Rep\t2\n",
      "1088\tPatterson\tWilliam\t1056\tSales Manager (APAC)\t2\n",
      "1401\tCastillo\tPamela\t1102\tSales Rep\t3\n",
      "1286\tTseng\tFoon Yue\t1143\tSales Rep\t3\n",
      "1501\tBott\tLarry\t1102\tSales Rep\t3\n",
      "1323\tVanauf\tGeorge\t1143\tSales Rep\t3\n",
      "1504\tJones\tBarry\t1102\tSales Rep\t3\n",
      "1625\tK\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' The hierarchical relationship between employees and their subordinates in the classicmodels.employees table is as follows: \\n\\nemployeeNumber\\tlastName\\tfirstName\\treportsTo\\tjobTitle\\tlevel\\n1002\\tMurphy\\tDiane\\tNone\\tPresident\\t0\\n1056\\tPatterson\\tMary\\t1002\\tVP Sales\\t1\\n1076\\tFirrelli\\tJeff\\t1002\\tVP Marketing\\t1\\n1102\\tBondur\\tGerard\\t1056\\tSale Manager (EMEA)\\t2\\n1143\\tBow\\tAnthony\\t1056\\tSales Manager (NA)\\t2\\n1621\\tNishi\\tMami\\t1056\\tSales Rep\\t2\\n1088\\tPatterson\\tWilliam\\t1056\\tSales Manager (APAC)\\t2\\n1401\\tCastillo\\tPamela\\t1102\\tSales Rep\\t3\\n1286\\tTseng\\tFoon Yue\\t1143\\tSales Rep\\t3\\n1501\\tBott\\tLarry\\t1102\\tSales Rep\\t3\\n1323\\tVanauf\\tGeorge\\t1143\\tSales Rep\\t3\\n1504\\tJones\\tBarry\\t1102\\tSales Rep\\t3\\n1625\\tK'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run('''\n",
    "write SQL to  Query the hierarchical relationship between employees and their subordinates in the classicmodels.employees table,sort by hierarchical, please use recursive CTE, and return the column name with data in the output, use table format.\n",
    "''')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f8001541",
   "metadata": {},
   "source": [
    "MultiLanguag Prompt Example: Chinese(Manderin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a3f90ee3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "\n",
      "写一个SQL查询，查找classicmodels.employees表中员工及其下属的层级关系，并按层级排序，请使用递归CTE，输出使用英文，并且识别列名放在输出里。\n",
      "\n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m \n",
      "WITH RECURSIVE employee_hierarchy (employeeNumber, lastName, firstName, reportsTo, level) AS (\n",
      "    SELECT employeeNumber, lastName, firstName, reportsTo, 0 AS level\n",
      "    FROM employees\n",
      "    WHERE reportsTo IS NULL\n",
      "    UNION ALL\n",
      "    SELECT e.employeeNumber, e.lastName, e.firstName, e.reportsTo, eh.level + 1\n",
      "    FROM employees e\n",
      "    INNER JOIN employee_hierarchy eh\n",
      "    ON e.reportsTo = eh.employeeNumber\n",
      ")\n",
      "SELECT employeeNumber, lastName, firstName, reportsTo, level\n",
      "FROM employee_hierarchy\n",
      "ORDER BY level;\n",
      "\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(1002, 'Murphy', 'Diane', None, 0), (1056, 'Patterson', 'Mary', 1002, 1), (1076, 'Firrelli', 'Jeff', 1002, 1), (1102, 'Bondur', 'Gerard', 1056, 2), (1143, 'Bow', 'Anthony', 1056, 2), (1621, 'Nishi', 'Mami', 1056, 2), (1088, 'Patterson', 'William', 1056, 2), (1401, 'Castillo', 'Pamela', 1102, 3), (1286, 'Tseng', 'Foon Yue', 1143, 3), (1501, 'Bott', 'Larry', 1102, 3), (1323, 'Vanauf', 'George', 1143, 3), (1504, 'Jones', 'Barry', 1102, 3), (1625, 'Kato', 'Yoshimi', 1621, 3), (1611, 'Fixter', 'Andy', 1088, 3), (1702, 'Gerard', 'Martin', 1102, 3), (1612, 'Marsh', 'Peter', 1088, 3), (1165, 'Jennings', 'Leslie', 1143, 3), (1619, 'King', 'Tom', 1088, 3), (1166, 'Thompson', 'Leslie', 1143, 3), (1337, 'Bondur', 'Loui', 1102, 3), (1188, 'Firrelli', 'Julie', 1143, 3), (1370, 'Hernandez', 'Gerard', 1102, 3), (1216, 'Patterson', 'Steve', 1143, 3)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m 员工及其下属的层级关系按层级排序为：1002 (Murphy, Diane), 1056 (Patterson, Mary), 1076 (Firrelli, Jeff), 1102 (Bondur, Gerard), 1143 (Bow, Anthony), 1621 (Nishi, Mami), 1088 (Patterson, William), 1401 (Castillo, Pamela), 1286 (Tseng, Foon Yue), 1501 (Bott, Larry), 1323 (Vanauf, George), 1504 (Jones, Barry), 1625 (Kato, Yoshimi), 1611 (Fixter, Andy), 1702 (Gerard, Martin), 1612 (Marsh, Peter), 1165 (Jennings, Leslie), 1619 (King, Tom), 1166 (Thompson, Leslie), 1337 (Bondur, Loui), 1188 (Firrelli, Julie), 1370 (Hernandez, Gerard), 1216 (Patterson, Steve).\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' 员工及其下属的层级关系按层级排序为：1002 (Murphy, Diane), 1056 (Patterson, Mary), 1076 (Firrelli, Jeff), 1102 (Bondur, Gerard), 1143 (Bow, Anthony), 1621 (Nishi, Mami), 1088 (Patterson, William), 1401 (Castillo, Pamela), 1286 (Tseng, Foon Yue), 1501 (Bott, Larry), 1323 (Vanauf, George), 1504 (Jones, Barry), 1625 (Kato, Yoshimi), 1611 (Fixter, Andy), 1702 (Gerard, Martin), 1612 (Marsh, Peter), 1165 (Jennings, Leslie), 1619 (King, Tom), 1166 (Thompson, Leslie), 1337 (Bondur, Loui), 1188 (Firrelli, Julie), 1370 (Hernandez, Gerard), 1216 (Patterson, Steve).'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run('''\n",
    "写一个SQL查询，查找classicmodels.employees表中员工及其下属的层级关系，并按层级排序，请使用递归CTE，输出使用英文，并且识别列名放在输出里。\n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd1fe8ea",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "31e799ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = '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'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "49eeeefc",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a8fc8f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "15ff81df",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "\n",
      "            write an SQL query to find the most frequently purchased product name for each customer. \n",
      "             If there is a tie, return the product name with the smaller lexicographic order.\n",
      "             \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m \n",
      "            SELECT c.name, p.name \n",
      "            FROM clients c \n",
      "            JOIN payments pa ON c.client_id = pa.client_id \n",
      "            JOIN payment_methods p ON pa.payment_method = p.payment_method_id \n",
      "            GROUP BY c.name, p.name \n",
      "            ORDER BY COUNT(*) DESC, p.name ASC \n",
      "            LIMIT 3;\n",
      "            \u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[('Topiclounge', 'Credit Card'), ('Yadel', 'Credit Card'), ('Topiclounge', 'Cash')]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m The most frequently purchased product name for each customer is Credit Card, followed by Cash.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' The most frequently purchased product name for each customer is Credit Card, followed by Cash.'"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run('''\n",
    "            write an SQL query to find the most frequently purchased product name for each customer. \n",
    "             If there is a tie, return the product name with the smaller lexicographic order.\n",
    "             ''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad2cba6",
   "metadata": {},
   "source": [
    "## Customize Prompt\n",
    "You can also customize the prompt that is used. Here is an example prompting it to understand that foobar is the same as the Employee table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8ca7bafb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts.prompt import PromptTemplate\n",
    "\n",
    "_DEFAULT_TEMPLATE = \"\"\"Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.\n",
    "Use the following format:\n",
    "\n",
    "Question: \"Question here\"\n",
    "SQLQuery: \"SQL Query to run\"\n",
    "SQLResult: \"Result of the SQLQuery\"\n",
    "Answer: \"Final answer here\"\n",
    "\n",
    "Only use the following tables:\n",
    "\n",
    "{table_info}\n",
    "\n",
    "If someone asks for the table foobar, they really mean the employee table.\n",
    "\n",
    "Question: {input}\"\"\"\n",
    "PROMPT = PromptTemplate(\n",
    "    input_variables=[\"input\", \"table_info\", \"dialect\"], template=_DEFAULT_TEMPLATE\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ec47a2bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, prompt=PROMPT, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ebb0674e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "How many client locate in NY?\n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM clients WHERE state = 'NY';\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(1,)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m There is 1 client located in NY.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' There is 1 client located in NY.'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run(\"How many client locate in NY?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88d8b969",
   "metadata": {},
   "source": [
    "## Return Intermediate Steps\n",
    "\n",
    "You can also return the intermediate steps of the SQLDatabaseChain. This allows you to access the SQL statement that was generated, as well as the result of running that against the SQL Database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "38559487",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, prompt=PROMPT, verbose=True, return_intermediate_steps=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "78b6af4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "How many client locate in NY?\n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM clients WHERE state = 'NY';\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(1,)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m There is 1 client located in NY.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[\" SELECT COUNT(*) FROM clients WHERE state = 'NY';\", '[(1,)]']"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = db_chain(\"How many client locate in NY?\")\n",
    "result[\"intermediate_steps\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b408f800",
   "metadata": {},
   "source": [
    "## Choosing how to limit the number of rows returned\n",
    "If you are querying for several rows of a table you can select the maximum number of results you want to get by using the 'top_k' parameter (default is 10). This is useful for avoiding query results that exceed the prompt max length or consume tokens unnecessarily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6adaa799",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True, top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "edfc8a8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "What are some example tracks by composer Johann Sebastian Bach? \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT Name, Composer FROM Track WHERE Composer LIKE '%Johann Sebastian Bach%' LIMIT 3;\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[('Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Johann Sebastian Bach'), ('Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria', 'Johann Sebastian Bach'), ('Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude', 'Johann Sebastian Bach')]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m Some example tracks by composer Johann Sebastian Bach are 'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria', and 'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude'.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Some example tracks by composer Johann Sebastian Bach are \\'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace\\', \\'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria\\', and \\'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude\\'.'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run(\"What are some example tracks by composer Johann Sebastian Bach?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcc5e936",
   "metadata": {},
   "source": [
    "## Adding example rows from each table\n",
    "Sometimes, the format of the data is not obvious and it is optimal to include a sample of rows from the tables in the prompt to allow the LLM to understand the data before providing a final query. Here we will use this feature to let the LLM know that artists are saved with their full names by providing two rows from the `Track` table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9a22ee47",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\n",
    "    \"sqlite:///../../../../notebooks/Chinook.db\",\n",
    "    include_tables=['Track'], # we include only one table to save tokens in the prompt :)\n",
    "    sample_rows_in_table_info=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952c0b4d",
   "metadata": {},
   "source": [
    "The sample rows are added to the prompt after each corresponding table's column information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9de86267",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CREATE TABLE \"Track\" (\n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(200) NOT NULL, \n",
      "\t\"AlbumId\" INTEGER, \n",
      "\t\"MediaTypeId\" INTEGER NOT NULL, \n",
      "\t\"GenreId\" INTEGER, \n",
      "\t\"Composer\" NVARCHAR(220), \n",
      "\t\"Milliseconds\" INTEGER NOT NULL, \n",
      "\t\"Bytes\" INTEGER, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"MediaTypeId\") REFERENCES \"MediaType\" (\"MediaTypeId\"), \n",
      "\tFOREIGN KEY(\"GenreId\") REFERENCES \"Genre\" (\"GenreId\"), \n",
      "\tFOREIGN KEY(\"AlbumId\") REFERENCES \"Album\" (\"AlbumId\")\n",
      ")\n",
      "/*\n",
      "2 rows from Track table:\n",
      "TrackId\tName\tAlbumId\tMediaTypeId\tGenreId\tComposer\tMilliseconds\tBytes\tUnitPrice\n",
      "1\tFor Those About To Rock (We Salute You)\t1\t1\t1\tAngus Young, Malcolm Young, Brian Johnson\t343719\t11170334\t0.99\n",
      "2\tBalls to the Wall\t2\t2\t1\tNone\t342562\t5510424\t0.99\n",
      "*/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jon/projects/langchain/langchain/sql_database.py:135: SAWarning: Dialect sqlite+pysqlite does *not* support Decimal objects natively, and SQLAlchemy must convert from floating point - rounding errors and other issues may occur. Please consider storing Decimal numbers as strings or integers on this platform for lossless storage.\n",
      "  sample_rows = connection.execute(command)\n"
     ]
    }
   ],
   "source": [
    "print(db.table_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "bcb7a489",
   "metadata": {},
   "outputs": [],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "81e05d82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "What are some example tracks by Bach? \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT Name FROM Track WHERE Composer LIKE '%Bach%' LIMIT 5;\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[('American Woman',), ('Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace',), ('Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria',), ('Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude',), ('Toccata and Fugue in D Minor, BWV 565: I. Toccata',)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m Some example tracks by Bach are 'American Woman', 'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria', 'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude', and 'Toccata and Fugue in D Minor, BWV 565: I. Toccata'.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Some example tracks by Bach are \\'American Woman\\', \\'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace\\', \\'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria\\', \\'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude\\', and \\'Toccata and Fugue in D Minor, BWV 565: I. Toccata\\'.'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain.run(\"What are some example tracks by Bach?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef94e948",
   "metadata": {},
   "source": [
    "### Custom Table Info\n",
    "In some cases, it can be useful to provide custom table information instead of using the automatically generated table definitions and the first `sample_rows_in_table_info` sample rows. For example, if you know that the first few rows of a table are uninformative, it could help to manually provide example rows that are more diverse or provide more information to the model. It is also possible to limit the columns that will be visible to the model if there are unnecessary columns. \n",
    "\n",
    "This information can be provided as a dictionary with table names as the keys and table information as the values. For example, let's provide a custom definition and sample rows for the Track table with only a few columns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2ad33ab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_table_info = {\n",
    "    \"Track\": \"\"\"CREATE TABLE Track (\n",
    "\t\"TrackId\" INTEGER NOT NULL, \n",
    "\t\"Name\" NVARCHAR(200) NOT NULL,\n",
    "\t\"Composer\" NVARCHAR(220),\n",
    "\tPRIMARY KEY (\"TrackId\")\n",
    ")\n",
    "/*\n",
    "3 rows from Track table:\n",
    "TrackId\tName\tComposer\n",
    "1\tFor Those About To Rock (We Salute You)\tAngus Young, Malcolm Young, Brian Johnson\n",
    "2\tBalls to the Wall\tNone\n",
    "3\tMy favorite song ever\tThe coolest composer of all time\n",
    "*/\"\"\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "db144352",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CREATE TABLE \"Playlist\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"PlaylistId\")\n",
      ")\n",
      "/*\n",
      "2 rows from Playlist table:\n",
      "PlaylistId\tName\n",
      "1\tMusic\n",
      "2\tMovies\n",
      "*/\n",
      "\n",
      "CREATE TABLE Track (\n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(200) NOT NULL,\n",
      "\t\"Composer\" NVARCHAR(220),\n",
      "\tPRIMARY KEY (\"TrackId\")\n",
      ")\n",
      "/*\n",
      "3 rows from Track table:\n",
      "TrackId\tName\tComposer\n",
      "1\tFor Those About To Rock (We Salute You)\tAngus Young, Malcolm Young, Brian Johnson\n",
      "2\tBalls to the Wall\tNone\n",
      "3\tMy favorite song ever\tThe coolest composer of all time\n",
      "*/\n"
     ]
    }
   ],
   "source": [
    "db = SQLDatabase.from_uri(\n",
    "    \"sqlite:///../../../../notebooks/Chinook.db\",\n",
    "    include_tables=['Track', 'Playlist'],\n",
    "    sample_rows_in_table_info=2,\n",
    "    custom_table_info=custom_table_info)\n",
    "\n",
    "print(db.table_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc6f507",
   "metadata": {},
   "source": [
    "Note how our custom table definition and sample rows for `Track` overrides the `sample_rows_in_table_info` parameter. Tables that are not overridden by `custom_table_info`, in this example `Playlist`, will have their table info gathered automatically as usual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "dfbda4e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "What are some example tracks by Bach? \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT Name, Composer FROM Track WHERE Composer LIKE '%Bach%' LIMIT 5;\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[('American Woman', 'B. Cummings/G. Peterson/M.J. Kale/R. Bachman'), ('Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Johann Sebastian Bach'), ('Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria', 'Johann Sebastian Bach'), ('Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude', 'Johann Sebastian Bach'), ('Toccata and Fugue in D Minor, BWV 565: I. Toccata', 'Johann Sebastian Bach')]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m Some example tracks by Bach are 'American Woman', 'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria', 'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude', and 'Toccata and Fugue in D Minor, BWV 565: I. Toccata'.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Some example tracks by Bach are \\'American Woman\\', \\'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace\\', \\'Aria Mit 30 Veränderungen, BWV 988 \"Goldberg Variations\": Aria\\', \\'Suite for Solo Cello No. 1 in G Major, BWV 1007: I. Prélude\\', and \\'Toccata and Fugue in D Minor, BWV 565: I. Toccata\\'.'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)\n",
    "db_chain.run(\"What are some example tracks by Bach?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c12ae15a",
   "metadata": {},
   "source": [
    "## SQLDatabaseSequentialChain\n",
    "\n",
    "Chain for querying SQL database that is a sequential chain.\n",
    "\n",
    "The chain is as follows:\n",
    "\n",
    "    1. Based on the query, determine which tables to use.\n",
    "    2. Based on those tables, call the normal SQL database chain.\n",
    "\n",
    "This is useful in cases where the number of tables in the database is large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e59a4740",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import SQLDatabaseSequentialChain\n",
    "db = SQLDatabase.from_uri(\"sqlite:///../../../../notebooks/Chinook.db\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "58bb49b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = SQLDatabaseSequentialChain.from_llm(llm, db, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "95017b1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseSequentialChain chain...\u001b[0m\n",
      "Table names to use:\n",
      "\u001b[33;1m\u001b[1;3m['Customer', 'Employee']\u001b[0m\n",
      "\n",
      "\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
      "How many employees are also customers? \n",
      "SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM Employee INNER JOIN Customer ON Employee.EmployeeId = Customer.SupportRepId;\u001b[0m\n",
      "SQLResult: \u001b[33;1m\u001b[1;3m[(59,)]\u001b[0m\n",
      "Answer:\u001b[32;1m\u001b[1;3m 59 employees are also customers.\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' 59 employees are also customers.'"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"How many employees are also customers?\")"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}