mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
refactor: Optimize question connection pool
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user