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

Author: dataeaseShu

.env.example

@@ -14,6 +14,12 @@ FIRST_SUPERUSER_PASSWORD=123456  # Change this to your pwd
 
 TOKEN_KEY="X-SQLBOT-TOKEN"
 DEFAULT_PWD="SQLBot@123456"
+ASSISTANT_TOKEN_KEY="X-SQLBOT-ASSISTANT-TOKEN"
+
+LOG_LEVEL="INFO"  # DEBUG, INFO, WARNING, ERROR
+LOG_DIR="/opt/sqlbot/logs"
+LOG_FORMAT="%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"
+SQL_DEBUG=False
 
 # Postgres
 POSTGRES_SERVER=localhost

backend/alembic/env.py

@@ -5,7 +5,6 @@
 sys.path.insert(0, dirname(dirname(abspath(__file__))))
 
 import os
-from logging.config import fileConfig
 
 from alembic import context
 from sqlalchemy import engine_from_config, pool
@@ -16,7 +15,6 @@
 
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.
-fileConfig(config.config_file_name)
 
 # add your model's MetaData object here
 # for 'autogenerate' support

backend/apps/chat/task/llm.py

@@ -1,6 +1,5 @@
 import json
 import sqlparse
-import logging
 import traceback
 import warnings
 from typing import Any, List, Optional, Union, Dict
@@ -35,7 +34,7 @@
 from apps.system.schemas.system_schema import AssistantOutDsSchema
 from common.core.config import settings
 from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
-from common.utils.utils import extract_nested_json
+from common.utils.utils import SQLBotLogUtil, extract_nested_json
 
 warnings.filterwarnings("ignore")
 
@@ -223,7 +222,7 @@ def generate_analysis(self):
         res = self.llm.stream(analysis_msg)
         token_usage = {}
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -273,7 +272,7 @@ def generate_predict(self):
         res = self.llm.stream(predict_msg)
         token_usage = {}
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -323,7 +322,7 @@ def generate_recommend_questions_task(self):
         token_usage = {}
         res = self.llm.stream(guess_msg)
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -390,7 +389,7 @@ def select_datasource(self):
         token_usage = {}
         res = self.llm.stream(datasource_msg)
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -474,7 +473,7 @@ def generate_sql(self):
         token_usage = {}
         res = self.llm.stream(self.sql_message)
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -543,7 +542,7 @@ def generate_filter(self, sql: str, tables: List):
         res = self.llm.stream(msg)
         token_usage = {}
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -567,7 +566,7 @@ def generate_filter(self, sql: str, tables: List):
         #                                                                                 'content': msg.content} for msg
         #                                                                                in
         #                                                                                analysis_msg]).decode())
-        print(full_filter_text)
+        SQLBotLogUtil.info(full_filter_text)
         return full_filter_text
 
     def generate_chart(self):
@@ -582,7 +581,7 @@ def generate_chart(self):
         token_usage = {}
         res = self.llm.stream(self.chart_message)
         for chunk in res:
-            print(chunk)
+            SQLBotLogUtil.info(chunk)
             reasoning_content_chunk = ''
             if 'reasoning_content' in chunk.additional_kwargs:
                 reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
@@ -691,7 +690,7 @@ def execute_sql(self, sql: str):
         Returns:
             Query results
         """
-        print(f"Executing SQL on ds_id {self.ds.id}: {sql}")
+        SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}")
         return exec_sql(self.ds, sql)
 
 
@@ -717,7 +716,7 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
 
     except Exception as e:
         error_msg = f"SQL execution failed: {str(e)}"
-        logging.error(error_msg)
+        SQLBotLogUtil.exception(error_msg)
         raise RuntimeError(error_msg)
 
 
@@ -741,7 +740,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
             ds_res = llm_service.select_datasource()
 
             for chunk in ds_res:
-                print(chunk)
+                SQLBotLogUtil.info(chunk)
                 if in_chat:
                     yield orjson.dumps(
                         {'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
@@ -765,7 +764,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
             yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
 
         # filter sql
-        print(full_sql_text)
+        SQLBotLogUtil.info(full_sql_text)
 
         # todo row permission
         sql_json_str = extract_nested_json(full_sql_text)
@@ -785,11 +784,11 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
             raise Exception("SQL query is empty")
 
         sql_result = llm_service.generate_filter(data.get('sql'), data.get('tables'))  # maybe no sql and tables
-        print(sql_result)
+        SQLBotLogUtil.info(sql_result)
         sql = llm_service.check_save_sql(res=sql_result)
         # sql = llm_service.check_save_sql(res=full_sql_text)
 
-        print(sql)
+        SQLBotLogUtil.info(sql)
         format_sql = sqlparse.format(sql, reindent=True)
         if in_chat:
             yield orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n'
@@ -815,9 +814,9 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
             yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
 
         # filter chart
-        print(full_chart_text)
+        SQLBotLogUtil.info(full_chart_text)
         chart = llm_service.check_save_chart(res=full_chart_text)
-        print(chart)
+        SQLBotLogUtil.info(chart)
         if in_chat:
             yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
         else:
@@ -856,7 +855,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
             if chart['type'] != 'table':
                 yield '### generated chart picture\n\n'
                 image_url = request_picture(llm_service.record.chat_id, llm_service.record.id, chart, result)
-                print(image_url)
+                SQLBotLogUtil.info(image_url)
                 yield f'![{chart["type"]}]({image_url})'
     except Exception as e:
         traceback.print_exc()

backend/apps/datasource/crud/datasource.py

@@ -15,7 +15,7 @@
 from apps.db.engine import get_engine_config, get_engine_conn
 from apps.db.type import db_type_relation
 from common.core.deps import SessionDep, CurrentUser, Trans
-from common.utils.utils import deepcopy_ignore_extra
+from common.utils.utils import SQLBotLogUtil, deepcopy_ignore_extra
 from .table import get_tables_by_ds_id
 from ..crud.field import delete_field_by_ds_id, update_field
 from ..crud.table import delete_table_by_ds_id, update_table
@@ -39,12 +39,12 @@ def check_status(session: SessionDep, ds: CoreDatasource, is_raise: bool = False
     conn = get_engine(ds)
     try:
         with conn.connect() as connection:
-            print("success")
+            SQLBotLogUtil.info("success")
             return True
     except Exception as e:
+        SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
         if is_raise:
             raise HTTPException(status_code=500, detail=f'Connect Failed: {e.args}')
-        print("Fail:", e)
         return False
 
 
@@ -86,6 +86,7 @@ def create_ds(session: SessionDep, trans: Trans, user: CurrentUser, create_ds: C
 
 def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]):
     ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
+    check_status(session, ds, True)
     sync_table(session, ds, tables)
     updateNum(session, ds)
 
@@ -131,6 +132,7 @@ def getTables(session: SessionDep, id: int):
 
 
 def getTablesByDs(session: SessionDep, ds: CoreDatasource):
+    check_status(session, ds, True)
     tables = get_tables(ds)
     return tables
 
@@ -244,6 +246,7 @@ def preview(session: SessionDep, id: int, data: TableObj):
         return {"fields": [], "data": [], "sql": ''}
 
     ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
+    check_status(session, ds, True)
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
     sql: str = ""
     if ds.type == "mysql":