fix: Return the link to the connection pool after the node execution is completed

Author: shaohuzhang1

apps/application/flow/step_node/tool_lib_node/i_tool_lib_node.py

@@ -8,6 +8,7 @@
 """
 from typing import Type
 
+from django.db import connection
 from django.db.models import QuerySet
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
@@ -31,6 +32,8 @@ class FunctionLibNodeParamsSerializer(serializers.Serializer):
     def is_valid(self, *, raise_exception=False):
         super().is_valid(raise_exception=True)
         f_lib = QuerySet(Tool).filter(id=self.data.get('tool_lib_id')).first()
+        # 归还链接到连接池
+        connection.close()
         if f_lib is None:
             raise Exception(_('The function has been deleted'))
 

apps/application/flow/workflow_manage.py

@@ -14,7 +14,7 @@
 from functools import reduce
 from typing import List, Dict
 
-from django.db import close_old_connections
+from django.db import close_old_connections, connection
 from django.utils import translation
 from django.utils.translation import get_language
 from langchain_core.prompts import PromptTemplate
@@ -433,6 +433,8 @@ def hand_event_node_result(self, current_node, node_result_future):
             return None
         finally:
             current_node.node_chunk.end()
+            # 归还链接到连接池
+            connection.close()
 
     def run_node_async(self, node):
         future = executor.submit(self.run_node, node)

apps/models_provider/tools.py

@@ -106,6 +106,8 @@ def is_valid_credential(provider, model_type, model_name, model_credential: Dict
 
 def get_model_by_id(_id, workspace_id):
     model = QuerySet(Model).filter(id=_id).first()
+    # 归还链接到连接池
+    connection.close()
     get_authorized_model = DatabaseModelManage.get_model("get_authorized_model")
     if model and model.workspace_id != workspace_id and get_authorized_model is not None:
         model = get_authorized_model(QuerySet(Model).filter(id=_id), workspace_id).first()