refactor: Optimize question connection pool

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -18,7 +18,7 @@
 from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, BaseMessageChunk
 from sqlalchemy import select
 from sqlalchemy.orm import sessionmaker
-from sqlmodel import create_engine, Session
+from sqlmodel import Session
 
 from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_config
 from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
@@ -37,6 +37,7 @@
 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.db import engine
 from common.core.deps import CurrentAssistant, CurrentUser
 from common.error import SingleMessageError, SQLBotDBError, ParseSQLResultError, SQLBotDBConnectionError
 from common.utils.utils import SQLBotLogUtil, extract_nested_json, prepare_for_orjson
@@ -50,6 +51,10 @@
 dynamic_ds_types = [1, 3]
 dynamic_subsql_prefix = 'select * from sqlbot_dynamic_temp_table_'
 
+session_maker = sessionmaker(bind=engine)
+db_session = session_maker()
+
+
 class LLMService:
     ds: CoreDatasource
     chat_question: ChatQuestion
@@ -59,7 +64,7 @@ class LLMService:
     sql_message: List[Union[BaseMessage, dict[str, Any]]] = []
     chart_message: List[Union[BaseMessage, dict[str, Any]]] = []
 
-    session: Session
+    session: Session = db_session
     current_user: CurrentUser
     current_assistant: Optional[CurrentAssistant] = None
     out_ds_instance: Optional[AssistantOutDs] = None
@@ -79,9 +84,9 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
                  current_assistant: Optional[CurrentAssistant] = None, no_reasoning: bool = False,
                  config: LLMConfig = None):
         self.chunk_list = []
-        engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
-        session_maker = sessionmaker(bind=engine)
-        self.session = session_maker()
+        # engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
+        # session_maker = sessionmaker(bind=engine)
+        # self.session = session_maker()
         self.session.exec = self.session.exec if hasattr(self.session, "exec") else self.session.execute
         self.current_user = current_user
         self.current_assistant = current_assistant