feat: improve error message display

Author: ulleo

backend/apps/chat/task/llm.py

@@ -15,7 +15,6 @@
 from langchain_community.utilities import SQLDatabase
 from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk
 from sqlalchemy import select
-from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import sessionmaker
 from sqlmodel import create_engine, Session
 
@@ -30,13 +29,13 @@
 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
-from apps.db.db import exec_sql, get_version
+from apps.db.db import exec_sql, get_version, check_connection
 from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds
 from apps.system.schemas.system_schema import AssistantOutDsSchema
 from apps.terminology.curd.terminology import get_terminology_template
 from common.core.config import settings
 from common.core.deps import CurrentAssistant, CurrentUser
-from common.error import SingleMessageError
+from common.error import SingleMessageError, SQLBotDBError, ParseSQLResultError, SQLBotDBConnectionError
 from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
 
 warnings.filterwarnings("ignore")
@@ -868,7 +867,14 @@ def execute_sql(self, sql: str):
             Query results
         """
         SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}")
-        return exec_sql(self.ds, sql)
+        try:
+            return exec_sql(self.ds, sql)
+        except Exception as e:
+            if isinstance(e, ParseSQLResultError):
+                raise e
+            else:
+                err = traceback.format_exc(limit=1, chain=True)
+                raise SQLBotDBError(err)
 
     def pop_chunk(self):
         try:
@@ -940,6 +946,12 @@ def run_task(self, in_chat: bool = True):
                                                                               ds=self.ds)
             else:
                 self.validate_history_ds()
+
+            # check connection
+            connected = check_connection(ds=self.ds, trans=None)
+            if not connected:
+                raise SQLBotDBConnectionError('Connect DB failed')
+
             # generate sql
             sql_res = self.generate_sql()
             full_sql_text = ''
@@ -1059,13 +1071,12 @@ def run_task(self, in_chat: bool = True):
             error_msg: str
             if isinstance(e, SingleMessageError):
                 error_msg = str(e)
-            elif isinstance(e, ConnectionError):
+            elif isinstance(e, SQLBotDBConnectionError):
                 error_msg = orjson.dumps(
-                    {'message': str(e), 'traceback': traceback.format_exc(limit=1),
-                     'type': 'db-connection-err'}).decode()
-            elif isinstance(e, DBAPIError):
+                    {'message': str(e), 'type': 'db-connection-err'}).decode()
+            elif isinstance(e, SQLBotDBError):
                 error_msg = orjson.dumps(
-                    {'message': str(e), 'traceback': traceback.format_exc(limit=1), 'type': 'exec-sql-err'}).decode()
+                    {'message': 'Execute SQL Failed', 'traceback': str(e), 'type': 'exec-sql-err'}).decode()
             else:
                 error_msg = orjson.dumps({'message': str(e), 'traceback': traceback.format_exc(limit=1)}).decode()
             self.save_error(message=error_msg)

backend/apps/db/db.py

@@ -3,8 +3,10 @@
 import platform
 import urllib.parse
 from decimal import Decimal
+from typing import Optional
 
 from apps.db.db_sql import get_table_sql, get_field_sql, get_version_sql
+from common.error import ParseSQLResultError
 
 if platform.system() != "Darwin":
     import dmPython
@@ -100,7 +102,7 @@ def get_session(ds: CoreDatasource | AssistantOutDsSchema):
     return session
 
 
-def check_connection(trans: Trans, ds: CoreDatasource, is_raise: bool = False):
+def check_connection(trans: Optional[Trans], ds: CoreDatasource, is_raise: bool = False):
     db = DB.get_db(ds.type)
     if db.connect_type == ConnectType.sqlalchemy:
         conn = get_engine(ds, 10)
@@ -209,7 +211,8 @@ def get_schema(ds: CoreDatasource):
             if ds.type == "sqlServer":
                 sql = f"""select name from sys.schemas"""
             elif ds.type == "pg" or ds.type == "excel":
-                sql = """SELECT nspname FROM pg_namespace"""
+                sql = """SELECT nspname
+                         FROM pg_namespace"""
             elif ds.type == "oracle":
                 sql = f"""select * from all_users"""
             with session.execute(text(sql)) as result:
@@ -325,7 +328,7 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                     return {"fields": columns, "data": result_list,
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                 except Exception as ex:
-                    raise ex
+                    raise ParseSQLResultError(str(ex))
     else:
         conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
         if ds.type == 'dm':
@@ -345,7 +348,7 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                     return {"fields": columns, "data": result_list,
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                 except Exception as ex:
-                    raise ex
+                    raise ParseSQLResultError(str(ex))
         elif ds.type == 'doris':
             with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
                                  port=conf.port, db=conf.database, connect_timeout=conf.timeout,
@@ -364,7 +367,7 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                     return {"fields": columns, "data": result_list,
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                 except Exception as ex:
-                    raise ex
+                    raise ParseSQLResultError(str(ex))
         elif ds.type == 'redshift':
             with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
                                             password=conf.password,
@@ -383,4 +386,4 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                     return {"fields": columns, "data": result_list,
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                 except Exception as ex:
-                    raise ex
+                    raise ParseSQLResultError(str(ex))

backend/common/error.py

@@ -4,4 +4,16 @@ def __init__(self, message):
         self.message = message
 
     def __str__(self):
-        return self.message
+        return self.message
+
+
+class SQLBotDBConnectionError(Exception):
+    pass
+
+
+class SQLBotDBError(Exception):
+    pass
+
+
+class ParseSQLResultError(Exception):
+    pass

frontend/src/views/chat/ErrorInfo.vue

@@ -17,7 +17,7 @@ const showBlock = computed(() => {
 })
 
 const errorMessage = computed(() => {
-  const obj = { message: props.error, showMore: false, traceback: '', type: '' }
+  const obj = { message: props.error, showMore: false, traceback: '', type: undefined }
   if (showBlock.value && props.error?.trim().startsWith('{') && props.error?.trim().endsWith('}')) {
     try {
       const json = JSON.parse(props.error?.trim())
@@ -44,7 +44,7 @@ function showTraceBack() {
 <template>
   <div v-if="showBlock">
     <div
-      v-if="!errorMessage.showMore"
+      v-if="!errorMessage.showMore && errorMessage.type == undefined"
       v-dompurify-html="errorMessage.message"
       class="error-container"
     ></div>
@@ -58,7 +58,9 @@ function showTraceBack() {
       <template v-else>
         {{ t('chat.error') }}
       </template>
-      <el-button text @click="showTraceBack">{{ t('chat.show_error_detail') }}</el-button>
+      <el-button v-if="errorMessage.showMore" text @click="showTraceBack">
+        {{ t('chat.show_error_detail') }}
+      </el-button>
     </div>
 
     <el-drawer