Quick tip: Using LangChain’s SQLDatabaseToolkit with SingleStoreDB

Abstract

LangChain is a powerful framework that includes a variety of tools, including the agent_toolkits. In this article, we'll use the SQLDatabaseToolkit to interact with SingleStoreDB by making a request in English.

Introduction

A recent blog post described an internal SingleStore Demothon. One of the cool demos showed how to use LangChain's SQLDatabaseToolkit with SingleStoreDB. In this short article, we'll replicate the demo.

Create a SingleStoreDB Cloud account

A previous article showed the steps to create a free SingleStoreDB Cloud account. We'll use LangChain Demo Group as our Workspace Group Name and langchain-demo as our Workspace Name. We'll make a note of our password and host name.

Create the TPC-H database

Once the workspace is available, we'll select the Tutorials (identified by the white question mark on a black background on the bottom right-hand side of the web page) and then choose Tutorials > Try with Sample Data > Benchmark Dataset - TPC-H. We'll check that the langchain-demo workspace is selected and then click the button to ▶ Load TPC-H Dataset, as shown in Figure 1.

Figure 1. Load TPC-H Dataset.

After a short time, the database will be ready. The database name will start with s2_dataset_tpch_ and be visible under the Workspace Overview tab. We'll make a note of the entire database name.

As described in a previous article, we'll follow the instructions to create a Notebook.

Fill out the Notebook

First, we'll install LangChain:

!pip install langchain --quiet

and add some imports:

from langchain import OpenAI, SQLDatabase
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import create_sql_agent

Next, we'll set our OpenAI API Key:

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

We'll now create the agent:

s2_password = "<password>"
s2_host = "<host>"
s2_db = "<database>"

db = SQLDatabase.from_uri(f"mysql+pymysql://admin:{s2_password}@{s2_host}:3306/{s2_db}")

llm = OpenAI(
    temperature = 0,
    verbose = False
)

toolkit = SQLDatabaseToolkit(
    db = db,
    llm = llm
)

agent_executor = create_sql_agent(
    llm = OpenAI(temperature = 0),
    toolkit = toolkit,
    verbose = True
)

We'll replace the <password>, <host> and <database> with the values we saved earlier.

Finally, we'll run the agent, as follows:

agent_executor.run("Write a SingleStore query that shows the highest paying customers per country, include how much they've spent, use the nation name")

The output should be similar to the following:

> Entering new AgentExecutor chain...
Thought: I should use the sql_db_list_tables tool to get a list of tables in the database.
Action: sql_db_list_tables
Action Input: ""
Observation: customer, lineitem, nation, orders, part, partsupp, region, supplier
Thought: I should query the schema of the customer, orders, and nation tables.
Action: sql_db_schema
Action Input: "customer, orders, nation"
Observation: 
CREATE TABLE customer (
 c_custkey INTEGER(11) NOT NULL, 
 c_name VARCHAR(25) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 c_address VARCHAR(40) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 c_nationkey INTEGER(11) NOT NULL, 
 c_phone VARCHAR(15) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 c_acctbal DECIMAL(15, 2) NOT NULL, 
 c_mktsegment VARCHAR(10) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 c_comment VARCHAR(117) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL
)

/*
3 rows from customer table:
c_custkey c_name c_address c_nationkey c_phone c_acctbal c_mktsegment c_comment
3 Customer#000000003 MG9kdTD2WBHm 1 11-719-748-3364 7498.12 AUTOMOBILE  deposits eat slyly ironic, even instructions. express foxes detect slyly. blithely even accounts ab
26 Customer#000000026 8ljrc5ZeMl7UciP 22 32-363-455-4837 5182.05 AUTOMOBILE c requests use furiously ironic requests. slyly ironic dependencies us
33 Customer#000000033 qFSlMuLucBmx9xnn5ib2csWUweg D 17 27-375-391-1280 -78.56 AUTOMOBILE s. slyly regular accounts are furiously. carefully pending requests
*/


CREATE TABLE nation (
 n_nationkey INTEGER(11) NOT NULL, 
 n_name VARCHAR(25) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 n_regionkey INTEGER(11) NOT NULL, 
 n_comment VARCHAR(152) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL
)

