diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index ea65ddc8..584e1dca 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -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