Text-to-SQL将自然语言转换为数据库查询语句

有关Text-To-SQL方法，可以查阅我的另一篇文章，Text-to-SQL方法研究

直接与数据库对话-text2sql

Text2sql就是把文本转换为sql语言，这段时间公司有这方面的需求，调研了一下市面上text2sql的方法，比如阿里的Chat2DB,麻省理工开源的Vanna。试验了一下，最终还是决定自研，基于Vanna的思想，RAG+大模型。

使用开源的Vanna实现text2sql比较方便，Vanna可以直接连接数据库，但是当用户权限能访问多个数据库的时候，就比较麻烦了，而且Vanna向量化存储之后，新的question作对比时没有区分数据库。因此自己实现了一下text2sq，仍然采用Vanna的思想，提前训练DDL,Sqlques,和数据库document。

这里简单做一下记录，以供后续学习使用。

基本思路

1、数据库DDL语句，SQL-Question，Dcoument信息获取

2、基于用户提问question和数据库Document锁定要分析的数据库

3、模型训练：借助数据库的DDL语句、元数据（描述数据库自身数据的信息）、相关文档说明、参考样例SQL等，训练一个RAG“模型”。

这一模型结合了embedding技术和向量数据库，使得数据库的结构和内容能够被高效地索引和检索。

4、语义检索： 当用户输入自然语言描述的问题时，①会从向量库里面检索，迅速找出与问题相关的内容；②进行BM25算法文本召回，找到与问题 最相关的内容；③分别使用RRF算法和Re-ranking重排序算法，锁定最相关内容

语义匹配：使用算法（如BERT等）来理解查询和文档的语义相似性

文本召回匹配：BM25算法文本召回，找到与问题最相关的内容

rerank结果重排序：对搜索结果进行排序。

5、Prompt构建： 检索到的相关信息会被组装进Prompt中，形成一个结构化的查询描述。这一Prompt随后会被传递给LLM（大型语言模型）用于生成准确的SQL查询。

实现逻辑图

实现架构图：

具体实现方式如下所示：

