feat: integrate ResourceMappingSerializer for resource count retrieval across multiple modules

Author: wxg0103

apps/application/serializers/application.py

@@ -49,6 +49,8 @@
 from models_provider.models import Model
 from models_provider.tools import get_model_instance_by_model_workspace_id
 from system_manage.models import WorkspaceUserResourcePermission, AuthTargetType
+from system_manage.models.resource_mapping import ResourceMapping
+from system_manage.serializers.resource_mapping_serializers import ResourceMappingSerializer
 from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer
 from tools.models import Tool, ToolScope
 from tools.serializers.tool import ToolExportModelSerializer
@@ -411,12 +413,15 @@ def page(self, current_page: int, page_size: int, instance: Dict):
         user_id = self.data.get("user_id")
         workspace_manage = is_workspace_manage(user_id, workspace_id)
         is_x_pack_ee = self.is_x_pack_ee()
-        return native_page_search(current_page, page_size, self.get_query_set(instance, workspace_manage, is_x_pack_ee),
-                                  get_file_content(
-                                      os.path.join(PROJECT_DIR, "apps", "application", 'sql',
-                                                   'list_application.sql' if workspace_manage else (
-                                                       'list_application_user_ee.sql' if is_x_pack_ee else 'list_application_user.sql'))),
-                                  )
+        result = native_page_search(current_page, page_size,
+                                    self.get_query_set(instance, workspace_manage, is_x_pack_ee),
+                                    get_file_content(
+                                        os.path.join(PROJECT_DIR, "apps", "application", 'sql',
+                                                     'list_application.sql' if workspace_manage else (
+                                                         'list_application_user_ee.sql' if is_x_pack_ee else 'list_application_user.sql'))),
+                                    )
+
+        return ResourceMappingSerializer().get_resource_count(result)
 
 
 class ApplicationImportRequest(serializers.Serializer):
@@ -778,9 +783,13 @@ def get_mcp_servers(self, instance, with_valid=True):
     def delete(self, with_valid=True):
         if with_valid:
             self.is_valid()
-        QuerySet(ApplicationVersion).filter(application_id=self.data.get('application_id')).delete()
-        QuerySet(ApplicationKnowledgeMapping).filter(application_id=self.data.get('application_id')).delete()
-        QuerySet(Application).filter(id=self.data.get('application_id')).delete()
+        application_id = self.data.get('application_id')
+        QuerySet(ApplicationVersion).filter(application_id=application_id).delete()
+        QuerySet(ApplicationKnowledgeMapping).filter(application_id=application_id).delete()
+        QuerySet(ResourceMapping).filter(
+            Q(target_id=application_id) | Q(source_id=application_id)
+        ).delete()
+        QuerySet(Application).filter(id=application_id).delete()
         return True
 
     def export(self, with_valid=True):

apps/knowledge/serializers/knowledge.py

@@ -46,6 +46,8 @@
 from maxkb.conf import PROJECT_DIR
 from models_provider.models import Model
 from system_manage.models import WorkspaceUserResourcePermission, AuthTargetType
+from system_manage.models.resource_mapping import ResourceMapping
+from system_manage.serializers.resource_mapping_serializers import ResourceMappingSerializer
 from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer
 from users.serializers.user import is_workspace_manage
 
@@ -186,7 +188,7 @@ def page(self, current_page: int, page_size: int):
                 raise serializers.ValidationError(_('Folder not found'))
             workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
             is_x_pack_ee = self.is_x_pack_ee()
-            return native_page_search(
+            result = native_page_search(
                 current_page,
                 page_size,
                 self.get_query_set(workspace_manage, is_x_pack_ee),
@@ -202,6 +204,7 @@ def page(self, current_page: int, page_size: int):
                 ),
                 post_records_handler=lambda r: r
             )
+            return ResourceMappingSerializer().get_resource_count(result)
 
         def list(self):
             self.is_valid(raise_exception=True)
@@ -435,6 +438,9 @@ def delete(self):
             File.objects.filter(
                 source_id=knowledge.id,
             ).delete()
+            QuerySet(ResourceMapping).filter(
+                Q(target_id=knowledge) | Q(source_id=knowledge)
+            ).delete()
             delete_embedding_by_knowledge(self.data.get('knowledge_id'))
             return True
 

apps/models_provider/serializers/model_serializer.py

@@ -26,6 +26,8 @@
 from models_provider.models import Model, Status
 from models_provider.tools import get_model_credential
 from system_manage.models import WorkspaceUserResourcePermission, AuthTargetType
+from system_manage.models.resource_mapping import ResourceMapping
+from system_manage.serializers.resource_mapping_serializers import ResourceMappingSerializer
 from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer
 from users.serializers.user import is_workspace_manage
 
@@ -190,6 +192,7 @@ def delete(self, with_valid=True):
             #     if dataset_count > 0:
             #         raise AppApiException(500, f"该模型关联了{dataset_count} 个应用，无法删除该模型。")
             model.delete()
+            ResourceMapping.objects.filter(target_id=model_id).delete()
             return True
 
         def edit(self, instance: Dict, user_id: str, with_valid=True):
