{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0e499e90-7a6d-4fab-8aab-31a4df417601",
   "metadata": {},
   "source": [
    "# SQL Database Agent\n",
    "\n",
    "This notebook showcases an agent designed to interact with a sql databases. The agent builds off of [SQLDatabaseChain](https://python.langchain.com/docs/use_cases/tabular/sqlite) and is designed to answer more general questions about a database, as well as recover from errors.\n",
    "\n",
    "Note that, as this agent is in active development, all answers might not be correct. Additionally, it is not guaranteed that the agent won't perform DML statements on your database given certain questions. Be careful running it on sensitive data!\n",
    "\n",
    "This uses the example Chinook database. 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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ec927ac6-9b2a-4e8a-9a6e-3e429191875c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "53422913-967b-4f2a-8022-00269c1be1b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.agents import create_sql_agent\n",
    "from langchain.agents.agent_toolkits import SQLDatabaseToolkit\n",
    "from langchain.sql_database import SQLDatabase\n",
    "from langchain.llms.openai import OpenAI\n",
    "from langchain.agents import AgentExecutor\n",
    "from langchain.agents.agent_types import AgentType\n",
    "from langchain.chat_models import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "65ec5bb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"sqlite:///../../../../../notebooks/Chinook.db\")\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=OpenAI(temperature=0))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f74d1792",
   "metadata": {},
   "source": [
    "## Using ZERO_SHOT_REACT_DESCRIPTION\n",
    "\n",
    "This shows how to initialize the agent using the ZERO_SHOT_REACT_DESCRIPTION agent type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "090f3699-79c6-4ce1-ab96-a94f0121fd64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent_executor = create_sql_agent(\n",
    "    llm=OpenAI(temperature=0),\n",
    "    toolkit=toolkit,\n",
    "    verbose=True,\n",
    "    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "971cc455",
   "metadata": {},
   "source": [
    "## Using OpenAI Functions\n",
    "\n",
    "This shows how to initialize the agent using the OPENAI_FUNCTIONS agent type. Note that this is an alternative to the above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6426a27d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# agent_executor = create_sql_agent(\n",
    "#     llm=ChatOpenAI(temperature=0, model=\"gpt-3.5-turbo-0613\"),\n",
    "#     toolkit=toolkit,\n",
    "#     verbose=True,\n",
    "#     agent_type=AgentType.OPENAI_FUNCTIONS\n",
    "# )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "54c01168",
   "metadata": {},
   "source": [
    "## Disclamer ⚠️\n",
    "\n",
    "The query chain may generate insert/update/delete queries. When this is not expected, use a custom prompt or create a SQL users without write permissions.\n",
    "\n",
    "The final user might overload your SQL database by asking a simple question such as \"run the biggest query possible\". The generated query might look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "949772b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "SELECT * FROM \"public\".\"users\"\n",
    "    JOIN \"public\".\"user_permissions\" ON \"public\".\"users\".id = \"public\".\"user_permissions\".user_id\n",
    "    JOIN \"public\".\"projects\" ON \"public\".\"users\".id = \"public\".\"projects\".user_id\n",
    "    JOIN \"public\".\"events\" ON \"public\".\"projects\".id = \"public\".\"events\".project_id;"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5a4a9455",
   "metadata": {},
   "source": [
    "For a transactional SQL database, if one of the table above contains millions of rows, the query might cause trouble to other applications using the same database.\n",
    "\n",
    "Most datawarehouse oriented databases support user-level quota, for limiting resource usage."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "36ae48c7-cb08-4fef-977e-c7d4b96a464b",
   "metadata": {},
   "source": [
    "## Example: describing a table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ff70e83d-5ad0-4fc7-bb96-27d82ac166d7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `list_tables_sql_db` with `{}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mAlbum, Artist, Track, PlaylistTrack, InvoiceLine, sales_table, Playlist, Genre, Employee, Customer, Invoice, MediaType\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `schema_sql_db` with `PlaylistTrack`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from PlaylistTrack table:\n",
      "PlaylistId\tTrackId\n",
      "1\t3402\n",
      "1\t3389\n",
      "1\t3390\n",
      "*/\u001b[0m\u001b[32;1m\u001b[1;3mThe `PlaylistTrack` table has two columns: `PlaylistId` and `TrackId`. It is a junction table that represents the relationship between playlists and tracks. \n",
      "\n",
      "Here is the schema of the `PlaylistTrack` table:\n",
      "\n",
      "```\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "```\n",
      "\n",
      "Here are three sample rows from the `PlaylistTrack` table:\n",
      "\n",
      "```\n",
      "PlaylistId   TrackId\n",
      "1            3402\n",
      "1            3389\n",
      "1            3390\n",
      "```\n",
      "\n",
      "Please let me know if there is anything else I can help you with.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The `PlaylistTrack` table has two columns: `PlaylistId` and `TrackId`. It is a junction table that represents the relationship between playlists and tracks. \\n\\nHere is the schema of the `PlaylistTrack` table:\\n\\n```\\nCREATE TABLE \"PlaylistTrack\" (\\n\\t\"PlaylistId\" INTEGER NOT NULL, \\n\\t\"TrackId\" INTEGER NOT NULL, \\n\\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \\n\\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \\n\\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\\n)\\n```\\n\\nHere are three sample rows from the `PlaylistTrack` table:\\n\\n```\\nPlaylistId   TrackId\\n1            3402\\n1            3389\\n1            3390\\n```\\n\\nPlease let me know if there is anything else I can help you with.'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\"Describe the playlisttrack table\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9abcfe8e-1868-42a4-8345-ad2d9b44c681",
   "metadata": {},
   "source": [
    "## Example: describing a table, recovering from an error\n",
    "\n",
    "In this example, the agent tries to search for a table that doesn't exist, but finds the next best result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bea76658-a65b-47e2-b294-6d52c5556246",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mAction: list_tables_sql_db\n",
      "Action Input: \"\"\u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mGenre, PlaylistTrack, MediaType, Invoice, InvoiceLine, Track, Playlist, Customer, Album, Employee, Artist\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should look at the schema of the PlaylistSong table\n",
      "Action: schema_sql_db\n",
      "Action Input: \"PlaylistSong\"\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3mError: table_names {'PlaylistSong'} not found in database\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should check the spelling of the table\n",
      "Action: list_tables_sql_db\n",
      "Action Input: \"\"\u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mGenre, PlaylistTrack, MediaType, Invoice, InvoiceLine, Track, Playlist, Customer, Album, Employee, Artist\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m The table is called PlaylistTrack\n",
      "Action: schema_sql_db\n",
      "Action Input: \"PlaylistTrack\"\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3m\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'PlaylistTrack' LIMIT 3;\n",
      "PlaylistId TrackId\n",
      "1 3402\n",
      "1 3389\n",
      "1 3390\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer\n",
      "Final Answer: The PlaylistTrack table contains two columns, PlaylistId and TrackId, which are both integers and are used to link Playlist and Track tables.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The PlaylistTrack table contains two columns, PlaylistId and TrackId, which are both integers and are used to link Playlist and Track tables.'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\"Describe the playlistsong table\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fbc26af-97e4-4a21-82aa-48bdc992da26",
   "metadata": {},
   "source": [
    "## Example: running queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "17bea710-4a23-4de0-b48e-21d57be48293",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mAction: list_tables_sql_db\n",
      "Action Input: \"\"\u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mInvoice, MediaType, Artist, InvoiceLine, Genre, Playlist, Employee, Album, PlaylistTrack, Track, Customer\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should look at the schema of the relevant tables to see what columns I can use.\n",
      "Action: schema_sql_db\n",
      "Action Input: \"Invoice, Customer\"\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3m\n",
      "CREATE TABLE \"Customer\" (\n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"FirstName\" NVARCHAR(40) NOT NULL, \n",
      "\t\"LastName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"Company\" NVARCHAR(80), \n",
      "\t\"Address\" NVARCHAR(70), \n",
      "\t\"City\" NVARCHAR(40), \n",
      "\t\"State\" NVARCHAR(40), \n",
      "\t\"Country\" NVARCHAR(40), \n",
      "\t\"PostalCode\" NVARCHAR(10), \n",
      "\t\"Phone\" NVARCHAR(24), \n",
      "\t\"Fax\" NVARCHAR(24), \n",
      "\t\"Email\" NVARCHAR(60) NOT NULL, \n",
      "\t\"SupportRepId\" INTEGER, \n",
      "\tPRIMARY KEY (\"CustomerId\"), \n",
      "\tFOREIGN KEY(\"SupportRepId\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'Customer' LIMIT 3;\n",
      "CustomerId FirstName LastName Company Address City State Country PostalCode Phone Fax Email SupportRepId\n",
      "1 Luís Gonçalves Embraer - Empresa Brasileira de Aeronáutica S.A. Av. Brigadeiro Faria Lima, 2170 São José dos Campos SP Brazil 12227-000 +55 (12) 3923-5555 +55 (12) 3923-5566 luisg@embraer.com.br 3\n",
      "2 Leonie Köhler None Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 +49 0711 2842222 None leonekohler@surfeu.de 5\n",
      "3 François Tremblay None 1498 rue Bélanger Montréal QC Canada H2G 1A7 +1 (514) 721-4711 None ftremblay@gmail.com 3\n",
      "\n",
      "\n",
      "CREATE TABLE \"Invoice\" (\n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceDate\" DATETIME NOT NULL, \n",
      "\t\"BillingAddress\" NVARCHAR(70), \n",
      "\t\"BillingCity\" NVARCHAR(40), \n",
      "\t\"BillingState\" NVARCHAR(40), \n",
      "\t\"BillingCountry\" NVARCHAR(40), \n",
      "\t\"BillingPostalCode\" NVARCHAR(10), \n",
      "\t\"Total\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceId\"), \n",
      "\tFOREIGN KEY(\"CustomerId\") REFERENCES \"Customer\" (\"CustomerId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'Invoice' LIMIT 3;\n",
      "InvoiceId CustomerId InvoiceDate BillingAddress BillingCity BillingState BillingCountry BillingPostalCode Total\n",
      "1 2 2009-01-01 00:00:00 Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 1.98\n",
      "2 4 2009-01-02 00:00:00 Ullevålsveien 14 Oslo None Norway 0171 3.96\n",
      "3 8 2009-01-03 00:00:00 Grétrystraat 63 Brussels None Belgium 1000 5.94\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should query the Invoice and Customer tables to get the total sales per country.\n",
      "Action: query_sql_db\n",
      "Action Input: SELECT c.Country, SUM(i.Total) AS TotalSales FROM Invoice i INNER JOIN Customer c ON i.CustomerId = c.CustomerId GROUP BY c.Country ORDER BY TotalSales DESC LIMIT 10\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m[('USA', 523.0600000000003), ('Canada', 303.9599999999999), ('France', 195.09999999999994), ('Brazil', 190.09999999999997), ('Germany', 156.48), ('United Kingdom', 112.85999999999999), ('Czech Republic', 90.24000000000001), ('Portugal', 77.23999999999998), ('India', 75.25999999999999), ('Chile', 46.62)]\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer\n",
      "Final Answer: The customers from the USA spent the most, with a total of $523.06.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The customers from the USA spent the most, with a total of $523.06.'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\n",
    "    \"List the total sales per country. Which country's customers spent the most?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "474dddda-c067-4eeb-98b1-e763ee78b18c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mAction: list_tables_sql_db\n",
      "Action Input: \"\"\u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mInvoice, MediaType, Artist, InvoiceLine, Genre, Playlist, Employee, Album, PlaylistTrack, Track, Customer\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should look at the schema of the Playlist and PlaylistTrack tables to see what columns I can use.\n",
      "Action: schema_sql_db\n",
      "Action Input: \"Playlist, PlaylistTrack\"\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3m\n",
      "CREATE TABLE \"Playlist\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'Playlist' LIMIT 3;\n",
      "PlaylistId Name\n",
      "1 Music\n",
      "2 Movies\n",
      "3 TV Shows\n",
      "\n",
      "\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'PlaylistTrack' LIMIT 3;\n",
      "PlaylistId TrackId\n",
      "1 3402\n",
      "1 3389\n",
      "1 3390\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I can use a SELECT statement to get the total number of tracks in each playlist.\n",
      "Action: query_checker_sql_db\n",
      "Action Input: SELECT Playlist.Name, COUNT(PlaylistTrack.TrackId) AS TotalTracks FROM Playlist INNER JOIN PlaylistTrack ON Playlist.PlaylistId = PlaylistTrack.PlaylistId GROUP BY Playlist.Name\u001b[0m\n",
      "Observation: \u001b[31;1m\u001b[1;3m\n",
      "\n",
      "SELECT Playlist.Name, COUNT(PlaylistTrack.TrackId) AS TotalTracks FROM Playlist INNER JOIN PlaylistTrack ON Playlist.PlaylistId = PlaylistTrack.PlaylistId GROUP BY Playlist.Name\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m The query looks correct, I can now execute it.\n",
      "Action: query_sql_db\n",
      "Action Input: SELECT Playlist.Name, COUNT(PlaylistTrack.TrackId) AS TotalTracks FROM Playlist INNER JOIN PlaylistTrack ON Playlist.PlaylistId = PlaylistTrack.PlaylistId GROUP BY Playlist.Name LIMIT 10\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m[('90’s Music', 1477), ('Brazilian Music', 39), ('Classical', 75), ('Classical 101 - Deep Cuts', 25), ('Classical 101 - Next Steps', 25), ('Classical 101 - The Basics', 25), ('Grunge', 15), ('Heavy Metal Classic', 26), ('Music', 6580), ('Music Videos', 1)]\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer.\n",
      "Final Answer: The total number of tracks in each playlist are: '90’s Music' (1477), 'Brazilian Music' (39), 'Classical' (75), 'Classical 101 - Deep Cuts' (25), 'Classical 101 - Next Steps' (25), 'Classical 101 - The Basics' (25), 'Grunge' (15), 'Heavy Metal Classic' (26), 'Music' (6580), 'Music Videos' (1).\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"The total number of tracks in each playlist are: '90’s Music' (1477), 'Brazilian Music' (39), 'Classical' (75), 'Classical 101 - Deep Cuts' (25), 'Classical 101 - Next Steps' (25), 'Classical 101 - The Basics' (25), 'Grunge' (15), 'Heavy Metal Classic' (26), 'Music' (6580), 'Music Videos' (1).\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\n",
    "    \"Show the total number of tracks in each playlist. The Playlist name should be included in the result.\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7c7503b5-d9d9-4faa-b064-29fcdb5ff213",
   "metadata": {},
   "source": [
    "## Recovering from an error\n",
    "\n",
    "In this example, the agent is able to recover from an error after initially trying to access an attribute (`Track.ArtistId`) which doesn't exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9fe4901e-f9e1-4022-b6bc-80e2b2d6a3a4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mAction: list_tables_sql_db\n",
      "Action Input: \"\"\u001b[0m\n",
      "Observation: \u001b[38;5;200m\u001b[1;3mMediaType, Track, Invoice, Album, Playlist, Customer, Employee, InvoiceLine, PlaylistTrack, Genre, Artist\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should look at the schema of the Artist, InvoiceLine, and Track tables to see what columns I can use.\n",
      "Action: schema_sql_db\n",
      "Action Input: \"Artist, InvoiceLine, Track\"\u001b[0m\n",
      "Observation: \u001b[33;1m\u001b[1;3m\n",
      "CREATE TABLE \"Artist\" (\n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"ArtistId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'Artist' LIMIT 3;\n",
      "ArtistId Name\n",
      "1 AC/DC\n",
      "2 Accept\n",
      "3 Aerosmith\n",
      "\n",
      "\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",
      "SELECT * FROM 'Track' LIMIT 3;\n",
      "TrackId Name AlbumId MediaTypeId GenreId Composer Milliseconds Bytes UnitPrice\n",
      "1 For Those About To Rock (We Salute You) 1 1 1 Angus Young, Malcolm Young, Brian Johnson 343719 11170334 0.99\n",
      "2 Balls to the Wall 2 2 1 None 342562 5510424 0.99\n",
      "3 Fast As a Shark 3 2 1 F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman 230619 3990994 0.99\n",
      "\n",
      "\n",
      "CREATE TABLE \"InvoiceLine\" (\n",
      "\t\"InvoiceLineId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\t\"Quantity\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceLineId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"InvoiceId\") REFERENCES \"Invoice\" (\"InvoiceId\")\n",
      ")\n",
      "\n",
      "SELECT * FROM 'InvoiceLine' LIMIT 3;\n",
      "InvoiceLineId InvoiceId TrackId UnitPrice Quantity\n",
      "1 1 2 0.99 1\n",
      "2 1 4 0.99 1\n",
      "3 2 6 0.99 1\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should query the database to get the top 3 best selling artists.\n",
      "Action: query_sql_db\n",
      "Action Input: SELECT Artist.Name, SUM(InvoiceLine.Quantity) AS TotalQuantity FROM Artist INNER JOIN Track ON Artist.ArtistId = Track.ArtistId INNER JOIN InvoiceLine ON Track.TrackId = InvoiceLine.TrackId GROUP BY Artist.Name ORDER BY TotalQuantity DESC LIMIT 3\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mError: (sqlite3.OperationalError) no such column: Track.ArtistId\n",
      "[SQL: SELECT Artist.Name, SUM(InvoiceLine.Quantity) AS TotalQuantity FROM Artist INNER JOIN Track ON Artist.ArtistId = Track.ArtistId INNER JOIN InvoiceLine ON Track.TrackId = InvoiceLine.TrackId GROUP BY Artist.Name ORDER BY TotalQuantity DESC LIMIT 3]\n",
      "(Background on this error at: https://sqlalche.me/e/14/e3q8)\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should double check my query before executing it.\n",
      "Action: query_checker_sql_db\n",
      "Action Input: SELECT Artist.Name, SUM(InvoiceLine.Quantity) AS TotalQuantity FROM Artist INNER JOIN Track ON Artist.ArtistId = Track.ArtistId INNER JOIN InvoiceLine ON Track.TrackId = InvoiceLine.TrackId GROUP BY Artist.Name ORDER BY TotalQuantity DESC LIMIT 3\u001b[0m\n",
      "Observation: \u001b[31;1m\u001b[1;3m\n",
      "\n",
      "SELECT Artist.Name, SUM(InvoiceLine.Quantity) AS TotalQuantity \n",
      "FROM Artist \n",
      "INNER JOIN Track ON Artist.ArtistId = Track.ArtistId \n",
      "INNER JOIN InvoiceLine ON Track.TrackId = InvoiceLine.TrackId \n",
      "GROUP BY Artist.Name \n",
      "ORDER BY TotalQuantity DESC \n",
      "LIMIT 3;\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer.\n",
      "Action: query_sql_db\n",
      "Action Input: SELECT Artist.Name, SUM(InvoiceLine.Quantity) AS TotalQuantity FROM Artist INNER JOIN Album ON Artist.ArtistId = Album.ArtistId INNER JOIN Track ON Album.AlbumId = Track.AlbumId INNER JOIN InvoiceLine ON Track.TrackId = InvoiceLine.TrackId GROUP BY Artist.Name ORDER BY TotalQuantity DESC LIMIT 3\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3m[('Iron Maiden', 140), ('U2', 107), ('Metallica', 91)]\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer.\n",
      "Final Answer: The top 3 best selling artists are Iron Maiden, U2, and Metallica.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The top 3 best selling artists are Iron Maiden, U2, and Metallica.'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\"Who are the top 3 best selling artists?\")"
   ]
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}