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: add SQL examples in prompt

Browse Source
This commit is contained in:
ulleo
2025-09-17 14:36:32 +08:00
parent a6c067a30f
commit 3ecf2a0c18
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+4
View File
@@ -29,6 +29,7 @@ from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \
get_last_execute_sql_error
from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum
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
from apps.datasource.models.datasource import CoreDatasource
@@ -935,6 +936,9 @@ class LLMService:
self.chat_question.terminologies = get_terminology_template(self.session, self.chat_question.question,
self.ds.oid if isinstance(self.ds,
CoreDatasource) else 1)
self.chat_question.data_training = get_training_template(self.session, self.chat_question.question,
self.ds.id, self.ds.oid)
self.init_messages()
# return id
backend/apps/data_training/curd/data_training.py
+2 -2
View File
@@ -204,7 +204,7 @@ LIMIT {settings.EMBEDDING_DATA_TRAINING_TOP_COUNT}
"""
def select_terminology_by_question(session: SessionDep, question: str, oid: int, datasource: int):
def select_training_by_question(session: SessionDep, question: str, oid: int, datasource: int):
if question.strip() == "":
return []
@@ -303,7 +303,7 @@ def get_training_template(session: SessionDep, question: str, datasource: int, o
oid = 1
if not datasource:
return ''
_results = select_terminology_by_question(session, question, oid, datasource)
_results = select_training_by_question(session, question, oid, datasource)
if _results and len(_results) > 0:
data_training = to_xml_string(_results)
template = get_base_data_training_template().format(data_training=data_training)