1.数据库的选择 class DataBaseSearch(object): def __init__(self, _model): self.name = 'DataBaseSearch' self.model = _model self.instruction = "为这段内容生成表示以用于匹配文本描述：" self.SIZE = 1024 self.index = faiss.IndexFlatL2(self.SIZE) self.textdata = [] self.subdata = {} self.i2key = {} self.id2ddls = {} self.id2sqlques = {} self.id2docs = {} self.strtexts = {} # self.ddldata = [] # self.sqlques_data = [] # self.document_data = [] self.load_textdata() # 加载text数据 self.load_textdata_vec() # text数据向量化 def load_textdata(self): try: response = requests.post( url="xxx", verify=False) print(response.text) jsonobj = json.loads(response.text) textdatas = jsonobj["data"] for textdata in textdatas: # 提取每一个数据库内容 cid = textdata["dataSetID"] cddls = textdata["ddl"] csql_ques = textdata["exp"] cdocuments = textdata["Intro"] self.textdata.append((cid, cddls, csql_ques, cdocuments)) # 整合所有数据 except Exception as e: print(e) # print("load textdata ", self.textdata) def load_textdata_vec(self): num0 = 0 for recode in self.textdata: _id = recode[0] _ddls = recode[1] _sql_ques = recode[2] _documents = recode[3] # _strtexts = str(_ddls) + str(_sql_ques) + str(_documents) _strtexts = str(_sql_ques) + str(_documents) text_embeddings = self.model.encode([_strtexts], normalize_embeddings=True) self.index.add(text_embeddings) self.i2key[num0] = _id self.strtexts[_id] = _strtexts self.id2ddls[_id] = _ddls self.id2sqlques[_id] = _sql_ques self.id2docs[_id] = _documents num0 += 1 # print("init instruction vec", num0) def calculate_score(self, score, question, kws): pass def find_vec_database(self, question, k, theata): # print(question) q_embeddings = self.model.encode([self.instruction + question], normalize_embeddings=True) D, I = self.index.search(q_embeddings, k) result = [] for i in range(k): sim_i = I[0][i] uuid = self.i2key.get(sim_i, "none") sim_v = D[0][i] database_texts = self.strtexts.get(uuid, "none") # score = self.calculate_score(sim_v, question, database_texts) # wait implement score = int(sim_v*1000) if score < theata: doc = {} doc["score"] = score doc["dataSetID"] = uuid result.append(doc) # print(result) return result if __name__ == '__main__': modelpath = "E:\module\bge-large-zh-v1.5" model = SentenceTransformer(modelpath) vs = DataBaseSearch(model) result = vs.find_vec_database("查询济南市第三幼儿园所有小班班级？", 1, 2000) print(result) 2.sql_ques：sql问题训练 class SqlQuesSearch(object): def __init__(self, _model): self.name = "SqlQuesSearch" self.model = _model self.instruction = "为这段内容生成表示以用于匹配文本描述：" self.SIZE = 1024 self.index = faiss.IndexFlatL2(self.SIZE) self.sqlquedata = [] self.i2dbid = {} self.i2sqlid = {} self.id2sqlque = {} self.id2que = {} self.id2sql = {} self.dbid2sqlques = {} # # self.sqlques = {} # # self.i2key = {} # # self.id2sqlques = {} # # self.num2sqlque = {} # self.ddldata = [] # self.sqlques_data = [] # self.document_data = [] self.load_textdata() # 加载text数据 self.load_textdata_vec() # text数据向量化 def load_textdata(self): try: response = requests.post( url="xxx", verify=False) print(response.text) jsonobj = json.loads(response.text) textdatas = jsonobj["data"] datadatas = jsonobj["data"] for datadata in datadatas: # 提取每一个数据库sql-ques内容 dbid = datadata["dataSetID"] sql_ques = datadata["exp"] self.sqlquedata.append((dbid, sql_ques)) # 整合sql数据 except Exception as e: print(e) # print("load textdata ", self.sqlquedata) def load_textdata_vec(self): num0 = 0 for recode in self.sqlquedata: db_id = recode[0] sql_ques = recode[1] for sql_que in sql_ques: sql_id = sql_que["sql_id"] question = sql_que["question"] sql = sql_que["sql"] ddl_embeddings = self.model.encode([question], normalize_embeddings=True) self.index.add(ddl_embeddings) self.i2dbid[num0] = db_id self.i2sqlid[num0] = sql_id self.id2que[sql_id] = question self.id2sql[sql_id] = sql num0 += 1 print("init sql-que vec", num0) def calculate_score(sim_v, question, sql_ques): pass def find_vec_sqlque(self, question, k, theta, dataSetID, number): q_embeddings = self.model.encode([self.instruction + question], normalize_embeddings=True) D, I = self.index.search(q_embeddings, k) result = [] for i in range(k): sim_i = I[0][i] dbid = self.i2dbid.get(sim_i, "none") # 获取数据库id sqlid = self.i2sqlid.get(sim_i, "none") question = self.id2que.get(sqlid, "none") sql = self.id2sql.get(sqlid, "none") if dbid == dataSetID: sim_v = D[0][i] score = int(sim_v * 1000) if score < theta: doc = {} doc["score"] = score doc["question"] = question doc["sql"] = sql result.append(doc) if len(result) == number: break return result if __name__ == '__main__': modelpath = "E:\module\bge-large-zh-v1.5" model = SentenceTransformer(modelpath) vs = SqlQuesSearch(model) result = vs.find_vec_sqlque("查询7月18日所有的儿童观察记录？", 3, 2000, dataSetID=111) print(result) 3.数据库DDL训练 class DdlQuesSearch(object): def __init__(self, _model): self.name = "DdlQuesSearch" self.model = _model self.instruction = "为这段内容生成表示以用于匹配文本描述：" self.SIZE = 1024 self.index = faiss.IndexFlatL2(self.SIZE) self.ddldata = [] self.sqlques = {} self.i2dbid = {} self.i2ddlid = {} self.dbid2ddls = {} self.id2ddl = {} self.ddlid2dbid = {} # self.ddldata = [] # self.sqlques_data = [] # self.document_data = [] self.load_ddldata() # 加载text数据 self.load_ddl_vec() # text数据向量化 def load_ddldata(self): try: response = requests.post( url="xxx", verify=False) print(response.text) jsonobj = json.loads(response.text) for database in databases: db_id = database["dataSetID"] ddls = database["ddl"] self.ddldata.append((db_id, ddls)) # print(db_id) # for ddl in database["ddl"]: # ddl_id = ddl["ddl_id"] # ddl = ddl['ddl'] # # self.id2ddl[ddl_id] = ddl # self.dbid2ddls[db_id] = self.id2ddl except Exception as e: print(e) # print("load textdata ", self.ddldata) def load_ddl_vec(self): num0 = 0 for recode in self.ddldata: db_id = recode[0] ddls = recode[1] for ddl in ddls: ddl_id = ddl["ddl_id"] ddl_name = ddl["TABLE"] ddl = ddl['ddl'] ddl_embeddings = self.model.encode([ddl], normalize_embeddings=True) self.index.add(ddl_embeddings) self.i2dbid[num0] = db_id self.i2ddlid[num0] = ddl_id self.id2ddl[ddl_id] = ddl self.ddlid2dbid[ddl_id] = db_id num0 += 1 self.dbid2ddls[db_id] = self.id2ddl print("init ddl vec", num0) def find_vec_ddl(self, question, k, theata, dataSetID, number): # dataSetID:数据库id # self.id2ddls.get(action_id) q_embeddings = self.model.encode([self.instruction + question], normalize_embeddings=True) D, I = self.index.search(q_embeddings, k) result = [] for i in range(k): sim_i = I[0][i] dbid = self.i2dbid.get(sim_i, "none") # 获取数据库id ddlid = self.i2ddlid.get(sim_i, "none") if dbid == dataSetID: sim_v = D[0][i] score = int(sim_v * 1000) if score < theata: doc = {} doc["score"] = score doc["ddl"] = self.id2ddl.get(ddlid, "none") result.append(doc) if len(result) == number: break return result if __name__ == '__main__': modelpath = "E:\module\bge-large-zh-v1.5" model = SentenceTransformer(modelpath) vs = DdlQuesSearch(model) ss = vs.find_vec_ddl("定时任务执行记录表", 2, 2000, 111) print(ss) 4.数据库document训练 class DocQuesSearch(object): def __init__(self): self.name = "TestDataSearch" self.docdata = [] self.load_doc_data() def load_doc_data(self): try: response = requests.post( url="xxx", verify=False) print(response.text) jsonobj = json.loads(response.text) databases = jsonobj["data"] for database in databases: db_id = database["dataSetID"] doc = database["Intro"] self.docdata.append((db_id, doc)) except Exception as e: print(e) # print("load ddldata ", self.docdata) def find_similar_doc(self, dataSetID): result = [] for recode in self.docdata: dbid = recode[0] doc = recode[1] if dbid == dataSetID: result.append(doc) return result if __name__ == '__main__': docques_search = DocQuesSearch() result = docques_search.find_similar_doc(222) print(result) 5.生成sql语句，这里使用的qwen-max模型 import re import random import os, json import dashscope from dashscope.api_entities.dashscope_response import Message from ddl_engine import DdlQuesSearch from dashscope import Generation from sqlques_engine import SqlQuesSearch from sentence_transformers import SentenceTransformer class Genarate(object): def __init__(self): self.api_key = os.environ.get('api_key') self.model_name = os.environ.get('model') def system_message(self, message): return {'role': 'system', 'content': message} def user_message(self, message): return {'role': 'user', 'content': message} def assistant_message(self, message): return {'role': 'assistant', 'content': message} def submit_prompt(self, prompt): resp = Generation.call( model=self.model_name, messages=prompt, seed=random.randint(1, 10000), result_format='message', api_key=self.api_key) if resp["status_code"] == 200: answer = resp.output.choices[0].message.content global DEBUG_INFO DEBUG_INFO = (prompt, answer) return answer else: answer = None return answer def generate_sql(self, question, sql_result, ddl_result, doc_result): prompt = self.get_sql_prompt( question = question, sql_result = sql_result, ddl_result = ddl_result, doc_result = doc_result) print("SQL Prompt:",prompt) llm_response = self.submit_prompt(prompt) sql = self.extrat_sql(llm_response) return sql def extrat_sql(self, llm_response): sqls = re.findall(r"WITH.*?;", llm_response, re.DOTALL) if sqls: sql = sqls[-1] return sql sqls = re.findall(r"SELECT.*?;", llm_response, re.DOTALL) if sqls: sql = sqls[-1] return sql sqls = re.findall(r"```sql (.*)```", llm_response, re.DOTALL) if sqls: sql = sqls[-1] return sql sqls = re.findall(r"```(.*)```", llm_response, re.DOTALL) if sqls: sql = sqls[-1] return sql return llm_response def get_sql_prompt(self, question, sql_result, ddl_result, doc_result): initial_prompt = "You are a SQL expert. " + "Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. " initial_prompt = self.add_ddl_to_prompt( initial_prompt, ddl_result) initial_prompt = self.add_documentation_to_prompt(initial_prompt, doc_result) initial_prompt += ( "===Response Guidelines " "1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. " "2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql " "3. If the provided context is insufficient, please explain why it can't be generated. " "4. Please use the most relevant table(s). " "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. " ) message_log = [self.system_message(initial_prompt)] message_log = self.add_sqlques_to_prompt(question, sql_result, message_log) return message_log def add_ddl_to_prompt(self, initial_prompt, ddl_result): """ :param initial_prompt: :param ddl_result: :return: """ ddl_list = [ ddl_['ddl'] for ddl_ in ddl_result] if len(ddl_list) > 0: initial_prompt += " ===Tables " for ddl in ddl_list: initial_prompt += f"{ddl} " return initial_prompt def add_sqlques_to_prompt(self, question, sql_result, message_log): """ :param sql_result: :return: """ if len(sql_result) > 0: for example in sql_result: if example is not None and "question" in example and "sql" in example: message_log.append(self.user_message(example["question"])) message_log.append(self.assistant_message(example["sql"])) message_log.append(self.user_message(question)) return message_log def add_documentation_to_prompt(self, initial_prompt, doc_result): if len(doc_result) > 0: initial_prompt += " ===Additional Context " for doc in doc_result: initial_prompt += f"{doc} " return initial_prompt if __name__ == '__main__': modelpath = "E:\module\bge-large-zh-v1.5" model = SentenceTransformer(modelpath) vs = DdlQuesSearch(model) ss = vs.find_vec_ddl("定时任务执行记录表", 1, 2000, 111) print(ss) 6.执行结果显示

如图可以看到正确生成了sql,可以正常执行，因为表是拉取到，没有数据，所以查询结果为空。

需要源码的同学，可以留言。