SQL Agent
Motivation
Enterprise customers often store and handle information in relational databases but querying such databases effectively requires bespoke knowledge of the underlying database structure as well as strong SQL coding skills. One way to address these challenges is to build an LLM agent capable of generating and executing SQL queries based on natural language. For example, if a user asks: what are the top 4 rows in table X
, the agent should be able to generate SELECT * FROM X LIMIT 4
, execute this query and return the output to the user.
Objective
In this notebook we explore how to setup a Cohere ReAct Agent to answer questions over SQL Databases. We show how this can be done seamlessly with langchain’s existing SQLDBToolkit.
Toolkit Setup
1 from langchain.agents import AgentExecutor 2 from langchain_cohere.react_multi_hop.agent import create_cohere_react_agent 3 from langchain_core.prompts import ChatPromptTemplate 4 from langchain_cohere.chat_models import ChatCohere 5 from langchain_community.utilities.sql_database import SQLDatabase 6 from langchain_community.agent_toolkits import SQLDatabaseToolkit 7 import os 8 import json
1 # Uncomment if you need to install the following packages 2 #!pip install --quiet langchain langchain_cohere langchain_experimental --upgrade
Langchain already has a SQLDBToolkit that consists of 4 tools to handle SQL query generation, execution and validation. To use this, you simply need to define your LLM and DB and pass these into the Toolkit.
These are the following tools:
- ‘sql_db_query’: executes SQL code on the database
- ’sql_db_schema’: returns the schema of tables given the list of tables
- ’sql_db_list_tables’: lists the tables in the database
- ’sql_db_query_checker’: validates the SQL query
1 # load the cohere api key 2 os.environ["COHERE_API_KEY"] = ""
1 DB_NAME='Chinook.db' 2 MODEL="command-r-plus" 3 llm = ChatCohere(model=MODEL, temperature=0.1,verbose=True) 4 db = SQLDatabase.from_uri(f"sqlite:///{DB_NAME}") 5 toolkit = SQLDatabaseToolkit(db=db, llm=llm) 6 context = toolkit.get_context() 7 tools = toolkit.get_tools() 8 9 print('**List of pre-defined Langchain Tools**') 10 print([tool.name for tool in tools])
**List of pre-defined Langchain Tools** ['sql_db_query', 'sql_db_schema', 'sql_db_list_tables', 'sql_db_query_checker']
SQL Agent
We follow the general cohere react agent setup in Langchain to build our SQL agent.
1 # define the prompt template 2 prompt = ChatPromptTemplate.from_template("{input}") 3 # instantiate the ReAct agent 4 agent = create_cohere_react_agent( 5 llm=llm, 6 tools=tools, 7 prompt=prompt, 8 ) 9 agent_executor = AgentExecutor(agent=agent, 10 tools=tools, 11 verbose=True, 12 return_intermediate_steps=True 13 )
1 output=agent_executor.invoke({ 2 "input": 'what tables are available?', 3 })
[1m> Entering new AgentExecutor chain...[0m [32;1m[1;3m I will use the sql_db_list_tables tool to find out which tables are available. {'tool_name': 'sql_db_list_tables', 'parameters': {'tool_input': ''}} [0m[38;5;200m[1;3mAlbum, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track[0m[32;1m[1;3mRelevant Documents: 0 Cited Documents: 0 Answer: The following tables are available: Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track. Grounded answer: The following tables are available: <co: 0="">Album</co:>, <co: 0="">Artist</co:>, <co: 0="">Customer</co:>, <co: 0="">Employee</co:>, <co: 0="">Genre</co:>, <co: 0="">Invoice</co:>, <co: 0="">InvoiceLine</co:>, <co: 0="">MediaType</co:>, <co: 0="">Playlist</co:>, <co: 0="">PlaylistTrack</co:>, <co: 0="">Track</co:>.[0m [1m> Finished chain.[0m
1 print(output['output'])
The following tables are available: Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track.
The agent uses the list_tables tool to effectively highlight all the tables in the DB.
1 output=agent_executor.invoke({ 2 "input": 'show the first row of the Playlist and Genre tables?', 3 })
[1m> Entering new AgentExecutor chain...[0m [32;1m[1;3m I will use the sql_db_schema tool to find the first row of the Playlist and Genre tables. {'tool_name': 'sql_db_schema', 'parameters': {'table_names': 'Playlist, Genre'}} [0m[33;1m[1;3m CREATE TABLE "Genre" ( "GenreId" INTEGER NOT NULL, "Name" NVARCHAR(120), PRIMARY KEY ("GenreId") ) /* 3 rows from Genre table: GenreId Name 1 Rock 2 Jazz 3 Metal */ CREATE TABLE "Playlist" ( "PlaylistId" INTEGER NOT NULL, "Name" NVARCHAR(120), PRIMARY KEY ("PlaylistId") ) /* 3 rows from Playlist table: PlaylistId Name 1 Music 2 Movies 3 TV Shows */[0m[32;1m[1;3mRelevant Documents: 0 Cited Documents: 0 Answer: Here is the first row of the Genre table: | GenreId | Name | |---|---| | 1 | Rock | Here is the first row of the Playlist table: | PlaylistId | Name | |---|---| | 1 | Music | Grounded answer: Here is the first row of the Genre table: | <co: 0="">GenreId</co:> | <co: 0="">Name</co:> | |---|---| | <co: 0="">1</co:> | <co: 0="">Rock</co:> | Here is the first row of the Playlist table: | <co: 0="">PlaylistId</co:> | <co: 0="">Name</co:> | |---|---| | <co: 0="">1</co:> | <co: 0="">Music</co:> |[0m [1m> Finished chain.[0m
1 print(output['output'])
Here is the first row of the Genre table: | GenreId | Name | |---|---| | 1 | Rock | Here is the first row of the Playlist table: | PlaylistId | Name | |---|---| | 1 | Music |
Here we see that the tool takes a list of tables to query the sql_db_schema tool to retrieve the various schemas.
1 output=agent_executor.invoke({ 2 "input": 'which countries have the most invoices?', 3 })
[1m> Entering new AgentExecutor chain...[0m [32;1m[1;3m I will search for the number of invoices per country and then write an answer. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT country, COUNT(*) AS invoice_count FROM invoices GROUP BY country ORDER BY invoice_count DESC'}} [0m[36;1m[1;3mError: (sqlite3.OperationalError) no such table: invoices [SQL: SELECT country, COUNT(*) AS invoice_count FROM invoices GROUP BY country ORDER BY invoice_count DESC] (Background on this error at: https://sqlalche.me/e/20/e3q8)[0m[32;1m[1;3m I received an error message. I will now use the sql_db_list_tables tool to find out the correct table name. {'tool_name': 'sql_db_list_tables', 'parameters': {'tool_input': ''}} [0m[38;5;200m[1;3mAlbum, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track[0m[32;1m[1;3m I found out that the correct table name is 'Invoice'. I will now rewrite my query and run it again. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT country, COUNT(*) AS invoice_count FROM Invoice GROUP BY country ORDER BY invoice_count DESC'}} [0m[36;1m[1;3mError: (sqlite3.OperationalError) no such column: country [SQL: SELECT country, COUNT(*) AS invoice_count FROM Invoice GROUP BY country ORDER BY invoice_count DESC] (Background on this error at: https://sqlalche.me/e/20/e3q8)[0m[32;1m[1;3m I received another error message. I will now use the sql_db_schema tool to find out the correct column name. {'tool_name': 'sql_db_schema', 'parameters': {'table_names': 'Invoice'}} [0m[33;1m[1;3m CREATE TABLE "Invoice" ( "InvoiceId" INTEGER NOT NULL, "CustomerId" INTEGER NOT NULL, "InvoiceDate" DATETIME NOT NULL, "BillingAddress" NVARCHAR(70), "BillingCity" NVARCHAR(40), "BillingState" NVARCHAR(40), "BillingCountry" NVARCHAR(40), "BillingPostalCode" NVARCHAR(10), "Total" NUMERIC(10, 2) NOT NULL, PRIMARY KEY ("InvoiceId"), FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId") ) /* 3 rows from Invoice table: InvoiceId CustomerId InvoiceDate BillingAddress BillingCity BillingState BillingCountry BillingPostalCode Total 1 2 2021-01-01 00:00:00 Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 1.98 2 4 2021-01-02 00:00:00 Ullevålsveien 14 Oslo None Norway 0171 3.96 3 8 2021-01-03 00:00:00 Grétrystraat 63 Brussels None Belgium 1000 5.94 */[0m[32;1m[1;3m I found out that the correct column name is 'BillingCountry'. I will now rewrite my query and run it again. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT BillingCountry AS country, COUNT(*) AS invoice_count FROM Invoice GROUP BY BillingCountry ORDER BY invoice_count DESC'}} [0m[36;1m[1;3m[('USA', 91), ('Canada', 56), ('France', 35), ('Brazil', 35), ('Germany', 28), ('United Kingdom', 21), ('Portugal', 14), ('Czech Republic', 14), ('India', 13), ('Sweden', 7), ('Spain', 7), ('Poland', 7), ('Norway', 7), ('Netherlands', 7), ('Italy', 7), ('Ireland', 7), ('Hungary', 7), ('Finland', 7), ('Denmark', 7), ('Chile', 7), ('Belgium', 7), ('Austria', 7), ('Australia', 7), ('Argentina', 7)][0m[32;1m[1;3mRelevant Documents: 1,3,4 Cited Documents: 4 Answer: The countries with the most invoices are the USA (91), Canada (56), and France (35). Grounded answer: The countries with the most invoices are the <co: 4="">USA (91</co:>), <co: 4="">Canada (56</co:>), and <co: 4="">France (35</co:>).[0m [1m> Finished chain.[0m
1 print(output['output'])
The countries with the most invoices are the USA (91), Canada (56), and France (35).
The agent initially makes some errors as it jumps to answer the question using the db_query tool, but it then realizes it needs to figure out what tables it has access to and what they look like. It then fixes the SQL code and is able to generate the right answer.
1 output=agent_executor.invoke({ 2 "input": 'who is the best customer? The customer who has spent the most money is the best.', 3 })
[1m> Entering new AgentExecutor chain...[0m [32;1m[1;3m I will run a SQL query to find the customer who has spent the most money. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT customer_name, SUM(total_cost) AS total_spent FROM orders GROUP BY customer_name ORDER BY total_spent DESC LIMIT 1;'}} [0m[36;1m[1;3mError: (sqlite3.OperationalError) no such table: orders [SQL: SELECT customer_name, SUM(total_cost) AS total_spent FROM orders GROUP BY customer_name ORDER BY total_spent DESC LIMIT 1;] (Background on this error at: https://sqlalche.me/e/20/e3q8)[0m[32;1m[1;3m I received an error message saying that there is no table called 'orders'. I will now use the sql_db_list_tables tool to find out what tables are available. {'tool_name': 'sql_db_list_tables', 'parameters': {'tool_input': ''}} [0m[38;5;200m[1;3mAlbum, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track[0m[32;1m[1;3m I found that there is a table called 'Customer'. I will now use the sql_db_schema tool to find out what columns are in the 'Customer' table. {'tool_name': 'sql_db_schema', 'parameters': {'table_names': 'Customer'}} [0m[33;1m[1;3m CREATE TABLE "Customer" ( "CustomerId" INTEGER NOT NULL, "FirstName" NVARCHAR(40) NOT NULL, "LastName" NVARCHAR(20) NOT NULL, "Company" NVARCHAR(80), "Address" NVARCHAR(70), "City" NVARCHAR(40), "State" NVARCHAR(40), "Country" NVARCHAR(40), "PostalCode" NVARCHAR(10), "Phone" NVARCHAR(24), "Fax" NVARCHAR(24), "Email" NVARCHAR(60) NOT NULL, "SupportRepId" INTEGER, PRIMARY KEY ("CustomerId"), FOREIGN KEY("SupportRepId") REFERENCES "Employee" ("EmployeeId") ) /* 3 rows from Customer table: CustomerId FirstName LastName Company Address City State Country PostalCode Phone Fax Email SupportRepId 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 2 Leonie Köhler None Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 +49 0711 2842222 None leonekohler@surfeu.de 5 3 François Tremblay None 1498 rue Bélanger Montréal QC Canada H2G 1A7 +1 (514) 721-4711 None ftremblay@gmail.com 3 */[0m[32;1m[1;3m I found that the 'Customer' table does not contain any information about how much money a customer has spent. I will now use the sql_db_schema tool to find out what columns are in the 'Invoice' table. {'tool_name': 'sql_db_schema', 'parameters': {'table_names': 'Invoice'}} [0m[33;1m[1;3m CREATE TABLE "Invoice" ( "InvoiceId" INTEGER NOT NULL, "CustomerId" INTEGER NOT NULL, "InvoiceDate" DATETIME NOT NULL, "BillingAddress" NVARCHAR(70), "BillingCity" NVARCHAR(40), "BillingState" NVARCHAR(40), "BillingCountry" NVARCHAR(40), "BillingPostalCode" NVARCHAR(10), "Total" NUMERIC(10, 2) NOT NULL, PRIMARY KEY ("InvoiceId"), FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId") ) /* 3 rows from Invoice table: InvoiceId CustomerId InvoiceDate BillingAddress BillingCity BillingState BillingCountry BillingPostalCode Total 1 2 2021-01-01 00:00:00 Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 1.98 2 4 2021-01-02 00:00:00 Ullevålsveien 14 Oslo None Norway 0171 3.96 3 8 2021-01-03 00:00:00 Grétrystraat 63 Brussels None Belgium 1000 5.94 */[0m[32;1m[1;3m I found that the 'Invoice' table contains a 'Total' column, which is likely to be the total amount spent by the customer on that invoice. I will now use the sql_db_query tool to find the customer who has spent the most money. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT c.FirstName, c.LastName, SUM(i.Total) AS total_spent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.CustomerId ORDER BY total_spent DESC LIMIT 1;'}} [0m[36;1m[1;3m[('Helena', 'Holý', 49.62)][0m[32;1m[1;3mRelevant Documents: 1,2,3,4 Cited Documents: 4 Answer: The best customer is Helena Holý, who has spent a total of 49.62. Grounded answer: The best customer is <co: 4="">Helena Holý</co:>, who has spent a total of <co: 4="">49.62</co:>.[0m [1m> Finished chain.[0m
1 print(output['output'])
The best customer is Helena Holý, who has spent a total of 49.62.
As you can see, the agent makes an error, but is able to rectify itself. It also manages to generate a SQL query over two tables in the database.
SQL Agent with context
Generally, passing in additional context to the preamble can help reduce the initial failures. This context is provided by the SQLDBToolkit and contains the first 3 rows of the tables in the Database.
1 print('**Context to pass to LLM on tables**') 2 print('Table Names') 3 print(context['table_names']) 4 print('Table Schemas') 5 print(context['table_info'])
1 **Context to pass to LLM on tables** 2 Table Names 3 Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track 4 Table Schemas 5 6 CREATE TABLE "Album" ( 7 "AlbumId" INTEGER NOT NULL, 8 "Title" NVARCHAR(160) NOT NULL, 9 "ArtistId" INTEGER NOT NULL, 10 PRIMARY KEY ("AlbumId"), 11 FOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId") 12 ) 13 14 /* 15 3 rows from Album table: 16 AlbumId Title ArtistId 17 1 For Those About To Rock We Salute You 1 18 2 Balls to the Wall 2 19 3 Restless and Wild 2 20 */ 21 22 23 CREATE TABLE "Artist" ( 24 "ArtistId" INTEGER NOT NULL, 25 "Name" NVARCHAR(120), 26 PRIMARY KEY ("ArtistId") 27 ) 28 29 /* 30 3 rows from Artist table: 31 ArtistId Name 32 1 AC/DC 33 2 Accept 34 3 Aerosmith 35 */ 36 37 38 CREATE TABLE "Customer" ( 39 "CustomerId" INTEGER NOT NULL, 40 "FirstName" NVARCHAR(40) NOT NULL, 41 "LastName" NVARCHAR(20) NOT NULL, 42 "Company" NVARCHAR(80), 43 "Address" NVARCHAR(70), 44 "City" NVARCHAR(40), 45 "State" NVARCHAR(40), 46 "Country" NVARCHAR(40), 47 "PostalCode" NVARCHAR(10), 48 "Phone" NVARCHAR(24), 49 "Fax" NVARCHAR(24), 50 "Email" NVARCHAR(60) NOT NULL, 51 "SupportRepId" INTEGER, 52 PRIMARY KEY ("CustomerId"), 53 FOREIGN KEY("SupportRepId") REFERENCES "Employee" ("EmployeeId") 54 ) 55 56 /* 57 3 rows from Customer table: 58 CustomerId FirstName LastName Company Address City State Country PostalCode Phone Fax Email SupportRepId 59 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 60 2 Leonie Köhler None Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 +49 0711 2842222 None leonekohler@surfeu.de 5 61 3 François Tremblay None 1498 rue Bélanger Montréal QC Canada H2G 1A7 +1 (514) 721-4711 None ftremblay@gmail.com 3 62 */ 63 64 65 CREATE TABLE "Employee" ( 66 "EmployeeId" INTEGER NOT NULL, 67 "LastName" NVARCHAR(20) NOT NULL, 68 "FirstName" NVARCHAR(20) NOT NULL, 69 "Title" NVARCHAR(30), 70 "ReportsTo" INTEGER, 71 "BirthDate" DATETIME, 72 "HireDate" DATETIME, 73 "Address" NVARCHAR(70), 74 "City" NVARCHAR(40), 75 "State" NVARCHAR(40), 76 "Country" NVARCHAR(40), 77 "PostalCode" NVARCHAR(10), 78 "Phone" NVARCHAR(24), 79 "Fax" NVARCHAR(24), 80 "Email" NVARCHAR(60), 81 PRIMARY KEY ("EmployeeId"), 82 FOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId") 83 ) 84 85 /* 86 3 rows from Employee table: 87 EmployeeId LastName FirstName Title ReportsTo BirthDate HireDate Address City State Country PostalCode Phone Fax Email 88 1 Adams Andrew General Manager None 1962-02-18 00:00:00 2002-08-14 00:00:00 11120 Jasper Ave NW Edmonton AB Canada T5K 2N1 +1 (780) 428-9482 +1 (780) 428-3457 andrew@chinookcorp.com 89 2 Edwards Nancy Sales Manager 1 1958-12-08 00:00:00 2002-05-01 00:00:00 825 8 Ave SW Calgary AB Canada T2P 2T3 +1 (403) 262-3443 +1 (403) 262-3322 nancy@chinookcorp.com 90 3 Peacock Jane Sales Support Agent 2 1973-08-29 00:00:00 2002-04-01 00:00:00 1111 6 Ave SW Calgary AB Canada T2P 5M5 +1 (403) 262-3443 +1 (403) 262-6712 jane@chinookcorp.com 91 */ 92 93 94 CREATE TABLE "Genre" ( 95 "GenreId" INTEGER NOT NULL, 96 "Name" NVARCHAR(120), 97 PRIMARY KEY ("GenreId") 98 ) 99 100 /* 101 3 rows from Genre table: 102 GenreId Name 103 1 Rock 104 2 Jazz 105 3 Metal 106 */ 107 108 109 CREATE TABLE "Invoice" ( 110 "InvoiceId" INTEGER NOT NULL, 111 "CustomerId" INTEGER NOT NULL, 112 "InvoiceDate" DATETIME NOT NULL, 113 "BillingAddress" NVARCHAR(70), 114 "BillingCity" NVARCHAR(40), 115 "BillingState" NVARCHAR(40), 116 "BillingCountry" NVARCHAR(40), 117 "BillingPostalCode" NVARCHAR(10), 118 "Total" NUMERIC(10, 2) NOT NULL, 119 PRIMARY KEY ("InvoiceId"), 120 FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId") 121 ) 122 123 /* 124 3 rows from Invoice table: 125 InvoiceId CustomerId InvoiceDate BillingAddress BillingCity BillingState BillingCountry BillingPostalCode Total 126 1 2 2021-01-01 00:00:00 Theodor-Heuss-Straße 34 Stuttgart None Germany 70174 1.98 127 2 4 2021-01-02 00:00:00 Ullevålsveien 14 Oslo None Norway 0171 3.96 128 3 8 2021-01-03 00:00:00 Grétrystraat 63 Brussels None Belgium 1000 5.94 129 */ 130 131 132 CREATE TABLE "InvoiceLine" ( 133 "InvoiceLineId" INTEGER NOT NULL, 134 "InvoiceId" INTEGER NOT NULL, 135 "TrackId" INTEGER NOT NULL, 136 "UnitPrice" NUMERIC(10, 2) NOT NULL, 137 "Quantity" INTEGER NOT NULL, 138 PRIMARY KEY ("InvoiceLineId"), 139 FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"), 140 FOREIGN KEY("InvoiceId") REFERENCES "Invoice" ("InvoiceId") 141 ) 142 143 /* 144 3 rows from InvoiceLine table: 145 InvoiceLineId InvoiceId TrackId UnitPrice Quantity 146 1 1 2 0.99 1 147 2 1 4 0.99 1 148 3 2 6 0.99 1 149 */ 150 151 152 CREATE TABLE "MediaType" ( 153 "MediaTypeId" INTEGER NOT NULL, 154 "Name" NVARCHAR(120), 155 PRIMARY KEY ("MediaTypeId") 156 ) 157 158 /* 159 3 rows from MediaType table: 160 MediaTypeId Name 161 1 MPEG audio file 162 2 Protected AAC audio file 163 3 Protected MPEG-4 video file 164 */ 165 166 167 CREATE TABLE "Playlist" ( 168 "PlaylistId" INTEGER NOT NULL, 169 "Name" NVARCHAR(120), 170 PRIMARY KEY ("PlaylistId") 171 ) 172 173 /* 174 3 rows from Playlist table: 175 PlaylistId Name 176 1 Music 177 2 Movies 178 3 TV Shows 179 */ 180 181 182 CREATE TABLE "PlaylistTrack" ( 183 "PlaylistId" INTEGER NOT NULL, 184 "TrackId" INTEGER NOT NULL, 185 PRIMARY KEY ("PlaylistId", "TrackId"), 186 FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"), 187 FOREIGN KEY("PlaylistId") REFERENCES "Playlist" ("PlaylistId") 188 ) 189 190 /* 191 3 rows from PlaylistTrack table: 192 PlaylistId TrackId 193 1 3402 194 1 3389 195 1 3390 196 */ 197 198 199 CREATE TABLE "Track" ( 200 "TrackId" INTEGER NOT NULL, 201 "Name" NVARCHAR(200) NOT NULL, 202 "AlbumId" INTEGER, 203 "MediaTypeId" INTEGER NOT NULL, 204 "GenreId" INTEGER, 205 "Composer" NVARCHAR(220), 206 "Milliseconds" INTEGER NOT NULL, 207 "Bytes" INTEGER, 208 "UnitPrice" NUMERIC(10, 2) NOT NULL, 209 PRIMARY KEY ("TrackId"), 210 FOREIGN KEY("MediaTypeId") REFERENCES "MediaType" ("MediaTypeId"), 211 FOREIGN KEY("GenreId") REFERENCES "Genre" ("GenreId"), 212 FOREIGN KEY("AlbumId") REFERENCES "Album" ("AlbumId") 213 ) 214 215 /* 216 3 rows from Track table: 217 TrackId Name AlbumId MediaTypeId GenreId Composer Milliseconds Bytes UnitPrice 218 1 For Those About To Rock (We Salute You) 1 1 1 Angus Young, Malcolm Young, Brian Johnson 343719 11170334 0.99 219 2 Balls to the Wall 2 2 1 U. Dirkschneider, W. Hoffmann, H. Frank, P. Baltes, S. Kaufmann, G. Hoffmann 342562 5510424 0.99 220 3 Fast As a Shark 3 2 1 F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman 230619 3990994 0.99 221 */
We can pass this context into the preamble and re-run a query to see how it performs.
1 preamble="""## Task And Context 2 You use your advanced complex reasoning capabilities to help people by answering their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You may need to use multiple tools in parallel or sequentially to complete your task. You should focus on serving the user's needs as best you can, which will be wide-ranging. 3 4 ## Style Guide 5 Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling. 6 7 ## Additional Information 8 You are an expert who answers the user's question by creating SQL queries and executing them. 9 You are equipped with a number of relevant SQL tools. 10 11 Here is information about the database: 12 {schema_info} 13 """.format(schema_info=context)
1 output=agent_executor.invoke({ 2 "input": 'provide the name of the best customer? The customer who has spent the most money is the best.', 3 "preamble": preamble 4 })
[1m> Entering new AgentExecutor chain...[0m [32;1m[1;3m I will write a SQL query to find the customer who has spent the most money. {'tool_name': 'sql_db_query', 'parameters': {'query': 'SELECT c.FirstName, c.LastName, SUM(i.Total) AS TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.CustomerId ORDER BY TotalSpent DESC LIMIT 1;'}} [0m[36;1m[1;3m[('Helena', 'Holý', 49.62)][0m[32;1m[1;3mRelevant Documents: 0 Cited Documents: 0 Answer: The customer who has spent the most money is Helena Holý. Grounded answer: The customer who has spent the most money is <co: 0="">Helena Holý</co:>.[0m [1m> Finished chain.[0m
1 print(output['output'])
The customer who has spent the most money is Helena Holý.
We can see that passing that additional context actually avoids the error seen in the previous section and gets to the answer in one tool call. This works as long as you have a few tables and a few columns per table. We will follow up with more techniques to improve stability and scalability in future work.