langchain/docs/extras/expression_language/cookbook/sql_db.ipynb

Cell:

[Cell type raw - unsupported, skipped]

We can replicate our SQLDatabaseChain with Runnables.

In [1]:

from langchain.prompts import ChatPromptTemplate

template = """Based on the table schema below, write a SQL query that would answer the user's question:
{schema}

Question: {question}
SQL Query:"""
prompt = ChatPromptTemplate.from_template(template)

In [2]:

from langchain.utilities import SQLDatabase

We'll need the Chinook sample DB for this example. There's many places to download it from, e.g. https://database.guide/2-sample-databases-sqlite/

In [20]:

db = SQLDatabase.from_uri("sqlite:///./Chinook.db")

In [21]:

def get_schema(_):
    return db.get_table_info()

In [22]:

def run_query(query):
    return db.run(query)

In [23]:

from operator import itemgetter

from langchain.chat_models import ChatOpenAI
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableLambda, RunnableMap

model = ChatOpenAI()

inputs = {
    "schema": RunnableLambda(get_schema),
    "question": itemgetter("question")
}
sql_response = (
        RunnableMap(inputs)
        | prompt
        | model.bind(stop=["\nSQLResult:"])
        | StrOutputParser()
    )

In [24]:

sql_response.invoke({"question": "How many employees are there?"})

Out [24]:

'SELECT COUNT(*) FROM Employee'

In [25]:

template = """Based on the table schema below, question, sql query, and sql response, write a natural language response:
{schema}

Question: {question}
SQL Query: {query}
SQL Response: {response}"""
prompt_response = ChatPromptTemplate.from_template(template)

In [26]:

full_chain = (
    RunnableMap({
        "question": itemgetter("question"),
        "query": sql_response,
    }) 
    | {
        "schema": RunnableLambda(get_schema),
        "question": itemgetter("question"),
        "query": itemgetter("query"),
        "response": lambda x: db.run(x["query"])    
    } 
    | prompt_response 
    | model
)

In [27]:

full_chain.invoke({"question": "How many employees are there?"})

Out [27]:

AIMessage(content='There are 8 employees.', additional_kwargs={}, example=False)

In [ ]: