optimize db connections

api/core/app/apps/advanced_chat/app_generator.py

@@ -95,12 +95,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
             extras=extras
         )
 
-        workflow = db.session.query(Workflow).filter(Workflow.id == workflow.id).first()
-        user = (db.session.query(Account).filter(Account.id == user.id).first()
-                if isinstance(user, Account)
-                else db.session.query(EndUser).filter(EndUser.id == user.id).first())
-        db.session.close()
-
         # init generate records
         (
             conversation,
@@ -206,7 +200,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         :return:
         """
         # init generate task pipeline
-
         generate_task_pipeline = AdvancedChatAppGenerateTaskPipeline(
             application_generate_entity=application_generate_entity,
             workflow=workflow,

api/core/app/apps/advanced_chat/generate_task_pipeline.py

@@ -122,13 +122,15 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
         self._output_moderation_handler = self._init_output_moderation()
         self._stream = stream
 
-        db.session.close()
-
     def process(self) -> Union[dict, Generator]:
         """
         Process generate task pipeline.
         :return:
         """
+        db.session.refresh(self._workflow)
+        db.session.refresh(self._user)
+        db.session.close()
+
         if self._stream:
             return self._process_stream_response()
         else:

api/core/app/apps/message_based_app_generator.py

@@ -176,9 +176,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
 
             db.session.add(conversation)
             db.session.commit()
-
-            conversation = db.session.query(Conversation).filter(Conversation.id == conversation.id).first()
-            db.session.close()
+            db.session.refresh(conversation)
 
         message = Message(
             app_id=app_config.app_id,
@@ -206,9 +204,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
 
         db.session.add(message)
         db.session.commit()
-
-        message = db.session.query(Message).filter(Message.id == message.id).first()
-        db.session.close()
+        db.session.refresh(message)
 
         for file in application_generate_entity.files:
             message_file = MessageFile(
@@ -224,8 +220,6 @@ class MessageBasedAppGenerator(BaseAppGenerator):
             db.session.add(message_file)
             db.session.commit()
 
-        db.session.close()
-
         return conversation, message
 
     def _get_conversation_introduction(self, application_generate_entity: AppGenerateEntity) -> str:

api/core/app/apps/workflow/generate_task_pipeline.py

@@ -99,13 +99,15 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
         self._output_moderation_handler = self._init_output_moderation()
         self._stream = stream
 
-        db.session.close()
-
     def process(self) -> Union[dict, Generator]:
         """
         Process generate task pipeline.
         :return:
         """
+        db.session.refresh(self._workflow)
+        db.session.refresh(self._user)
+        db.session.close()
+
         if self._stream:
             return self._process_stream_response()
         else:

api/core/app/apps/workflow_based_generate_task_pipeline.py

@@ -60,8 +60,7 @@ class WorkflowBasedGenerateTaskPipeline:
 
         db.session.add(workflow_run)
         db.session.commit()
-
-        workflow_run = db.session.query(WorkflowRun).filter(WorkflowRun.id == workflow_run.id).first()
+        db.session.refresh(workflow_run)
         db.session.close()
 
         return workflow_run
@@ -155,9 +154,7 @@ class WorkflowBasedGenerateTaskPipeline:
 
         db.session.add(workflow_node_execution)
         db.session.commit()
-
-        workflow_node_execution = (db.session.query(WorkflowNodeExecution)
-                                   .filter(WorkflowNodeExecution.id == workflow_node_execution.id).first())
+        db.session.refresh(workflow_node_execution)
         db.session.close()
 
         return workflow_node_execution