klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-09 09:46:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: improve generate sql quality

Browse Source
This commit is contained in:
ulleo
2026-04-07 19:06:17 +08:00
committed by ulleo
parent f75770f945
commit 594df91235
4 changed files with 164 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/models/chat_model.py
+56 -12
View File
@@ -1,8 +1,9 @@
from datetime import datetime
from enum import Enum
from typing import List, Optional, Union
from typing import List, Optional, Any, Union
from fastapi import Body
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from pydantic import BaseModel
from sqlalchemy import Column, Integer, Text, BigInteger, DateTime, Identity, Boolean
from sqlalchemy import Enum as SQLAlchemyEnum
@@ -230,6 +231,7 @@ class AiModelQuestion(BaseModel):
regenerate_record_id: Optional[int] = None
def sql_sys_question(self, db_type: Union[str, DB], enable_query_limit: bool = True):
templates: dict[str, str] = {}
_sql_template = get_sql_example_template(db_type)
_base_template = get_sql_template()
_process_check = _sql_template.get('process_check') if _sql_template.get('process_check') else _base_template[
@@ -245,22 +247,37 @@ class AiModelQuestion(BaseModel):
'example_answer_2']
_example_answer_3 = _sql_template['example_answer_3_with_limit'] if enable_query_limit else _sql_template[
'example_answer_3']
return _base_template['system'].format(engine=self.engine, schema=self.db_schema, question=self.question,
lang=self.lang, terminologies=self.terminologies,
data_training=self.data_training, custom_prompt=self.custom_prompt,
process_check=_process_check,
base_sql_rules=_base_sql_rules,
basic_sql_examples=_sql_examples,
example_engine=_example_engine,
example_answer_1=_example_answer_1,
example_answer_2=_example_answer_2,
example_answer_3=_example_answer_3)
templates['system'] = _base_template['system'].format(process_check=_process_check)
templates['rules'] = _base_template['generate_rules'].format(lang=self.lang,
base_sql_rules=_base_sql_rules,
basic_sql_examples=_sql_examples,
example_engine=_example_engine,
example_answer_1=_example_answer_1,
example_answer_2=_example_answer_2,
example_answer_3=_example_answer_3)
templates['schema'] = _base_template['generate_basic_info'].format(engine=self.engine, schema=self.db_schema)
if self.terminologies:
templates['terminologies'] = _base_template['generate_terminologies_info'].format(
terminologies=self.terminologies)
if self.data_training:
templates['data_training'] = _base_template['generate_data_training_info'].format(
data_training=self.data_training)
if self.custom_prompt:
templates['custom_prompt'] = _base_template['generate_custom_prompt_info'].format(
custom_prompt=self.custom_prompt)
return templates
def sql_user_question(self, current_time: str, change_title: bool):
_question = self.question
if self.regenerate_record_id:
_question = get_sql_template()['regenerate_hint'] + self.question
return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=_question,
return get_sql_template()['user'].format(lang=self.lang, engine=self.engine, schema=self.db_schema,
question=_question,
rule=self.rule, current_time=current_time, error_msg=self.error_msg,
change_title=change_title)
@@ -358,3 +375,30 @@ class McpAssistant(BaseModel):
url: str = Body(description='第三方数据接口')
authorization: str = Body(description='第三方接口凭证')
stream: Optional[bool] = Body(description='是否流式输出，默认为true开启, 关闭false则返回JSON对象', default=True)
class SystemPromptMessage(SystemMessage):
sqlbot_system: bool = True
def __init__(
self, content: Union[str, list[Union[str, dict]]], **kwargs: Any
) -> None:
super().__init__(content=content, **kwargs)
class HumanPromptMessage(HumanMessage):
sqlbot_system: bool = True
def __init__(
self, content: Union[str, list[Union[str, dict]]], **kwargs: Any
) -> None:
super().__init__(content=content, **kwargs)
class AIPromptMessage(AIMessage):
sqlbot_system: bool = True
def __init__(
self, content: Union[str, list[Union[str, dict]]], **kwargs: Any
) -> None:
super().__init__(content=content, **kwargs)
backend/apps/chat/task/llm.py
+79 -15
View File
@@ -14,7 +14,7 @@ import requests
import sqlparse
from langchain.chat_models.base import BaseChatModel
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, BaseMessageChunk
from sqlalchemy import and_, select
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlbot_xpack.config.model import SysArgModel
@@ -33,7 +33,7 @@ from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate, get_chat_predict_data, \
get_chat_chart_config, trigger_log_error
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \
ChatFinishStep, AxisObj
ChatFinishStep, AxisObj, SystemPromptMessage, HumanPromptMessage, AIPromptMessage
from apps.data_training.curd.data_training import get_training_template
from apps.datasource.crud.datasource import get_table_schema
from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user
@@ -216,12 +216,31 @@ class LLMService:
filter(lambda obj: obj.pid == self.chat_question.regenerate_record_id, self.generate_sql_logs), None)
last_sql_messages: List[dict[str, Any]] = _temp_log.messages if _temp_log else []
# 排除所有的系统提示词
last_sql_messages = [obj for obj in last_sql_messages if obj.get("sqlbot_system") != True]
count_limit = self.base_message_round_count_limit
self.sql_message = []
# add sys prompt
self.sql_message.append(SystemMessage(
content=self.chat_question.sql_sys_question(self.ds.type, self.enable_sql_row_limit)))
_system_templates = self.chat_question.sql_sys_question(self.ds.type, self.enable_sql_row_limit)
self.sql_message.append(SystemPromptMessage(content=_system_templates['system']))
self.sql_message.append(HumanPromptMessage(content=_system_templates['rules']))
self.sql_message.append(
AIPromptMessage(content='我已掌握所有规则，包括表结构、SQL规范、安全限制和输出格式，我会严格遵守这些规则。'))
self.sql_message.append(HumanPromptMessage(content=_system_templates['schema']))
self.sql_message.append(
AIPromptMessage(content='我已确认您提供的数据库信息与表结构schema，我生成的SQL不会超出您提供的范围。'))
if _system_templates.get('custom_prompt'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['custom_prompt']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的额外信息，我会进行参考。'))
if _system_templates.get('terminologies'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['terminologies']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的术语信息，我会进行参考。'))
if _system_templates.get('data_training'):
self.sql_message.append(HumanPromptMessage(content=_system_templates['data_training']))
self.sql_message.append(AIPromptMessage(content='我已确认您提供的SQL示例，我会进行参考。'))
if last_sql_messages is not None and len(last_sql_messages) > 0:
last_rounds = get_last_conversation_rounds(last_sql_messages, rounds=count_limit)
@@ -246,7 +265,7 @@ class LLMService:
self.chart_message = []
# add sys prompt
self.chart_message.append(SystemMessage(content=self.chat_question.chart_sys_question()))
self.chart_message.append(SystemPromptMessage(content=self.chat_question.chart_sys_question()))
if last_chart_messages is not None and len(last_chart_messages) > 0:
last_rounds = get_last_conversation_rounds(last_chart_messages, rounds=count_chart_limit)
@@ -369,7 +388,7 @@ class LLMService:
self.filter_custom_prompts(_session, CustomPromptTypeEnum.ANALYSIS, self.current_user.oid, ds_id)
analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question()))
analysis_msg.append(SystemPromptMessage(content=self.chat_question.analysis_sys_question()))
analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question()))
self.current_logs[OperationEnum.ANALYSIS] = start_log(session=_session,
@@ -379,6 +398,8 @@ class LLMService:
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in analysis_msg])
@@ -400,6 +421,8 @@ class LLMService:
OperationEnum.ANALYSIS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in analysis_msg],
reasoning_content=full_thinking_text,
@@ -417,7 +440,7 @@ class LLMService:
self.filter_custom_prompts(_session, CustomPromptTypeEnum.PREDICT_DATA, self.current_user.oid, ds_id)
predict_msg: List[Union[BaseMessage, dict[str, Any]]] = []
predict_msg.append(SystemMessage(content=self.chat_question.predict_sys_question()))
predict_msg.append(SystemPromptMessage(content=self.chat_question.predict_sys_question()))
predict_msg.append(HumanMessage(content=self.chat_question.predict_user_question()))
self.current_logs[OperationEnum.PREDICT_DATA] = start_log(session=_session,
@@ -427,6 +450,8 @@ class LLMService:
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in predict_msg])
@@ -449,6 +474,8 @@ class LLMService:
OperationEnum.PREDICT_DATA],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in predict_msg],
reasoning_content=full_thinking_text,
@@ -466,7 +493,7 @@ class LLMService:
embedding=False)
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question(self.articles_number)))
guess_msg.append(SystemPromptMessage(content=self.chat_question.guess_sys_question(self.articles_number)))
old_questions = list(map(lambda q: q.strip(), get_old_questions(_session, self.record.datasource)))
guess_msg.append(
@@ -479,6 +506,9 @@ class LLMService:
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in guess_msg])
@@ -500,6 +530,9 @@ class LLMService:
OperationEnum.GENERATE_RECOMMENDED_QUESTIONS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in guess_msg],
reasoning_content=full_thinking_text,
@@ -512,7 +545,7 @@ class LLMService:
def select_datasource(self, _session: Session):
datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
datasource_msg.append(SystemPromptMessage(self.chat_question.datasource_sys_question()))
if self.current_assistant and self.current_assistant.type != 4:
_ds_list = get_assistant_ds(session=_session, llm_service=self)
else:
@@ -552,6 +585,9 @@ class LLMService:
operate=OperationEnum.CHOOSE_DATASOURCE,
record_id=self.record.id,
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for
msg in datasource_msg])
@@ -571,6 +607,9 @@ class LLMService:
OperationEnum.CHOOSE_DATASOURCE],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in datasource_msg],
reasoning_content=full_thinking_text,
@@ -661,7 +700,10 @@ class LLMService:
operate=OperationEnum.GENERATE_SQL,
record_id=self.record.id,
full_message=[
{'type': msg.type, 'content': msg.content} for msg
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for msg
in self.sql_message])
full_thinking_text = ''
full_sql_text = ''
@@ -678,7 +720,11 @@ class LLMService:
self.current_logs[OperationEnum.GENERATE_SQL] = end_log(session=_session,
log=self.current_logs[OperationEnum.GENERATE_SQL],
full_message=[{'type': msg.type, 'content': msg.content}
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in self.sql_message],
reasoning_content=full_thinking_text,
token_usage=token_usage)
@@ -690,7 +736,7 @@ class LLMService:
self.chat_question.sql = sql
self.chat_question.sub_query = sub_query
dynamic_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = []
dynamic_sql_msg.append(SystemMessage(content=self.chat_question.dynamic_sys_question()))
dynamic_sql_msg.append(SystemPromptMessage(content=self.chat_question.dynamic_sys_question()))
dynamic_sql_msg.append(HumanMessage(content=self.chat_question.dynamic_user_question()))
self.current_logs[OperationEnum.GENERATE_DYNAMIC_SQL] = start_log(session=session,
@@ -699,6 +745,9 @@ class LLMService:
operate=OperationEnum.GENERATE_DYNAMIC_SQL,
record_id=self.record.id,
full_message=[{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for
msg in dynamic_sql_msg])
@@ -720,6 +769,9 @@ class LLMService:
OperationEnum.GENERATE_DYNAMIC_SQL],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in dynamic_sql_msg],
reasoning_content=full_thinking_text,
@@ -748,7 +800,7 @@ class LLMService:
self.chat_question.sql = sql
self.chat_question.filter = filter
permission_sql_msg: List[Union[BaseMessage, dict[str, Any]]] = []
permission_sql_msg.append(SystemMessage(content=self.chat_question.filter_sys_question()))
permission_sql_msg.append(SystemPromptMessage(content=self.chat_question.filter_sys_question()))
permission_sql_msg.append(HumanMessage(content=self.chat_question.filter_user_question()))
self.current_logs[OperationEnum.GENERATE_SQL_WITH_PERMISSIONS] = start_log(session=session,
@@ -758,6 +810,9 @@ class LLMService:
record_id=self.record.id,
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in permission_sql_msg])
@@ -778,6 +833,9 @@ class LLMService:
OperationEnum.GENERATE_SQL_WITH_PERMISSIONS],
full_message=[
{'type': msg.type,
'sqlbot_system': getattr(msg,
'sqlbot_system',
False) is True,
'content': msg.content}
for msg in permission_sql_msg],
reasoning_content=full_thinking_text,
@@ -813,7 +871,10 @@ class LLMService:
operate=OperationEnum.GENERATE_CHART,
record_id=self.record.id,
full_message=[
{'type': msg.type, 'content': msg.content} for
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content} for
msg
in self.chart_message])
full_thinking_text = ''
@@ -834,7 +895,10 @@ class LLMService:
self.current_logs[OperationEnum.GENERATE_CHART] = end_log(session=_session,
log=self.current_logs[OperationEnum.GENERATE_CHART],
full_message=[
{'type': msg.type, 'content': msg.content}
{'type': msg.type,
'sqlbot_system': getattr(msg, 'sqlbot_system',
False) is True,
'content': msg.content}
for msg in self.chart_message],
reasoning_content=full_thinking_text,
token_usage=token_usage)
backend/templates/template.yaml
+25 -15
View File
@@ -71,18 +71,21 @@ template:
<Info>内有<db-engine><m-schema><terminologies>等信息；
其中，<db-engine>：提供数据库引擎及版本信息；
<m-schema>：以 M-Schema 格式提供数据库表结构信息；
<terminologies>：提供一组术语，块内每一个<terminology>就是术语，其中同一个<words>内的多个<word>代表术语的多种叫法，也就是术语与它的同义词，<description>即该术语对应的描述，其中也可能是能够用来参考的计算公式，或者是一些其他的查询条件；
<sql-examples>：提供一组SQL示例，你可以参考这些示例来生成你的回答，其中<question>内是提问，<suggestion-answer>内是对于该<question>提问的解释或者对应应该回答的SQL示例。
<terminologies>：提供一组术语，块内每一个<terminology>就是术语，其中同一个<words>内的多个<word>代表术语的多种叫法，也就是术语与它的同义词，<description>即该术语对应的描述，其中也可能是能够用来参考的计算公式，或者是一些其他的查询条件；
<sql-examples>：提供一组SQL示例，你可以参考这些示例来生成你的回答，其中<question>内是提问，<suggestion-answer>内是对于该<question>提问的解释或者对应应该回答的SQL示例。
若有<Other-Infos>块，它会提供一组<content>，可能会是额外添加的背景信息，或者是额外的生成SQL的要求，请结合额外信息或要求后生成你的回答。
用户的提问在<user-question>内，<error-msg>内则会提供上次执行你提供的SQL时会出现的错误信息，<background-infos>内的<current-time>会告诉你用户当前提问的时间
你必须遵守<Rules>内规定的生成SQL规则
你必须遵守<SQL-Generation-Process>内规定的检查步骤生成你的回答
用户的提问在<user-question>内，<error-msg>内则会提供上次执行你提供的SQL时会出现的错误信息，<background-infos>内的<current-time>会告诉你用户当前提问的时间
</Instruction>
{process_check}
generate_rules: |
以下是你必须遵守的规则和可以参考的基础示例：
<Rules>
<rule>
请使用语言：{lang} 回答，若有深度思考过程，则思考过程也需要使用 {lang} 输出
<note>即使示例中显示了中文消息，在实际生成时也必须翻译成英文</note>
</rule>
<rule>
你只能生成查询用的SQL语句，不得生成增删改相关或操作数据库以及操作数据库数据的SQL
@@ -186,8 +189,6 @@ template:
</rule>
</Rules>
{process_check}
{basic_sql_examples}
<example>
@@ -298,23 +299,32 @@ template:
</example>
</chat-examples>
</example>
以下是正式的信息：
generate_terminologies_info: |
以下是你可以参考的术语：
{terminologies}
generate_data_training_info: |
以下是你可以参考的SQL示例：
{data_training}
generate_custom_prompt_info: |
以下是你可以参考的额外信息：
{custom_prompt}
generate_basic_info: |
以下是数据库与表结构信息，你生成的SQL使用到的表名与字段必须在提供的范围内
<Info>
<db-engine> {engine} </db-engine>
<m-schema>
{schema}
</m-schema>
{terminologies}
{data_training}
</Info>
{custom_prompt}
### 响应, 请根据上述要求直接返回JSON结果:
```json
user: |
请根据上述要求，用语言：{lang} 进行回答
如果<user-question>内的提问与上述要求冲突，你需要停止生成SQL并告知生成SQL失败的原因
请输出符合要求的JSON回答
<background-infos>
<current-time>
{current_time}
frontend/src/entity/supplier.ts
+4
View File
@@ -46,6 +46,10 @@ export const supplierList: Array<{
{ key: 'extra_body', val: '{"enable_thinking": false}', type: 'json' },
],
model_options: [
{ name: 'qwen3.6-plus' },
{ name: 'qwen3.5-plus' },
{ name: 'qwen3.5-flash' },
{ name: 'qwen3-coder-next' },
{ name: 'qwen3-coder-plus' },
{ name: 'qwen3-coder-flash' },
{ name: 'qwen-plus' },