92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
import os
|
||
from langchain_community.llms import QianfanLLMEndpoint,Tongyi
|
||
from langchain.chains import LLMChain
|
||
from langchain.prompts import PromptTemplate
|
||
import psycopg2
|
||
import re
|
||
|
||
def getSql(response):
|
||
pattern = r"#(.*?)#"
|
||
sql_list = re.findall(pattern, response)
|
||
sql_str = ''.join(sql_list)
|
||
print(sql_str)
|
||
|
||
return sql_str
|
||
|
||
|
||
def query_database(text, schema):
|
||
# 将表结构信息包含在提示中
|
||
schema_info = "\n".join([f"Table {table}: {', '.join(columns)}" for table, columns in schema.items()])
|
||
# print(schema_info)
|
||
prompt = PromptTemplate(
|
||
input_variables=["schema_info", "prompt"],
|
||
template='''以下是数据库的表结构信息：{schema_info}。
|
||
分析表结构信息，请根据以下描述生成一个符合SQLite数据库语法的SQL查询，描述：{prompt}。
|
||
并且要求输出的SQL以#开头,以#结尾，样例如下：
|
||
#SELECT * FROM table#
|
||
#SELECT COUNT(*) FROM table#
|
||
注意不要输出分析过程和其他内容，直接给出SQL语句。
|
||
'''
|
||
)
|
||
chain = LLMChain(llm=llm, prompt=prompt)
|
||
query = chain.run(
|
||
{"schema_info": schema_info, "prompt": text}
|
||
)
|
||
|
||
# print(query)
|
||
sql = getSql(query)
|
||
# result = execute_query(sql)
|
||
return sql
|
||
|
||
|
||
os.environ['DASHSCOPE_API_KEY'] = 'sk-c5f441f863f44094b0ddb96c831b5002'
|
||
# os.environ["QIANFAN_ACCESS_KEY"] = "zLiAbXCHVdJqQaZAREMpXxAc"
|
||
# os.environ["QIANFAN_SECRET_KEY"] = "wod9Pc2bZAyFVVOC9ypleJOAc4GoA4rT"
|
||
|
||
api = "zLiAbXCHVdJqQaZAREMpXxAc"
|
||
sk = "wod9Pc2bZAyFVVOC9ypleJOAc4GoA4rT"
|
||
|
||
llm = QianfanLLMEndpoint(model="ERNIE-4.0-8K", qianfan_ak=api, qianfan_sk=sk)
|
||
|
||
# 数据库连接参数
|
||
dbname = "gis_lca"
|
||
user = "postgres"
|
||
password = "Qibebt+123"
|
||
host = "localhost" # 或者是你 Docker 容器的 IP 地址，如果你在不同的机器上
|
||
port = "5432"
|
||
|
||
# 连接字符串
|
||
conn_string = f"host={host} dbname={dbname} user={user} password={password} port={port}"
|
||
# 连接到数据库
|
||
conn = psycopg2.connect(conn_string)
|
||
cur = conn.cursor()
|
||
|
||
# 获取数据库表结构，表名对应列名
|
||
schema = dict()
|
||
|
||
table_names = get_table_name(cur)
|
||
for name in table_names:
|
||
schema[name] = get_table_columns(cur, name)
|
||
|
||
|
||
# 读取数据
|
||
|
||
df = pd.read_excel(r"D:\python\LCA-GPT\LLM-SQL\Log\sql_ques_all.xlsx")
|
||
question = df['Question'].tolist()
|
||
|
||
sqllist = []
|
||
result = []
|
||
for ques in question:
|
||
prompt = ques
|
||
sql = query_database(prompt, schema)
|
||
sql = sql.lower()
|
||
sqllist.append(sql)
|
||
|
||
for sql in sqllist:
|
||
sql = sql.lower()
|
||
sql_data = pd.DataFrame(sqllist,columns=["SQL_pred"])
|
||
path = "./data/SQL_pred.xlsx"
|
||
sql_data.to_excel(path,index = False)
|
||
|
||
|