perf: Memory optimization

.github/workflows/build-and-push-vector-model.yml

@@ -5,7 +5,7 @@ on:
     inputs:
       dockerImageTag:
         description: 'Docker Image Tag'
-        default: 'v2.0.2'
+        default: 'v2.0.3'
         required: true
       architecture:
         description: 'Architecture'

.gitignore

@@ -188,4 +188,5 @@ apps/models_provider/impl/*/icon/
 apps/models_provider/impl/tencent_model_provider/credential/stt.py
 apps/models_provider/impl/tencent_model_provider/model/stt.py
 tmp/
-config.yml
+config.yml
+.SANDBOX_BANNED_HOSTS

apps/application/flow/step_node/loop_node/impl/base_loop_node.py

@@ -224,11 +224,17 @@ def handler(self, workflow):
 
 class BaseLoopNode(ILoopNode):
     def save_context(self, details, workflow_manage):
+        self.context['loop_context_data'] = details.get('loop_context_data')
+        self.context['loop_answer_data'] = details.get('loop_answer_data')
+        self.context['loop_node_data'] = details.get('loop_node_data')
         self.context['result'] = details.get('result')
-        for key, value in details['context'].items():
-            if key not in self.context:
-                self.context[key] = value
-        self.answer_text = str(details.get('result'))
+        self.context['params'] = details.get('params')
+        self.context['run_time'] = details.get('run_time')
+        self.context['index'] = details.get('current_index')
+        self.context['item'] = details.get('current_item')
+        for key, value in (details.get('loop_context_data') or {}).items():
+            self.context[key] = value
+        self.answer_text = ""
 
     def get_answer_list(self) -> List[Answer] | None:
         result = []

apps/application/flow/tools.py

@@ -8,7 +8,8 @@
 """
 import asyncio
 import json
-import traceback
+import queue
+import threading
 from typing import Iterator
 
 from django.http import StreamingHttpResponse
@@ -227,6 +228,30 @@ def generate_tool_message_template(name, context):
         return tool_message_template % (name, tool_message_json_template % (context))
 
 
+# 全局单例事件循环
+_global_loop = None
+_loop_thread = None
+_loop_lock = threading.Lock()
+
+
+def get_global_loop():
+    """获取全局共享的事件循环"""
+    global _global_loop, _loop_thread
+
+    with _loop_lock:
+        if _global_loop is None:
+            _global_loop = asyncio.new_event_loop()
+
+            def run_forever():
+                asyncio.set_event_loop(_global_loop)
+                _global_loop.run_forever()
+
+            _loop_thread = threading.Thread(target=run_forever, daemon=True, name="GlobalAsyncLoop")
+            _loop_thread.start()
+
+    return _global_loop
+
+
 async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable=True):
     client = MultiServerMCPClient(json.loads(mcp_servers))
     tools = await client.get_tools()
@@ -242,19 +267,31 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
 
 
 def mcp_response_generator(chat_model, message_list, mcp_servers, mcp_output_enable=True):
-    loop = asyncio.new_event_loop()
-    try:
-        async_gen = _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable)
-        while True:
-            try:
-                chunk = loop.run_until_complete(anext_async(async_gen))
-                yield chunk
-            except StopAsyncIteration:
-                break
-    except Exception as e:
-        maxkb_logger.error(f'Exception: {e}', exc_info=True)
-    finally:
-        loop.close()
+    """使用全局事件循环，不创建新实例"""
+    result_queue = queue.Queue()
+    loop = get_global_loop()  # 使用共享循环
+
+    async def _run():
+        try:
+            async_gen = _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable)
+            async for chunk in async_gen:
+                result_queue.put(('data', chunk))
+        except Exception as e:
+            maxkb_logger.error(f'Exception: {e}', exc_info=True)
+            result_queue.put(('error', e))
+        finally:
+            result_queue.put(('done', None))
+
+    # 在全局循环中调度任务
+    asyncio.run_coroutine_threadsafe(_run(), loop)
+
+    while True:
+        msg_type, data = result_queue.get()
+        if msg_type == 'done':
+            break
+        if msg_type == 'error':
+            raise data
+        yield data
 
 
 async def anext_async(agen):

apps/application/flow/workflow_manage.py

@@ -234,25 +234,62 @@ def run_block(self, language='zh'):
         非流式响应
         @return: 结果
         """
-        self.run_chain_async(None, None, language)
-        while self.is_run():
-            pass
-        details = self.get_runtime_details()
-        message_tokens = sum([row.get('message_tokens') for row in details.values() if
-                              'message_tokens' in row and row.get('message_tokens') is not None])
-        answer_tokens = sum([row.get('answer_tokens') for row in details.values() if
-                             'answer_tokens' in row and row.get('answer_tokens') is not None])
-        answer_text_list = self.get_answer_text_list()
-        answer_text = '\n\n'.join(
-            '\n\n'.join([a.get('content') for a in answer]) for answer in
-            answer_text_list)
-        answer_list = reduce(lambda pre, _n: [*pre, *_n], answer_text_list, [])
-        self.work_flow_post_handler.handler(self)
-        return self.base_to_response.to_block_response(self.params['chat_id'],
-                                                       self.params['chat_record_id'], answer_text, True
-                                                       , message_tokens, answer_tokens,
-                                                       _status=status.HTTP_200_OK if self.status == 200 else status.HTTP_500_INTERNAL_SERVER_ERROR,
-                                                       other_params={'answer_list': answer_list})
+        try:
+            self.run_chain_async(None, None, language)
+            while self.is_run():
+                pass
+            details = self.get_runtime_details()
+            message_tokens = sum([row.get('message_tokens') for row in details.values() if
+                                  'message_tokens' in row and row.get('message_tokens') is not None])
+            answer_tokens = sum([row.get('answer_tokens') for row in details.values() if
+                                 'answer_tokens' in row and row.get('answer_tokens') is not None])
+            answer_text_list = self.get_answer_text_list()
+            answer_text = '\n\n'.join(
+                '\n\n'.join([a.get('content') for a in answer]) for answer in
+                answer_text_list)
+            answer_list = reduce(lambda pre, _n: [*pre, *_n], answer_text_list, [])
+            self.work_flow_post_handler.handler(self)
+
+            res = self.base_to_response.to_block_response(self.params['chat_id'],
+                                                          self.params['chat_record_id'], answer_text, True
+                                                          , message_tokens, answer_tokens,
+                                                          _status=status.HTTP_200_OK if self.status == 200 else status.HTTP_500_INTERNAL_SERVER_ERROR,
+                                                          other_params={'answer_list': answer_list})
+        finally:
+            self._cleanup()
+        return res
+
+    def _cleanup(self):
+        """清理所有对象引用"""
+        # 清理列表
+        self.future_list.clear()
+        self.field_list.clear()
+        self.global_field_list.clear()
+        self.chat_field_list.clear()
+        self.image_list.clear()
+        self.video_list.clear()
+        self.document_list.clear()
+        self.audio_list.clear()
+        self.other_list.clear()
+        if hasattr(self, 'node_context'):
+            self.node_context.clear()
+
+        # 清理字典
+        self.context.clear()
+        self.chat_context.clear()
+        self.form_data.clear()
+
+        # 清理对象引用
+        self.node_chunk_manage = None
+        self.work_flow_post_handler = None
+        self.flow = None
+        self.start_node = None
+        self.current_node = None
+        self.current_result = None
+        self.chat_record = None
+        self.base_to_response = None
+        self.params = None
+        self.lock = None
 
     def run_stream(self, current_node, node_result_future, language='zh'):
         """
@@ -307,6 +344,7 @@ def await_result(self):
                                                                  '',
                                                                  [],
                                                                  '', True, message_tokens, answer_tokens, {})
+            self._cleanup()
 
     def run_chain_async(self, current_node, node_result_future, language='zh'):
         future = executor.submit(self.run_chain_manage, current_node, node_result_future, language)

apps/application/serializers/common.py

@@ -6,7 +6,6 @@
     @date：2025/6/9 13:42
     @desc:
 """
-from datetime import datetime
 from typing import List
 
 from django.core.cache import cache
@@ -226,10 +225,66 @@ def append_chat_record(self, chat_record: ChatRecord):
             chat_record.save()
             ChatCountSerializer(data={'chat_id': self.chat_id}).update_chat()
 
+    def to_dict(self):
+
+        return {
+            'chat_id': self.chat_id,
+            'chat_user_id': self.chat_user_id,
+            'chat_user_type': self.chat_user_type,
+            'knowledge_id_list': self.knowledge_id_list,
+            'exclude_document_id_list': self.exclude_document_id_list,
+            'application_id': self.application_id,
+            'chat_record_list': [self.chat_record_to_map(c) for c in self.chat_record_list],
+            'debug': self.debug
+        }
+
+    def chat_record_to_map(self, chat_record):
+        return {'id': chat_record.id,
+                'chat_id': chat_record.chat_id,
+                'vote_status': chat_record.vote_status,
+                'problem_text': chat_record.problem_text,
+                'answer_text': chat_record.answer_text,
+                'answer_text_list': chat_record.answer_text_list,
+                'message_tokens': chat_record.message_tokens,
+                'answer_tokens': chat_record.answer_tokens,
+                'const': chat_record.const,
+                'details': chat_record.details,
+                'improve_paragraph_id_list': chat_record.improve_paragraph_id_list,
+                'run_time': chat_record.run_time,
+                'index': chat_record.index}
+
+    @staticmethod
+    def map_to_chat_record(chat_record_dict):
+        ChatRecord(id=chat_record_dict.get('id'),
+                   chat_id=chat_record_dict.get('chat_id'),
+                   vote_status=chat_record_dict.get('vote_status'),
+                   problem_text=chat_record_dict.get('problem_text'),
+                   answer_text=chat_record_dict.get('answer_text'),
+                   answer_text_list=chat_record_dict.get('answer_text_list'),
+                   message_tokens=chat_record_dict.get('message_tokens'),
+                   answer_tokens=chat_record_dict.get('answer_tokens'),
+                   const=chat_record_dict.get('const'),
+                   details=chat_record_dict.get('details'),
+                   improve_paragraph_id_list=chat_record_dict.get('improve_paragraph_id_list'),
+                   run_time=chat_record_dict.get('run_time'),
+                   index=chat_record_dict.get('index'), )
+
     def set_cache(self):
-        cache.set(Cache_Version.CHAT.get_key(key=self.chat_id), self, version=Cache_Version.CHAT.get_version(),
+        cache.set(Cache_Version.CHAT.get_key(key=self.chat_id), self.to_dict(),
+                  version=Cache_Version.CHAT_INFO.get_version(),
                   timeout=60 * 30)
 
+    @staticmethod
+    def map_to_chat_info(chat_info_dict):
+        return ChatInfo(chat_info_dict.get('chat_id'), chat_info_dict.get('chat_user_id'),
+                        chat_info_dict.get('chat_user_type'), chat_info_dict.get('knowledge_id_list'),
+                        chat_info_dict.get('exclude_document_id_list'),
+                        chat_info_dict.get('application_id'),
+                        [ChatInfo.map_to_chat_record(c_r) for c_r in chat_info_dict.get('chat_record_list')])
+
     @staticmethod
     def get_cache(chat_id):
-        return cache.get(Cache_Version.CHAT.get_key(key=chat_id), version=Cache_Version.CHAT.get_version())
+        chat_info_dict = cache.get(Cache_Version.CHAT.get_key(key=chat_id), version=Cache_Version.CHAT_INFO.get_version())
+        if chat_info_dict:
+            return ChatInfo.map_to_chat_info(chat_info_dict)
+        return None

apps/application/views/application_chat.py

@@ -125,8 +125,8 @@ class OpenView(APIView):
         responses=None,
         tags=[_('Application')]  # type: ignore
     )
-    @has_permissions(PermissionConstants.APPLICATION_EDIT.get_workspace_application_permission(),
-                     PermissionConstants.APPLICATION_EDIT.get_workspace_permission_workspace_manage_role(),
+    @has_permissions(PermissionConstants.APPLICATION_READ.get_workspace_application_permission(),
+                     PermissionConstants.APPLICATION_READ.get_workspace_permission_workspace_manage_role(),
                      ViewPermission([RoleConstants.USER.get_workspace_role()],
                                     [PermissionConstants.APPLICATION.get_workspace_application_permission()],
                                     CompareConstants.AND),
@@ -167,8 +167,8 @@ class PromptGenerateView(APIView):
         responses=None,
         tags=[_('Application')]  # type: ignore
     )
-    @has_permissions(PermissionConstants.APPLICATION_EDIT.get_workspace_application_permission(),
-                     PermissionConstants.APPLICATION_EDIT.get_workspace_permission_workspace_manage_role(),
+    @has_permissions(PermissionConstants.APPLICATION_READ.get_workspace_application_permission(),
+                     PermissionConstants.APPLICATION_READ.get_workspace_permission_workspace_manage_role(),
                      ViewPermission([RoleConstants.USER.get_workspace_role()],
                                     [PermissionConstants.APPLICATION.get_workspace_application_permission()],
                                     CompareConstants.AND),

apps/application/views/application_chat_record.py

@@ -93,8 +93,8 @@ class ApplicationChatRecordOperateAPI(APIView):
     )
     @has_permissions(PermissionConstants.APPLICATION_CHAT_LOG_READ.get_workspace_application_permission(),
                      PermissionConstants.APPLICATION_CHAT_LOG_READ.get_workspace_permission_workspace_manage_role(),
-                     PermissionConstants.APPLICATION_EDIT.get_workspace_application_permission(),
-                     PermissionConstants.APPLICATION_EDIT.get_workspace_permission_workspace_manage_role(),
+                     PermissionConstants.APPLICATION_READ.get_workspace_application_permission(),
+                     PermissionConstants.APPLICATION_READ.get_workspace_permission_workspace_manage_role(),
                      ViewPermission([RoleConstants.USER.get_workspace_role()],
                                     [PermissionConstants.APPLICATION.get_workspace_application_permission()],
                                     CompareConstants.AND),

apps/common/auth/common.py

@@ -45,7 +45,7 @@ def to_string(self):
         value = json.dumps(self.to_dict())
         authentication = encrypt(value)
         cache_key = hashlib.sha256(authentication.encode()).hexdigest()
-        authentication_cache.set(cache_key, value, version=Cache_Version.CHAT.value, timeout=60 * 60 * 2)
+        authentication_cache.set(cache_key, value, version=Cache_Version.CHAT.get_version(), timeout=60 * 60 * 2)
         return authentication
 
     @staticmethod

apps/common/config/tokenizer_manage_config.py

@@ -7,18 +7,24 @@
     @desc:
 """
 
+import os
+
+class MKTokenizer:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def encode(self, text):
+        return self.tokenizer.encode(text).ids
+
 
 class TokenizerManage:
     tokenizer = None
 
     @staticmethod
     def get_tokenizer():
-        from transformers import BertTokenizer
-        if TokenizerManage.tokenizer is None:
-            TokenizerManage.tokenizer = BertTokenizer.from_pretrained(
-                'bert-base-cased',
-                cache_dir="/opt/maxkb-app/model/tokenizer",
-                local_files_only=True,
-                resume_download=False,
-                force_download=False)
-        return TokenizerManage.tokenizer
+        from tokenizers import Tokenizer
+        # 创建Tokenizer
+        model_path = os.path.join("/opt/maxkb-app", "model", "tokenizer", "models--bert-base-cased")
+        with open(f"{model_path}/refs/main", encoding="utf-8") as f: snapshot = f.read()
+        TokenizerManage.tokenizer = Tokenizer.from_file(f"{model_path}/snapshots/{snapshot}/tokenizer.json")
+        return MKTokenizer(TokenizerManage.tokenizer)

apps/common/constants/cache_version.py

@@ -30,6 +30,8 @@ class Cache_Version(Enum):
     # 对话
     CHAT = "CHAT", lambda key: key
 
+    CHAT_INFO = "CHAT_INFO", lambda key: key
+
     CHAT_VARIABLE = "CHAT_VARIABLE", lambda key: key
 
     # 应用API KEY

apps/common/event/__init__.py

@@ -7,9 +7,10 @@
     @desc:
 """
 from django.core.cache import cache
+from django.db.models import QuerySet
 from django.utils.translation import gettext as _
 
-from .listener_manage import *
+
 from ..constants.cache_version import Cache_Version
 from ..db.sql_execute import update_execute
 from ..utils.lock import RedisLock