Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/chat/task/llm.py

@@ -38,6 +38,7 @@
 from apps.system.schemas.system_schema import AssistantOutDsSchema
 from common.core.config import settings
 from common.core.deps import CurrentAssistant, CurrentUser
+from common.error import SingleMessageError
 from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
 
 warnings.filterwarnings("ignore")
@@ -78,7 +79,7 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
         chat_id = chat_question.chat_id
         chat: Chat | None = self.session.get(Chat, chat_id)
         if not chat:
-            raise Exception(f"Chat with id {chat_id} not found")
+            raise SingleMessageError(f"Chat with id {chat_id} not found")
         ds: CoreDatasource | AssistantOutDsSchema | None = None
         if chat.datasource:
             # Get available datasource
@@ -87,13 +88,13 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
                 self.out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
                 ds = self.out_ds_instance.get_ds(chat.datasource)
                 if not ds:
-                    raise Exception("No available datasource configuration found")
+                    raise SingleMessageError("No available datasource configuration found")
                 chat_question.engine = ds.type
                 chat_question.db_schema = self.out_ds_instance.get_db_schema(ds.id)
             else:
                 ds = self.session.get(CoreDatasource, chat.datasource)
                 if not ds:
-                    raise Exception("No available datasource configuration found")
+                    raise SingleMessageError("No available datasource configuration found")
                 chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
                 chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
 
@@ -446,10 +447,11 @@ def select_datasource(self):
                     _ds = self.session.get(CoreDatasource, _datasource)
                     if not _ds:
                         _datasource = None
-                        raise Exception(f"Datasource configuration with id {_datasource} not found")
+                        raise SingleMessageError(f"Datasource configuration with id {_datasource} not found")
                     self.ds = CoreDatasource(**_ds.model_dump())
                     self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
-                    self.chat_question.db_schema = get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds)
+                    self.chat_question.db_schema = get_table_schema(session=self.session,
+                                                                    current_user=self.current_user, ds=self.ds)
                     _engine_type = self.chat_question.engine
                     _chat.engine_type = _ds.type_name
                 # save chat
@@ -459,9 +461,9 @@ def select_datasource(self):
                 self.session.commit()
 
             elif data['fail']:
-                raise Exception(data['fail'])
+                raise SingleMessageError(data['fail'])
             else:
-                raise Exception('No available datasource configuration found')
+                raise SingleMessageError('No available datasource configuration found')
 
         except Exception as e:
             _error = e
@@ -542,7 +544,7 @@ def generate_with_sub_sql(self, sql, sub_mappings: list):
 
         SQLBotLogUtil.info(full_dynamic_text)
         return full_dynamic_text
-    
+
     def generate_assistant_dynamic_sql(self, sql, tables: List):
         ds: AssistantOutDsSchema = self.ds
         sub_query = []
@@ -552,7 +554,7 @@ def generate_assistant_dynamic_sql(self, sql, tables: List):
         if not sub_query:
             return None
         return self.generate_with_sub_sql(sql=sql, sub_mappings=sub_query)
-        
+
     def build_table_filter(self, sql: str, filters: list):
         filter = json.dumps(filters, ensure_ascii=False)
         self.chat_question.sql = sql
@@ -672,23 +674,25 @@ def generate_chart(self):
                                                          full_message=orjson.dumps(
                                                              [{'type': msg.type, 'content': msg.content} for msg in
                                                               self.chart_message]).decode())
+
     def check_sql(self, res: str) -> tuple[any]:
         json_str = extract_nested_json(res)
         if json_str is None:
-            raise Exception("Cannot parse sql from answer")
+            raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer',
+                                                   'traceback': "Cannot parse sql from answer:\n" + res}).decode())
         data: dict = orjson.loads(json_str)
         sql = ''
         message = ''
         if data['success']:
             sql = data['sql']
         else:
             message = data['message']
-            raise Exception(message)
-        
+            raise SingleMessageError(message)
+
         if sql.strip() == '':
-            raise Exception("SQL query is empty")
+            raise SingleMessageError("SQL query is empty")
         return sql, data.get('tables')
