text to sql prompt

text-to-sql.py

from langchain.utilities import SQLDatabase
from langchain.chat_models import ChatOpenAI
from langchain.schema.output_parser import StrOutputParser
from langchain import hub
from keys import OPENAI_API_KEY
import os

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

# Initialize database
db = SQLDatabase.from_uri("sqlite:///Chinook.db")
# Pull down prompt
prompt = hub.pull("rlm/text-to-sql")
# Initialize model
model = ChatOpenAI()

# Create chain with LangChain Expression Language
inputs = {
    "table_info": lambda x: db.get_table_info(),
    "input": lambda x: x["question"],
    "few_shot_examples": lambda x: "",
    "dialect": lambda x: db.dialect,
}
sql_response = (
    inputs
    | prompt
    | model.bind(stop=["\nSQLResult:"])
    | StrOutputParser()
)

# Call with a given question
response = sql_response.invoke({"question": "How many customers are there?"})

print(response)