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

feat(X-Pack): add custom prompt

Browse Source 
#213
This commit is contained in:
ulleo
2025-09-28 19:19:48 +08:00
committed by ulleo
parent 1b25652ba4
commit 75bd13a54f
10 changed files with 123 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/alembic/env.py
+1
View File
@@ -26,6 +26,7 @@ config = context.config
# from apps.settings.models.setting_models import SQLModel
# from apps.chat.models.chat_model import SQLModel
from apps.terminology.models.terminology_model import SQLModel
#from apps.custom_prompt.models.custom_prompt_model import SQLModel
# from apps.data_training.models.data_training_model import SQLModel
# from apps.dashboard.models.dashboard_model import SQLModel
from common.core.config import settings # noqa
backend/alembic/versions/046_add_custom_prompt.py
+39
View File
@@ -0,0 +1,39 @@
"""046_add_custom_prompt
Revision ID: 8855aea2dd61
Revises: 45e7e52bf2b8
Create Date: 2025-09-28 13:57:01.509249
"""
from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision = '8855aea2dd61'
down_revision = '45e7e52bf2b8'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('custom_prompt',
sa.Column('id', sa.BigInteger(), sa.Identity(always=True), nullable=False),
sa.Column('oid', sa.BigInteger(), nullable=True),
sa.Column('type', sa.Enum('GENERATE_SQL', 'ANALYSIS', 'PREDICT_DATA', name='customprompttypeenum', native_enum=False, length=20), nullable=True),
sa.Column('create_time', sa.DateTime(), nullable=True),
sa.Column('name', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=True),
sa.Column('prompt', sa.Text(), nullable=True),
sa.Column('specific_ds', sa.Boolean(), nullable=True),
sa.Column('datasource_ids', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table('custom_prompt')
# ### end Alembic commands ###
backend/apps/chat/models/chat_model.py
+7 -3
View File
@@ -40,11 +40,13 @@ class OperationEnum(Enum):
CHOOSE_DATASOURCE = '6'
GENERATE_DYNAMIC_SQL = '7'
class ChatFinishStep(Enum):
GENERATE_SQL = 1
QUERY_DATA = 2
GENERATE_CHART = 3
# TODO choose table / check connection / generate description
class ChatLog(SQLModel, table=True):
@@ -177,12 +179,13 @@ class AiModelQuestion(BaseModel):
sub_query: Optional[list[dict]] = None
terminologies: str = ""
data_training: str = ""
custom_prompt: str = ""
error_msg: str = ""
def sql_sys_question(self):
return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question,
lang=self.lang, terminologies=self.terminologies,
data_training=self.data_training)
data_training=self.data_training, custom_prompt=self.custom_prompt)
def sql_user_question(self, current_time: str):
return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question,
@@ -196,13 +199,14 @@ class AiModelQuestion(BaseModel):
chart_type=chart_type)
def analysis_sys_question(self):
return get_analysis_template()['system'].format(lang=self.lang, terminologies=self.terminologies)
return get_analysis_template()['system'].format(lang=self.lang, terminologies=self.terminologies,
custom_prompt=self.custom_prompt)
def analysis_user_question(self):
return get_analysis_template()['user'].format(fields=self.fields, data=self.data)
def predict_sys_question(self):
return get_predict_template()['system'].format(lang=self.lang)
return get_predict_template()['system'].format(lang=self.lang, custom_prompt=self.custom_prompt)
def predict_user_question(self):
return get_predict_template()['user'].format(fields=self.fields, data=self.data)
backend/apps/chat/task/llm.py
+18
View File
@@ -30,6 +30,9 @@ from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
get_last_execute_sql_error
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \
ChatFinishStep
from sqlbot_xpack.license.license_manage import SQLBotLicenseUtil
from sqlbot_xpack.custom_prompt.curd.custom_prompt import find_custom_prompts
from sqlbot_xpack.custom_prompt.models.custom_prompt_model import CustomPromptTypeEnum
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
@@ -244,6 +247,9 @@ class LLMService:
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question,
self.current_user.oid, ds_id)
if SQLBotLicenseUtil.valid():
self.chat_question.custom_prompt = find_custom_prompts(self.session, CustomPromptTypeEnum.ANALYSIS,
self.current_user.oid, ds_id)
analysis_msg.append(SystemMessage(content=self.chat_question.analysis_sys_question()))
analysis_msg.append(HumanMessage(content=self.chat_question.analysis_user_question()))
@@ -288,6 +294,12 @@ class LLMService:
self.chat_question.fields = orjson.dumps(fields).decode()
data = get_chat_chart_data(self.session, self.record.id)
self.chat_question.data = orjson.dumps(data.get('data')).decode()
if SQLBotLicenseUtil.valid():
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
self.chat_question.custom_prompt = find_custom_prompts(self.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(HumanMessage(content=self.chat_question.predict_user_question()))
@@ -509,6 +521,9 @@ class LLMService:
ds_id)
self.chat_question.data_training = get_training_template(self.session, self.chat_question.question, ds_id,
oid)
if SQLBotLicenseUtil.valid():
self.chat_question.custom_prompt = find_custom_prompts(self.session, CustomPromptTypeEnum.GENERATE_SQL,
oid, ds_id)
self.init_messages()
@@ -902,6 +917,9 @@ class LLMService:
oid, ds_id)
self.chat_question.data_training = get_training_template(self.session, self.chat_question.question,
ds_id, oid)
if SQLBotLicenseUtil.valid():
self.chat_question.custom_prompt = find_custom_prompts(self.session, CustomPromptTypeEnum.GENERATE_SQL,
oid, ds_id)
self.init_messages()
backend/apps/terminology/curd/terminology.py
+1 -3
View File
@@ -5,9 +5,7 @@ from typing import List, Optional, Any
from xml.dom.minidom import parseString
import dicttoxml
from sqlalchemy import BigInteger
from sqlalchemy import and_, or_, select, func, delete, update, union
from sqlalchemy import text
from sqlalchemy import and_, or_, select, func, delete, update, union, text, BigInteger
from sqlalchemy.orm import aliased
from sqlalchemy.orm.session import Session
backend/locales/en.json
+4
View File
@@ -47,6 +47,10 @@
"datasource_cannot_be_none": "Datasource cannot be none",
"data_training_not_exists": "Example does not exists",
"exists_in_db": "Question exists"
},
"i18n_custom_prompt": {
"exists_in_db": "Prompt name exists",
"not_exists": "Prompt does not exists"
},
"i18n_excel_export": {
"data_is_empty": "The form data is empty, cannot export data"
backend/locales/zh-CN.json
+4
View File
@@ -48,6 +48,10 @@
"data_training_not_exists": "该示例不存在",
"exists_in_db": "该问题已存在"
},
"i18n_custom_prompt": {
"exists_in_db": "模版名称已存在",
"not_exists": "该模版不存在"
},
"i18n_excel_export": {
"data_is_empty": "表单数据为空，无法导出数据"
}
backend/pyproject.toml
+1 -1
View File
@@ -39,7 +39,7 @@ dependencies = [
"pyyaml (>=6.0.2,<7.0.0)",
"fastapi-mcp (>=0.3.4,<0.4.0)",
"tabulate>=0.9.0",
"sqlbot-xpack>=0.0.3.31,<1.0.0",
"sqlbot-xpack>=0.0.3.36,<1.0.0",
"fastapi-cache2>=0.2.2",
"sqlparse>=0.5.3",
"redis>=6.2.0",
backend/template.yaml
+47 -14
View File
@@ -14,8 +14,9 @@ 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>会告诉你用户当前提问的时间
</Instruction>
@@ -219,7 +220,6 @@ template:
</chat-examples>
</example>
### 下面是提供的信息
<Info>
<db-engine> {engine} </db-engine>
<m-schema>
@@ -229,6 +229,7 @@ template:
{terminologies}
{data_training}
</Info>
{custom_prompt}
### 响应, 请根据上述要求直接返回JSON结果:
```json
@@ -394,7 +395,11 @@ template:
你当前的任务是根据给定的数据分析数据，并给出你的分析结果。
我们会在<Info>块内提供给你信息，帮助你进行分析：
<Info>内有<terminologies>等信息；
<terminologies>：提供一组术语，块内每一个<terminology>就是术语，其中同一个<words>内的多个<word>代表术语的多种叫法，也就是术语与它的同义词，<description>即该术语对应的描述，其中也可能是能够用来参考的计算公式，或者是一些其他的查询条件
<terminologies>：提供一组术语，块内每一个<terminology>就是术语，其中同一个<words>内的多个<word>代表术语的多种叫法，也就是术语与它的同义词，<description>即该术语对应的描述，其中也可能是能够用来参考的计算公式，或者是一些其他的查询条件。
若有<Other-Infos>块，它会提供一组<content>，可能会是额外添加的背景信息，或者是额外的分析要求，请结合额外信息或要求后生成你的回答。
用户会在提问中提供给你信息：
<data>块内是提供给你的数，以JSON格式给出；
<fields>块内提供给你对应的字段或字段别名。
</Instruction>
你必须遵守以下规则:
@@ -404,32 +409,60 @@ template:
</rule>
</Rules>
### 下面是提供的信息
<Info>
{terminologies}
</Info>
{custom_prompt}
user: |
### 字段(字段别名):
<fields>
{fields}
</fields>
### 数据:
<data>
{data}
</data>
predict:
system: |
### 请使用语言：{lang} 回答，若有深度思考过程，则思考过程也需要使用 {lang} 输出
<Instruction>
你是"SQLBOT"，智能问数小助手，可以根据用户提问，专业生成SQL与可视化图表。
你当前的任务是根据给定的数据进行数据预测，并给出你的预测结果。
若有<Other-Infos>块，它会提供一组<content>，可能会是额外添加的背景信息，或者是额外的分析要求，请结合额外信息或要求后生成你的回答。
用户会在提问中提供给你信息：
<data>块内是提供给你的数据，以JSON格式给出；
<fields>块内提供给你对应的字段或字段别名。
</Instruction>
### 说明：
你是一个数据分析师，你的任务是根据给定的数据进行数据预测，我将以JSON格式给你一组数据，你帮我预测之后的数据（一段可以展示趋势的数据，至少2个周期），用json数组的格式返回，返回的格式需要与传入的数据格式保持一致。
你必须遵守以下规则:
<Rules>
<rule>
请使用语言：{lang} 回答，若有深度思考过程，则思考过程也需要使用 {lang} 输出
</rule>
<rule>
预测的数据是一段可以展示趋势的数据，至少2个周期
</rule>
<rule>
返回的预测数据必须与用户提供的数据同样的格式，使用JSON数组的形式返回
</rule>
<rule>
无法预测或者不支持预测的数据请直接返回(不需要返回JSON格式)："抱歉，该数据无法进行预测。"(若有原因，则额外返回无法预测的原因)
</rule>
<rule>
预测的数据不需要返回用户提供的原有数据，请直接返回你预测的部份
</rule>
</Rules>
{custom_prompt}
### 响应, 请根据上述要求直接返回JSON结果:
```json
无法预测或者不支持预测的数据请直接返回(不需要返回JSON格式，需要翻译为 {lang} 输出)："抱歉，该数据无法进行预测。(有原因则返回无法预测的原因)"
如果可以预测，则不需要返回原有数据，直接返回预测的部份
user: |
### 字段(字段别名):
<fields>
{fields}
</fields>
### 数据:
<data>
{data}
</data>
datasource:
system: |
### 请使用语言：{lang} 回答
frontend/src/views/system/prompt/index.vue
+1 -1
View File
@@ -202,7 +202,7 @@ const search = () => {
oldKeywords.value = keywords.value
promptApi
.getList(pageInfo.currentPage, pageInfo.pageSize, currentType.value, {
question: keywords.value,
name: keywords.value,
})
.then((res: any) => {
toggleRowLoading.value = true