-        
+
     def check_save_sql(self, res: str) -> str:
         sql, *_ = self.check_sql(res=res)
         save_sql(session=self.session, sql=sql, record_id=self.record.id)
@@ -701,7 +705,8 @@ def check_save_chart(self, res: str) -> Dict[str, Any]:
 
         json_str = extract_nested_json(res)
         if json_str is None:
-            raise Exception("Cannot parse chart config from answer")
+            raise SingleMessageError(orjson.dumps({'message': 'Cannot parse chart config from answer',
+                                                   'traceback': "Cannot parse chart config from answer:\n" + res}).decode())
         data = orjson.loads(json_str)
 
         chart: Dict[str, Any] = {}
@@ -729,7 +734,7 @@ def check_save_chart(self, res: str) -> Dict[str, Any]:
             error = True
 
         if error:
-            raise Exception(message)
+            raise SingleMessageError(message)
 
         save_chart(session=self.session, chart=orjson.dumps(chart).decode(), record_id=self.record.id)
 
@@ -865,7 +870,7 @@ def run_task(self, in_chat: bool = True):
                     if dynamic_sql_result:
                         SQLBotLogUtil.info(dynamic_sql_result)
                         sql, *_ = self.check_sql(res=dynamic_sql_result)
-                        
+
                     sql_result = self.generate_assistant_filter(sql, tables)
                 else:
                     sql_result = self.generate_filter(sql, tables)  # maybe no sql and tables
@@ -951,12 +956,16 @@ def run_task(self, in_chat: bool = True):
                     SQLBotLogUtil.info(image_url)
                     yield f'![{chart["type"]}]({image_url})'
         except Exception as e:
-            traceback.print_exc()
-            self.save_error(message=str(e))
+            error_msg: str
+            if isinstance(e, SingleMessageError):
+                error_msg = str(e)
+            else:
+                error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
+            self.save_error(message=error_msg)
             if in_chat:
-                yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+                yield orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
             else:
-                yield f'> ❌ **ERROR**\n\n> \n\n> {str(e)}。'
+                yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。'
 
     def run_recommend_questions_task_async(self):
         self.future = executor.submit(self.run_recommend_questions_task_cache)
@@ -1022,31 +1031,37 @@ def run_analysis_or_predict_task(self, action_type: str):
 
             self.finish()
         except Exception as e:
-            traceback.print_exc()
-            self.save_error(message=str(e))
-            yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+            error_msg: str
+            if isinstance(e, SingleMessageError):
+                error_msg = str(e)
+            else:
+                error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
+            self.save_error(message=error_msg)
+            yield orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n'
         finally:
             # end
             pass
-    
+
     def validate_history_ds(self):
         _ds = self.ds
         if not self.current_assistant:
-            current_ds = self.session.get(CoreDatasource, _ds.id)
-            if not current_ds:
-                raise Exception('ds is invalid')
+            try:
+                current_ds = self.session.get(CoreDatasource, _ds.id)
+                if not current_ds:
+                    raise SingleMessageError('chat.ds_is_invalid')
+            except Exception as e:
+                raise SingleMessageError("chat.ds_is_invalid")
         else:
             try:
                 _ds_list: list[dict] = get_assistant_ds(session=self.session, llm_service=self)
                 match_ds = any(item.get("id") == _ds.id for item in _ds_list)
                 if not match_ds:
                     type = self.current_assistant.type
                     msg = f"ds is invalid [please check ds list and public ds list]" if type == 0 else f"ds is invalid [please check ds api]"
-                    raise Exception(msg)
+                    raise SingleMessageError(msg)
             except Exception as e:
