我们需要下载一个 LangChain
官方提供的本地小数据库。
安装依赖
SQL:
https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql
Shell:
pip install --upgrade --quiet langchain-core langchain-community langchain-openai
导入数据
我这里使用 Navicat
导入数据,你也可以通过别的方式导入(当然你有现成的数据库也可以,但是不要太大了,不然会消耗很多Token
)。
编写代码
这里我使用了 GPR 3.5 Turbo
,效果不理想的话可以试试GPT 4
或者 GPT 4 Turbo
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAItemplate = """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)db = SQLDatabase.from_uri("sqlite:///./Chinook.db")def get_schema(_):return db.get_table_info()def run_query(query):return db.run(query)model = ChatOpenAI(model="gpt-3.5-turbo",
)sql_response = (RunnablePassthrough.assign(schema=get_schema)| prompt| model.bind(stop=["\nSQLResult:"])| StrOutputParser()
)message = sql_response.invoke({"question": "How many employees are there?"})
print(f"message: {message}")
运行结果
➜ python3 test08.py
message: SELECT COUNT(*) AS totalEmployees
FROM Employee;