refactor: Optimize question connection pool

This commit is contained in:
junjun
2025-09-12 12:52:22 +08:00
parent 6e24ea91f9
commit 1e1138263c
+10 -5
View File
@@ -18,7 +18,7 @@ from langchain_community.utilities import SQLDatabase
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.crud.assistant import AssistantOutDs, AssistantOutDsFactory, ge
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 @@ executor = ThreadPoolExecutor(max_workers=200)
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 @@ class LLMService:
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