BBIT-Kai
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
|
||
from config.llm import llm
|
||
from langchain.prompts import PromptTemplate
|
||
from config.ssDb import ssDBLC
|
||
from langchain_community.agent_toolkits import create_sql_agent
|
||
|
||
#______________________________________________________________SQL描述_____________________________________________________________________
|
||
sqlDescriptionPrompt = PromptTemplate(
|
||
input_variables=["sql"],
|
||
template = """
|
||
你是一个SQL专家，精通SQLServer数据库。请把一下SQL查询语句用通俗易懂的中文进行总结。
|
||
SQL语句：{sql}
|
||
有以下要求：
|
||
1. 不要任何解释
|
||
2. 不能有标点符号
|
||
3. 不能有markdown语法
|
||
4. 要用业务语言描述，不能有专业语句例如SQL表名等
|
||
请生成你认为合适的标题，：
|
||
"""
|
||
)
|
||
sqlDescriptionChain = sqlDescriptionPrompt | llm
|
||
|
||
def get_sql_description_response( sql: str) -> str:
|
||
return sqlDescriptionChain.invoke({
|
||
"sql": sql
|
||
})
|
||
|
||
#______________________________________________________________第一次生成SQL_____________________________________________________________________
|
||
sqlPrompt = PromptTemplate(
|
||
input_variables=["userInput"],
|
||
template = """
|
||
你是一个SQL专家，精通SQLServer数据库。
|
||
请根据用户的需求，生成相应的SQL查询语句。
|
||
只需要返回SQL语句，不要任何解释。
|
||
用户需求：{userInput}
|
||
请生成SQL语句：
|
||
"""
|
||
)
|
||
sqlChain = sqlPrompt | llm
|
||
agent = create_sql_agent(
|
||
llm=llm,
|
||
db=ssDBLC,
|
||
agent_type="tool-calling",
|
||
verbose=True
|
||
)
|
||
# def get_chat_sql_response2( userInput: str) -> str:
|
||
# return sqlChain.invoke({
|
||
# "userInput": userInput
|
||
# })
|
||
def get_chat_sql_response( userInput: str) -> str:
|
||
return agent.invoke({"input": userInput})["output"]
|
||
|
||
#______________________________________________________________改进SQL_____________________________________________________________________
|
||
sqlImprovePrompt = PromptTemplate(
|
||
input_variables=["userInput", "sql"],
|
||
template = """
|
||
你是一个SQL专家，精通SQLServer数据库。
|
||
请根据用户的需求，改进已有的SQL查询语句。
|
||
只需要返回改进后的SQL语句，不要任何解释。
|
||
已有SQL：{sql}
|
||
用户需求：{userInput}
|
||
"""
|
||
)
|
||
sqlImproveChain = sqlImprovePrompt | llm
|
||
|
||
def get_chat_sql_improve_response( userInput: str) -> str:
|
||
return sqlImproveChain.invoke({
|
||
"userInput": userInput
|
||
})
|