From 2221eefff6799cb0ccdd35afdd72617c385aea54 Mon Sep 17 00:00:00 2001 From: wxg0103 <727495428@qq.com> Date: Mon, 29 Dec 2025 14:41:22 +0800 Subject: [PATCH] feat: integrate ResourceMappingSerializer for resource count retrieval across multiple modules --- apps/application/serializers/application.py | 27 +++++++++----- apps/knowledge/serializers/knowledge.py | 8 ++++- .../serializers/model_serializer.py | 19 ++++++---- .../resource_mapping_serializers.py | 36 +++++++++++++++++++ apps/tools/serializers/tool.py | 7 ++-- 5 files changed, 78 insertions(+), 19 deletions(-) diff --git a/apps/application/serializers/application.py b/apps/application/serializers/application.py index 137169384..404a6e37a 100644 --- a/apps/application/serializers/application.py +++ b/apps/application/serializers/application.py @@ -49,6 +49,8 @@ from maxkb.conf import PROJECT_DIR 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 @@ class Query(serializers.Serializer): 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 @@ class ApplicationOperateSerializer(serializers.Serializer): 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): diff --git a/apps/knowledge/serializers/knowledge.py b/apps/knowledge/serializers/knowledge.py index 070676822..00614b9a5 100644 --- a/apps/knowledge/serializers/knowledge.py +++ b/apps/knowledge/serializers/knowledge.py @@ -46,6 +46,8 @@ from knowledge.task.sync import sync_web_knowledge, sync_replace_web_knowledge 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 @@ class KnowledgeSerializer(serializers.Serializer): 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 @@ class KnowledgeSerializer(serializers.Serializer): ), 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 @@ class KnowledgeSerializer(serializers.Serializer): 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 diff --git a/apps/models_provider/serializers/model_serializer.py b/apps/models_provider/serializers/model_serializer.py index adaafb2a7..e3512e325 100644 --- a/apps/models_provider/serializers/model_serializer.py +++ b/apps/models_provider/serializers/model_serializer.py @@ -26,6 +26,8 @@ from models_provider.constants.model_provider_constants import ModelProvideConst 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 @@ class ModelSerializer(serializers.Serializer): # 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 @@ class ModelSerializer(serializers.Serializer): 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: diff --git a/apps/system_manage/serializers/resource_mapping_serializers.py b/apps/system_manage/serializers/resource_mapping_serializers.py index 2ee4634b5..9f600de65 100644 --- a/apps/system_manage/serializers/resource_mapping_serializers.py +++ b/apps/system_manage/serializers/resource_mapping_serializers.py @@ -14,8 +14,10 @@ from django.utils.translation import gettext_lazy as _ 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 @@ class ResourceMappingSerializer(serializers.Serializer): 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 + diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index cac977941..e39e7476b 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -31,6 +31,8 @@ from common.utils.common import get_file_content 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 @@ class ToolSerializer(serializers.Serializer): 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 @@ class ToolTreeSerializer(serializers.Serializer): 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 @@ class ToolTreeSerializer(serializers.Serializer): '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)