-                raise Exception(f"ds is invalid [{str(e)}]")
-           
-            
+                raise SingleMessageError(f"ds is invalid [{str(e)}]")
+
 
 def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
     """Execute SQL query using SQLDatabase

backend/common/error.py

@@ -0,0 +1,7 @@
+class SingleMessageError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message
+
+    def __str__(self):
+        return self.message

frontend/src/i18n/en.json

@@ -391,8 +391,8 @@
     "add_rule_group": "Add rule group",
     "permission_rule": "Permission rules",
     "restricted_user": "Restricted user",
-    "set_rule": "Setup rule",
-    "set_user": "Setup user",
+    "set_rule": "Set rule",
+    "set_user": "Set user",
     "2": "{msg}",
     "238_people": "{msg} people",
     "no_permission_rule": "No permission rule",
@@ -525,7 +525,10 @@
     "inference_process": "Thought Process",
     "thinking": "Thinking",
     "data_analysis": "Data Analysis",
-    "data_predict": "Data Prediction"
+    "data_predict": "Data Prediction",
+    "ds_is_invalid": "Datasource is invalid",
+    "error": "Error",
+    "show_error_detail": "Show error info"
   },
   "about": {
     "title": "About",
@@ -548,4 +551,4 @@
     "back_community": "Restore to Community Edition",
     "confirm_tips": "Are you sure to restore to Community Edition?"
   }
-}
+}

frontend/src/i18n/zh-CN.json

@@ -525,7 +525,10 @@
     "inference_process": "思考过程",
     "thinking": "思考中",
     "data_analysis": "数据分析",
-    "data_predict": "数据预测"
+    "data_predict": "数据预测",
+    "ds_is_invalid": "数据源无效",
+    "error": "错误",
+    "show_error_detail": "查看具体信息"
   },
   "about": {
     "title": "关于",
@@ -548,4 +551,4 @@
     "back_community": "还原至社区版",
     "confirm_tips": "确定还原至社区版？"
   }
-}
+}

frontend/src/views/chat/ErrorInfo.vue

@@ -0,0 +1,92 @@
+<script setup lang="ts">
+import { ref, computed } from 'vue'
+import { useI18n } from 'vue-i18n'
+import { useAssistantStore } from '@/stores/assistant.ts'
+
+const props = defineProps<{
+  error?: string
+}>()
+
+const { t } = useI18n()
+
+const assistantStore = useAssistantStore()
+const isAssistant = computed(() => assistantStore.getAssistant)
+
+const showBlock = computed(() => {
+  return props.error && props.error?.trim().length > 0
+})
+
+const errorMessage = computed(() => {
+  const obj = { message: props.error, showMore: false, traceback: '' }
+  if (showBlock.value && props.error?.trim().startsWith('{') && props.error?.trim().endsWith('}')) {
+    try {
+      const json = JSON.parse(props.error?.trim())
+      obj.message = json['message']
+      obj.traceback = json['traceback']
+      if (obj.traceback?.trim().length > 0) {
+        obj.showMore = true
+      }
+    } catch (e) {
+      console.error(e)
+    }
+  }
+  return obj
+})
+
+const show = ref(false)
+
+function showTraceBack() {
+  show.value = true
+}
+</script>
+
+<template>
+  <div v-if="showBlock">
+    <div
+      v-if="!errorMessage.showMore"
+      v-dompurify-html="errorMessage.message"
+      class="error-container"
+    ></div>
+    <div v-else class="error-container row">
+      {{ t('chat.error') }}
+      <el-button text @click="showTraceBack">{{ t('chat.show_error_detail') }}</el-button>
+    </div>
+
+    <el-drawer
+      v-model="show"
+      :size="isAssistant ? '100%' : '600px'"
+      :title="t('chat.error')"
+      direction="rtl"
+      body-class="chart-sql-error-body"
+    >
+      <el-main>
+        <div v-dompurify-html="errorMessage.traceback" class="error-container open"></div>
+      </el-main>
+    </el-drawer>
+  </div>
+</template>
+
+<style lang="less">
+.chart-sql-error-body {
+  padding: 0;
+}
+</style>
+<style scoped lang="less">
+.error-container {
+  font-weight: 400;
+  font-size: 16px;
+  line-height: 24px;
+  color: rgba(31, 35, 41, 1);
+  white-space: pre-wrap;
+
+  &.row {
+    display: flex;
+    flex-direction: row;
+    align-items: center;
+  }
+  &.open {
+    font-size: 14px;
+    line-height: 20px;
+  }
+}
+</style>