我们需要下载一个 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 ChatOpenAI 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) 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;
