feat: support carry last execute SQL error info to AI

Author: ulleo

backend/apps/chat/curd/chat.py

@@ -71,6 +71,21 @@ def get_chart_config(session: SessionDep, chart_record_id: int):
     return {}
 
 
+def get_last_execute_sql_error(session: SessionDep, chart_id: int):
+    stmt = select(ChatRecord.error).where(and_(ChatRecord.chat_id == chart_id)).order_by(
+        ChatRecord.create_time.desc()).limit(1)
+    res = session.execute(stmt).scalar()
+    if res:
+        try:
+            obj = orjson.loads(res)
+            if obj.get('type') and obj.get('type') == 'exec-sql-err':
+                return obj.get('traceback')
+        except Exception:
+            pass
+
+    return None
+
+
 def get_chat_chart_data(session: SessionDep, chart_record_id: int):
     stmt = select(ChatRecord.data).where(and_(ChatRecord.id == chart_record_id))
     res = session.execute(stmt)
@@ -701,7 +716,8 @@ def get_old_questions(session: SessionDep, datasource: int):
     if not datasource:
         return records
     stmt = select(ChatRecord.question).where(
-        and_(ChatRecord.datasource == datasource, ChatRecord.question.isnot(None), ChatRecord.error.is_(None))).order_by(
+        and_(ChatRecord.datasource == datasource, ChatRecord.question.isnot(None),
+             ChatRecord.error.is_(None))).order_by(
         ChatRecord.create_time.desc()).limit(20)
     result = session.execute(stmt)
     for r in result:

backend/apps/chat/models/chat_model.py

@@ -159,6 +159,7 @@ class ChatInfo(BaseModel):
 
 
 class AiModelQuestion(BaseModel):
+    question: str = None
     ai_modal_id: int = None
     ai_modal_name: str = None  # Specific model name
     engine: str = ""
@@ -171,14 +172,15 @@ class AiModelQuestion(BaseModel):
     filter: str = []
     sub_query: Optional[list[dict]] = None
     terminologies: 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)
 
     def sql_user_question(self, current_time: str):
         return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question,
-                                                 rule=self.rule, current_time=current_time)
+                                                 rule=self.rule, current_time=current_time, error_msg=self.error_msg)
 
     def chart_sys_question(self):
         return get_chart_template()['system'].format(sql=self.sql, question=self.question, lang=self.lang)
@@ -226,7 +228,6 @@ def dynamic_user_question(self):
 
 
 class ChatQuestion(AiModelQuestion):
-    question: str
     chat_id: int
 
 

backend/apps/chat/task/llm.py

@@ -24,7 +24,8 @@
     finish_record, save_analysis_answer, save_predict_answer, save_predict_data, \
     save_select_datasource_answer, save_recommend_question_answer, \
     get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \
-    get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log
+    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.datasource.crud.datasource import get_table_schema
 from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user
@@ -70,6 +71,8 @@ class LLMService:
     chunk_list: List[str] = []
     future: Future
 
+    last_execute_sql_error: str = None
+
     def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
                  current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False,
                  config: LLMConfig = None):
@@ -127,6 +130,15 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
         llm_instance = LLMFactory.create_llm(self.config)
         self.llm = llm_instance.llm
 
+        # get last_execute_sql_error
+        last_execute_sql_error = get_last_execute_sql_error(self.session, self.chat_question.chat_id)
+        if last_execute_sql_error:
+            self.chat_question.error_msg = f'''<error-msg>
+{last_execute_sql_error}
+</error-msg>'''
+        else:
+            self.chat_question.error_msg = ''
+
         self.init_messages()
 
     @classmethod

backend/template.yaml

@@ -98,7 +98,7 @@ template:
             - 柱状图或折线图：适合展示在横轴的字段优先排序，若SQL包含分类字段，则分类字段次一级排序
         </rule>
         <rule>
-          如果用户没有指定数据条数的限制，输出的查询SQL需要加上1000条的数据条数限制
+          如果用户没有指定数据条数的限制，输出的查询SQL必须加上1000条的数据条数限制
           如果用户指定的限制大于1000，则按1000处理
           <example>
           以PostgreSQL为例，查询Schema为TEST表TABLE下id字段，则生成的SQL为：
@@ -235,6 +235,7 @@ template:
         {current_time}
         </current-time>
       <background-infos>
+      {error_msg}
       <user-question>
       {question}
       </user-question>