@@ -366,23 +369,25 @@ def list(self, workspace_id, with_valid):
             workspace_manage = is_workspace_manage(user_id, workspace_id)
             query_params = self._build_query_params(workspace_id, workspace_manage, user_id)
             is_x_pack_ee = self.is_x_pack_ee()
-            return native_search(query_params,
-                                 select_string=get_file_content(
-                                     os.path.join(PROJECT_DIR, "apps", "models_provider", 'sql',
-                                                  'list_model.sql' if workspace_manage else (
-                                                      'list_model_user_ee.sql' if is_x_pack_ee else 'list_model_user.sql')
-                                                  )))
+            result = native_search(query_params,
+                                   select_string=get_file_content(
+                                       os.path.join(PROJECT_DIR, "apps", "models_provider", 'sql',
+                                                    'list_model.sql' if workspace_manage else (
+                                                        'list_model_user_ee.sql' if is_x_pack_ee else 'list_model_user.sql')
+                                                    )))
+            return ResourceMappingSerializer().get_resource_count(result)
 
         def share_list(self, workspace_id, with_valid=True):
             if with_valid:
                 self.is_valid(raise_exception=True)
             user_id = self.data.get("user_id")
             query_params = self._build_query_params(workspace_id, False, user_id)
-            return [
+            result = [
                 self._build_model_data(
                     model
                 ) for model in query_params.get('model_query_set')
             ]
+            return ResourceMappingSerializer().get_resource_count(result)
 
         def model_list(self, workspace_id, with_valid=True):
             if with_valid:

apps/system_manage/serializers/resource_mapping_serializers.py

@@ -14,8 +14,10 @@
 from rest_framework import serializers
 
 from common.db.search import native_page_search, get_dynamics_model
+from common.result import Page
 from common.utils.common import get_file_content
 from maxkb.conf import PROJECT_DIR
+from system_manage.models.resource_mapping import ResourceMapping
 
 
 class ResourceMappingSerializer(serializers.Serializer):
@@ -51,3 +53,37 @@ def page(self, current_page, page_size):
         return native_page_search(current_page, page_size, self.get_query_set(), get_file_content(
             os.path.join(PROJECT_DIR, "apps", "system_manage",
                          'sql', 'list_resource_mapping.sql')), with_table_name=False)
+
+    def get_resource_count(self, result_list):
+        """
+        获取资源映射计数
+        """
+        if not result_list:
+            return result_list
+        is_paginated = isinstance(result_list, Page)
+
+        data_to_process = result_list.get('records') if is_paginated else result_list
+
+        if isinstance(data_to_process, list) and data_to_process:
+            # 提取ID列表，确保每个项目都是字典且包含'id'键
+            ids = [item['id'] for item in data_to_process
+                   if isinstance(item, dict) and 'id' in item and item['id']]
+
+            if ids:  # 只有在ids非空时才执行查询
+                mapping_counts = ResourceMapping.objects.filter(
+                    target_id__in=ids
+                ).values('target_id').annotate(
+                    count=models.Count('id')
+                )
+
+                # 构建目标ID到计数的映射
+                count_dict = {str(item['target_id']): item['count'] for item in mapping_counts}
+
+                # 为每个结果项添加资源计数
+                for model in data_to_process:
+                    if isinstance(model, dict) and 'id' in model:
+                        model_id = str(model['id'])
+                        model['resource_count'] = count_dict.get(model_id, 0)
+
+        return result_list
+

apps/tools/serializers/tool.py

@@ -31,6 +31,8 @@
 from common.utils.logger import maxkb_logger
 from common.utils.rsa_util import rsa_long_decrypt, rsa_long_encrypt
 from common.utils.tool_code import ToolExecutor
+from system_manage.models.resource_mapping import ResourceMapping
+from system_manage.serializers.resource_mapping_serializers import ResourceMappingSerializer
 from knowledge.models import File, FileSourceType
 from maxkb.const import CONFIG, PROJECT_DIR
 from system_manage.models import AuthTargetType, WorkspaceUserResourcePermission
@@ -523,6 +525,7 @@ def delete(self):
                 QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
             QuerySet(WorkspaceUserResourcePermission).filter(target=tool.id).delete()
             QuerySet(Tool).filter(id=self.data.get('id')).delete()
+            ResourceMapping.objects.filter(target_id=self.data.get('id')).delete()
 
         def one(self):
             self.is_one_valid(raise_exception=True)
@@ -964,8 +967,7 @@ def page_tool_with_folders(self, current_page: int, page_size: int):
 
             workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
             is_x_pack_ee = self.is_x_pack_ee()
-
-            return native_page_search(
+            result = native_page_search(
                 current_page, page_size, self.get_query_set(workspace_manage, is_x_pack_ee),
                 get_file_content(
                     os.path.join(
@@ -982,6 +984,7 @@ def page_tool_with_folders(self, current_page: int, page_size: int):
                     'init_field_list': json.loads(record.get('init_field_list', '[]')),
                 },
             )
+            return ResourceMappingSerializer().get_resource_count(result)
 
         def get_tools(self):
             self.is_valid(raise_exception=True)