langchain/docs/modules/chains/examples/sqlite.ipynb

SQLite example

This example showcases hooking up an LLM to answer questions over a database.

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.

In [1]:

from langchain import OpenAI, SQLDatabase, SQLDatabaseChain

In [2]:

db = SQLDatabase.from_uri("sqlite:///../../../../notebooks/Chinook.db")
llm = OpenAI(temperature=0)

In [3]:

db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)

In [4]:

db_chain.run("How many employees are there?")

Out [4]:

> Entering new SQLDatabaseChain chain...
How many employees are there? 
SQLQuery: SELECT COUNT(*) FROM Employee;
SQLResult: [(9,)]
Answer: There are 9 employees.
> Finished chain.

' There are 9 employees.'

Customize Prompt

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

In [5]:

from langchain.prompts.prompt import PromptTemplate

_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.
Use the following format:

Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"

Only use the following tables:

{table_info}

If someone asks for the table foobar, they really mean the employee table.

Question: {input}"""
PROMPT = PromptTemplate(
    input_variables=["input", "table_info", "dialect"], template=_DEFAULT_TEMPLATE
)

In [6]:

db_chain = SQLDatabaseChain(llm=llm, database=db, prompt=PROMPT, verbose=True)

In [7]:

db_chain.run("How many employees are there in the foobar table?")

Out [7]:

> Entering new SQLDatabaseChain chain...
How many employees are there in the foobar table? 
SQLQuery: SELECT COUNT(*) FROM Employee;
SQLResult: [(9,)]
Answer: There are 9 employees in the foobar table.
> Finished chain.

' There are 9 employees in the foobar table.'

Return Intermediate Steps

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.

In [8]:

db_chain = SQLDatabaseChain(llm=llm, database=db, prompt=PROMPT, verbose=True, return_intermediate_steps=True)

In [10]:

result = db_chain("How many employees are there in the foobar table?")
result["intermediate_steps"]

Out [10]:

> Entering new SQLDatabaseChain chain...
How many employees are there in the foobar table? 
SQLQuery: SELECT COUNT(*) FROM Employee;
SQLResult: [(9,)]
Answer: There are 9 employees in the foobar table.
> Finished chain.

[' SELECT COUNT(*) FROM Employee;', '[(9,)]']

Choosing how to limit the number of rows returned

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.

In [3]:

db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True, top_k=3)

In [8]:

db_chain.run("What are some example tracks by composer Johann Sebastian Bach?")

Out [8]:

> Entering new SQLDatabaseChain chain...
What are some example tracks by composer Johann Sebastian Bach? 
SQLQuery: SELECT Name FROM Track WHERE Composer = 'Johann Sebastian Bach' LIMIT 3;
SQLResult: [('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',)]
Answer: Examples of tracks by Johann Sebastian Bach include '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'.
> Finished chain.

' Examples of tracks by Johann Sebastian Bach include \'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\'.'

Adding first row of each table

Sometimes, the format of the data is not obvious and it is optimal to include the first row of the table 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.

In [11]:

db = SQLDatabase.from_uri(
    "sqlite:///../../../../notebooks/Chinook.db", 
    include_tables=['Track'], # we include only one table to save tokens in the prompt :)
    sample_row_in_table_info=True)

In [12]:

db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)

In [13]:

db_chain.run("What are some example tracks by Bach?")

Out [13]:

> Entering new SQLDatabaseChain chain...
What are some example tracks by Bach? 
SQLQuery:Table 'Track' has columns: TrackId (INTEGER), Name (NVARCHAR(200)), AlbumId (INTEGER), MediaTypeId (INTEGER), GenreId (INTEGER), Composer (NVARCHAR(220)), Milliseconds (INTEGER), Bytes (INTEGER), UnitPrice (NUMERIC(10, 2)). Here is an example row for this table (long strings are truncated): ['1', 'For Those About To Rock (We Salute You)', '1', '1', '1', 'Angus Young, Malcolm Young, Brian Johnson', '343719', '11170334', '0.99'].
 SELECT TrackId, Name, Composer FROM Track WHERE Composer LIKE '%Bach%' ORDER BY Name LIMIT 5;
SQLResult: [(1709, 'American Woman', 'B. Cummings/G. Peterson/M.J. Kale/R. Bachman'), (3408, 'Aria Mit 30 Veränderungen, BWV 988 "Goldberg Variations": Aria', 'Johann Sebastian Bach'), (3433, 'Concerto No.2 in F Major, BWV1047, I. Allegro', 'Johann Sebastian Bach'), (3407, 'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', 'Johann Sebastian Bach'), (3490, 'Partita in E Major, BWV 1006A: I. Prelude', 'Johann Sebastian Bach')]
Answer: Some example tracks by Bach are 'American Woman', 'Aria Mit 30 Veränderungen, BWV 988 "Goldberg Variations": Aria', 'Concerto No.2 in F Major, BWV1047, I. Allegro', 'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace', and 'Partita in E Major, BWV 1006A: I. Prelude'.
> Finished chain.

' Some example tracks by Bach are \'American Woman\', \'Aria Mit 30 Veränderungen, BWV 988 "Goldberg Variations": Aria\', \'Concerto No.2 in F Major, BWV1047, I. Allegro\', \'Concerto for 2 Violins in D Minor, BWV 1043: I. Vivace\', and \'Partita in E Major, BWV 1006A: I. Prelude\'.'

SQLDatabaseSequentialChain

Chain for querying SQL database that is a sequential chain.

The chain is as follows:

1. Based on the query, determine which tables to use.
2. Based on those tables, call the normal SQL database chain.

This is useful in cases where the number of tables in the database is large.

In [3]:

from langchain.chains import SQLDatabaseSequentialChain

In [4]:

chain = SQLDatabaseSequentialChain.from_llm(llm, db, verbose=True)

In [5]:

chain.run("How many employees are also customers?")

Out [5]:

> Entering new SQLDatabaseSequentialChain chain...
Table names to use:
['Customer', 'Employee']

> Entering new SQLDatabaseChain chain...
How many employees are also customers? 
SQLQuery: SELECT COUNT(*) FROM Customer c INNER JOIN Employee e ON c.SupportRepId = e.EmployeeId;
SQLResult: [(59,)]
Answer: There are 59 employees who are also customers.
> Finished chain.

> Finished chain.

' There are 59 employees who are also customers.'