|
|
from util import use_pgvector, use_opanai, use_mysql
|
|
|
from util.use_mysql import search_desc_by_table_names
|
|
|
|
|
|
|
|
|
async def rag_generate_rule(query : str):
|
|
|
if not query:
|
|
|
return "请输入问题"
|
|
|
# 连接数据库
|
|
|
pgvector_conn = use_pgvector.connect_to_db()
|
|
|
|
|
|
# 将问题向量化
|
|
|
query_emb = await use_opanai.generation_vector("query")
|
|
|
|
|
|
# 根据问题关联数据库表，得到table_name_list
|
|
|
similar_docs = use_pgvector.search_similar_table(pgvector_conn, query_emb, limit=4)
|
|
|
rerank_docs = await use_opanai.rerank_documents(query, similar_docs, top_n=2)
|
|
|
|
|
|
table_name_list = [similar_docs[int(index.strip())-1][0] for index in rerank_docs.strip('[]').split(',')]
|
|
|
print(f"[table_name_list]: {table_name_list}")
|
|
|
|
|
|
# 获得相关表的schema
|
|
|
db_client = await use_mysql.get_db()
|
|
|
schema = await search_desc_by_table_names(table_name_list, db_client)
|
|
|
print(f"[schema]: {schema}")
|
|
|
|
|
|
# 根据问题搜索相关案例
|
|
|
similar_case = use_pgvector.search_similar_case(pgvector_conn, query_emb, limit=3)
|
|
|
print(f"[similar_case]: {similar_case}")
|
|
|
|
|
|
# 询问大模型生成SQL
|
|
|
prompt = [
|
|
|
{"role": "system", "content": f"""
|
|
|
你是精通postgresql的专家，十分擅长利用用户的问题，和库表schema，生成合适的SQL。最终结果里只需要返回SQL，不要解释。
|
|
|
|
|
|
请根据用户问题，生成SQL，请勿返回其他内容。
|
|
|
|
|
|
schema: {schema}
|
|
|
|
|
|
参考案例: {similar_case}
|
|
|
"""},
|
|
|
{"role": "user", "content": query}
|
|
|
]
|
|
|
print(f"[prompt]: {prompt}")
|
|
|
ans = await use_opanai.generation_rule(prompt)
|
|
|
|
|
|
return ans
|
|
|
|
|
|
async def main():
|
|
|
ans = await rag_generate_rule("“投资规模结构分解批次（备份）”表BAK_PC_IP_SCALE_STRUCT_RESOLVE中“业务板块类型”不合规的记录")
|
|
|
print(f"[answer]: {ans}")
|
|
|
|
|
|
# 使用asyncio运行异步函数
|
|
|
import asyncio
|
|
|
if __name__ == "__main__":
|
|
|
asyncio.run(main())
|