From 3ecf2a0c182a0aa45b6d2b30a6a8f1559cea7eb3 Mon Sep 17 00:00:00 2001 From: ulleo Date: Wed, 17 Sep 2025 14:36:32 +0800 Subject: [PATCH] feat: add SQL examples in prompt --- backend/apps/chat/task/llm.py | 4 ++++ backend/apps/data_training/curd/data_training.py | 4 ++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index ae2eb9f9..5455187d 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -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 diff --git a/backend/apps/data_training/curd/data_training.py b/backend/apps/data_training/curd/data_training.py index 2ef63361..c99b02ee 100644 --- a/backend/apps/data_training/curd/data_training.py +++ b/backend/apps/data_training/curd/data_training.py @@ -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)