langchain结合智谱slm大模型操作数据库

本文介绍如何使用langchain结合大模型操作数据库

from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI


model = ChatOpenAI(model="glm-4-plus",
api_key='自己的api-key',
base_url='https://open.bigmodel.cn/api/paas/v4/', temperature=0)

HOSTNAME = "host"
PORT = "port"
DATABASE = "database"
USERNAME = "username"
PASSWORD = r"password"
MYSQL_URI = "mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8mb4".format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)
print(MYSQL_URI)

db = SQLDatabase.from_uri(MYSQL_URI)
# print(db.dialect)
# print(db.get_usable_table_names())


create_sql_chain = create_sql_query_chain(llm=model, db=db)

tool_chain = QuerySQLDataBaseTool(db=db)

prompt = PromptTemplate.from_template("""Given the following user question, corresponding SQL query, and SQL result, answer the user question,用中文回答
Question:{question}
SQL Query:{query}
SQL Result: {result}
Answer: """)

create_sql_chain = create_sql_chain | (lambda s: s.replace('```sql', '').replace('```', ''))


answer_chain = prompt | model | StrOutputParser()

chain = RunnablePassthrough.assign(query=create_sql_chain).assign(result=itemgetter('query') | tool_chain) | answer_chain

resp = chain.invoke({"question": "请问有多少个已经启用的自动化?"})
print(resp)

如果要替换成其他大模型,只要把模型参数,请求url替换成对应的即可