Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

.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
 
 

backend/apps/system/crud/user.py

@@ -1,13 +1,13 @@
 
 from typing import Optional
 from sqlmodel import Session, select, delete as sqlmodel_delete
-import logging
 from apps.system.models.system_model import UserWsModel, WorkspaceModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import BaseUserDTO, UserInfoDTO, UserWs
 from common.core.deps import SessionDep
 from common.core.sqlbot_cache import cache, clear_cache
 from common.utils.locale import I18n
+from common.utils.utils import SQLBotLogUtil
 from ..models.user import UserModel
 from common.core.security import verify_md5pwd
 
@@ -68,4 +68,4 @@ async def single_delete(session: SessionDep, id: int):
 
 #@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")    
 async def clean_user_cache(id: int):
-    logging.info(f"User cache for [{id}] has been cleaned")
+    SQLBotLogUtil.info(f"User cache for [{id}] has been cleaned")

backend/apps/system/middleware/auth.py

@@ -1,6 +1,6 @@
 
 from typing import Optional
-from fastapi import HTTPException, Request
+from fastapi import Request
 from fastapi.responses import JSONResponse
 import jwt
 from sqlmodel import Session
@@ -13,6 +13,7 @@
 from common.core import security
 from common.core.config import settings
 from common.core.schemas import TokenPayload
+from common.utils.utils import SQLBotLogUtil
 from common.utils.whitelist import whiteUtils
 from fastapi.security.utils import get_authorization_scheme_param
 class TokenMiddleware(BaseHTTPMiddleware):
@@ -70,6 +71,7 @@ async def validateToken(self, token: Optional[str]):
                     ) """
                 return True, session_user
         except Exception as e:
+            SQLBotLogUtil.exception(f"Token validation error: {str(e)}", exc_info=True)
             return False, e
 
 
@@ -95,4 +97,6 @@ async def validateAssistant(self, assistantToken: Optional[str]) -> tuple[any]:
                 assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
                 return True, session_user, assistant_info
         except Exception as e:
+            SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}", exc_info=True)
+            # Return False and the exception message
             return False, e

backend/common/core/config.py

@@ -60,6 +60,11 @@ def all_cors_origins(self) -> list[str]:
     TOKEN_KEY: str =  "X-SQLBOT-TOKEN"
     DEFAULT_PWD: str = "SQLBot@123456"
     ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN"
+    
+    LOG_LEVEL: str = "INFO"  # DEBUG, INFO, WARNING, ERROR
+    LOG_DIR: str = "logs"
+    LOG_FORMAT: str = "%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"
+    SQL_DEBUG: bool = False
 
     @computed_field  # type: ignore[prop-decorator]
     @property

backend/common/core/response_middleware.py

@@ -1,10 +1,10 @@
 import json
-import logging
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import JSONResponse
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from common.core.config import settings
+from common.utils.utils import SQLBotLogUtil
 class ResponseMiddleware(BaseHTTPMiddleware):
     def __init__(self, app):
         super().__init__(app)
@@ -48,7 +48,7 @@ async def dispatch(self, request, call_next):
                     }
                 )
             except Exception as e:
-                logging.error(f"Response processing error: {str(e)}", exc_info=True)
+                SQLBotLogUtil.error(f"Response processing error: {str(e)}", exc_info=True)
                 return JSONResponse(
                     status_code=500,
                     content={
@@ -68,6 +68,7 @@ async def dispatch(self, request, call_next):
 class exception_handler():
     @staticmethod
     async def http_exception_handler(request: Request, exc: HTTPException):
+        SQLBotLogUtil.exception(f"HTTP Exception: {exc.detail}", exc_info=True)
         return JSONResponse(
             status_code=exc.status_code,
             content={
@@ -81,8 +82,7 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 
     @staticmethod
     async def global_exception_handler(request: Request, exc: Exception):
-        logging.error(f"Error info: {str(exc)}", exc_info=True)
-        
+        SQLBotLogUtil.exception(f"Unhandled Exception: {str(exc)}", exc_info=True)
         return JSONResponse(
             status_code=500,
             content={

backend/common/core/sqlbot_cache.py

@@ -1,12 +1,13 @@
-import logging
 from fastapi_cache import FastAPICache
 from fastapi_cache.decorator import cache as original_cache
 from functools import partial, wraps
-from typing import Optional, Set, Any, Dict, Tuple, Callable
+from typing import Optional, Set, Any, Dict, Tuple
 from inspect import Parameter, signature
 from contextlib import asynccontextmanager
 import asyncio
 
+from common.utils.utils import SQLBotLogUtil
+
 
 # 锁管理
 _cache_locks = {}
@@ -65,7 +66,7 @@ def custom_key_builder(
                     base_key += f"{value}:"
 
             except (IndexError, KeyError, AttributeError) as e:
-                logging.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
+                SQLBotLogUtil.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
 
         return base_key
 
@@ -127,19 +128,16 @@ async def wrapper(*args, **kwargs):
             )
 
             async with _get_cache_lock(cache_key):
-                logging.info(f"Using cache key: {cache_key}")
-                print(f"Using cache key: {cache_key}")
+                SQLBotLogUtil.info(f"Using cache key: {cache_key}")
                 backend = FastAPICache.get_backend()
                 cached_value = await backend.get(cache_key)
                 if cached_value is not None:
-                    logging.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
-                    print(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
+                    SQLBotLogUtil.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
                     return cached_value
 
                 result = await func(*args, **kwargs)
                 await backend.set(cache_key, result, expire)
-                logging.info(f"Cache miss for key: {cache_key}, result cached.")
-                print(f"Cache miss for key: {cache_key}, result cached.")
+                SQLBotLogUtil.info(f"Cache miss for key: {cache_key}, result cached.")
                 return result
 
         return wrapper
@@ -167,8 +165,7 @@ async def wrapper(*args, **kwargs):
                 if await FastAPICache.get_backend().get(cache_key):
                     await FastAPICache.clear(key=cache_key)
                     result = await func(*args, **kwargs)
-                    logging.info(f"Clearing cache for key: {cache_key}")
-                    print(f"Clearing cache for key: {cache_key}")
+                    SQLBotLogUtil.info(f"Clearing cache for key: {cache_key}")
                     return result
                 return await func(*args, **kwargs)