feat: optimize db connection when llm invoking (#2774)

api/core/app_runner/assistant_app_runner.py

@@ -195,6 +195,10 @@ class AssistantApplicationRunner(AppRunner):
         if set([ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL]).intersection(model_schema.features or []):
             agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING
 
+        db.session.refresh(conversation)
+        db.session.refresh(message)
+        db.session.close()
+
         # start agent runner
         if agent_entity.strategy == AgentEntity.Strategy.CHAIN_OF_THOUGHT:
             assistant_cot_runner = AssistantCotApplicationRunner(

api/core/app_runner/basic_app_runner.py

@@ -192,6 +192,8 @@ class BasicApplicationRunner(AppRunner):
             model=app_orchestration_config.model_config.model
         )
 
+        db.session.close()
+
         invoke_result = model_instance.invoke_llm(
             prompt_messages=prompt_messages,
             model_parameters=app_orchestration_config.model_config.parameters,

api/core/app_runner/generate_task_pipeline.py

@@ -89,6 +89,10 @@ class GenerateTaskPipeline:
         Process generate task pipeline.
         :return:
         """
+        db.session.refresh(self._conversation)
+        db.session.refresh(self._message)
+        db.session.close()
+
         if stream:
             return self._process_stream_response()
         else:
@@ -303,6 +307,7 @@ class GenerateTaskPipeline:
                     .first()
                 )
                 db.session.refresh(agent_thought)
+                db.session.close()
 
                 if agent_thought:
                     response = {
@@ -330,6 +335,8 @@ class GenerateTaskPipeline:
                     .filter(MessageFile.id == event.message_file_id)
                     .first()
                 )
+                db.session.close()
+
                 # get extension
                 if '.' in message_file.url:
                     extension = f'.{message_file.url.split(".")[-1]}'
@@ -413,6 +420,7 @@ class GenerateTaskPipeline:
         usage = llm_result.usage
 
         self._message = db.session.query(Message).filter(Message.id == self._message.id).first()
+        self._conversation = db.session.query(Conversation).filter(Conversation.id == self._conversation.id).first()
 
         self._message.message = self._prompt_messages_to_prompt_for_saving(self._task_state.llm_result.prompt_messages)
         self._message.message_tokens = usage.prompt_tokens