File size: 22,779 Bytes
cfd3735
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
{
 "cells": [
  {
   "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://langchain.readthedocs.io/en/latest/modules/chains/examples/sqlite.html) 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."
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "090f3699-79c6-4ce1-ab96-a94f0121fd64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"sqlite:///../../../../notebooks/Chinook.db\")\n",
    "toolkit = SQLDatabaseToolkit(db=db)\n",
    "\n",
    "agent_executor = create_sql_agent(\n",
    "    llm=OpenAI(temperature=0),\n",
    "    toolkit=toolkit,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36ae48c7-cb08-4fef-977e-c7d4b96a464b",
   "metadata": {},
   "source": [
    "## Example: describing a table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ff70e83d-5ad0-4fc7-bb96-27d82ac166d7",
   "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;3mArtist, Invoice, Playlist, Genre, Album, PlaylistTrack, Track, InvoiceLine, MediaType, Employee, Customer\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I should look at the schema of the playlisttrack table\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 has two columns, PlaylistId and TrackId, and is linked to the Playlist and Track tables.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The PlaylistTrack table has two columns, PlaylistId and TrackId, and is linked to the Playlist and Track tables.'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\"Describe the playlisttrack table\")"
   ]
  },
  {
   "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\")"
   ]
  },
  {
   "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 [email protected] 3\n",
      "2 Leonie Köhler None Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 +49 0711 2842222 None [email protected] 5\n",
      "3 François Tremblay None 1498 rue Bélanger Montréal QC Canada H2G 1A7 +1 (514) 721-4711 None [email protected] 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(\"List the total sales per country. Which country's customers spent the most?\")"
   ]
  },
  {
   "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(\"Show the total number of tracks in each playlist. The Playlist name should be included in the result.\")"
   ]
  },
  {
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}