/*
3 rows from nation table:
n_nationkey n_name n_regionkey n_comment
3 CANADA 1 eas hang ironic, silent packages. slyly regular packages are furiously over the tithes. fluffily bol
22 RUSSIA 3  requests against the platelets use never according to the quickly regular pint
1 ARGENTINA 1 al foxes promise slyly according to the regular accounts. bold requests alon
*/


CREATE TABLE orders (
 o_orderkey BIGINT(11) NOT NULL, 
 o_custkey INTEGER(11) NOT NULL, 
 o_orderstatus CHAR(1) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 o_totalprice DECIMAL(15, 2) NOT NULL, 
 o_orderdate DATE NOT NULL, 
 o_orderpriority VARCHAR(15) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 o_clerk VARCHAR(15) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL, 
 o_shippriority INTEGER(11) NOT NULL, 
 o_comment VARCHAR(79) CHARACTER SET utf8 COLLATE utf8_bin NOT NULL
)

/*
3 rows from orders table:
o_orderkey o_custkey o_orderstatus o_totalprice o_orderdate o_orderpriority o_clerk o_shippriority o_comment
17349095 193084 O 73125.77 1996-07-04 4-NOT SPECIFIED Clerk#000002841 0 hely special somas cajole-
17349349 468115 F 230397.76 1992-04-16 1-URGENT Clerk#000002856 0 s. deposits cajole carefully regular ac
17349376 758954 F 199122.15 1993-10-27 1-URGENT Clerk#000000408 0  wake quickly. carefully pending deposits wak
*/
Thought: I should use the sql_db_query tool to query the database.
Action: sql_db_query
Action Input: SELECT c.c_name, c.c_acctbal, n.n_name FROM customer c JOIN nation n ON c.c_nationkey = n.n_nationkey ORDER BY c.c_acctbal DESC LIMIT 10
Observation: [('Customer#000061453', Decimal('9999.99'), 'MOROCCO'), ('Customer#000508503', Decimal('9999.99'), 'UNITED STATES'), ('Customer#001123705', Decimal('9999.98'), 'MOZAMBIQUE'), ('Customer#000399453', Decimal('9999.97'), 'UNITED KINGDOM'), ('Customer#000242308', Decimal('9999.96'), 'RUSSIA'), ('Customer#000348586', Decimal('9999.96'), 'MOROCCO'), ('Customer#000069321', Decimal('9999.96'), 'MOROCCO'), ('Customer#000991983', Decimal('9999.95'), 'IRAN'), ('Customer#001165335', Decimal('9999.94'), 'EGYPT'), ('Customer#000652672', Decimal('9999.93'), 'CANADA')]
Thought: I now know the final answer
Final Answer: The highest paying customers per country are: Customer#000061453 from Morocco with $9999.99, Customer#000508503 from United States with $9999.99, Customer#001123705 from Mozambique with $9999.98, Customer#000399453 from United Kingdom with $9999.97, Customer#000242308 from Russia with $9999.96, Customer#000348586 from Morocco with $9999.96, Customer#000069321 from Morocco with $9999.96, Customer#000991983 from Iran with $9999.95, Customer#001165335 from Egypt with $9999.94, and Customer#000652672 from Canada with $9999.93.

> Finished chain.
'The highest paying customers per country are: Customer#000061453 from Morocco with $9999.99, Customer#000508503 from United States with $9999.99, Customer#001123705 from Mozambique with $9999.98, Customer#000399453 from United Kingdom with $9999.97, Customer#000242308 from Russia with $9999.96, Customer#000348586 from Morocco with $9999.96, Customer#000069321 from Morocco with $9999.96, Customer#000991983 from Iran with $9999.95, Customer#001165335 from Egypt with $9999.94, and Customer#000652672 from Canada with $9999.93.'

Summary

Using the SQLDatabaseToolkit we can interact with a SingleStoreDB database using English. This provides an alternative to using SQL.

Acknowledgements

I thank my colleague Justin Kuntz for the demo code, which was adapted for this article.