From aa593f3790bf88ae598f5266e368fe7ee2f7c275 Mon Sep 17 00:00:00 2001 From: CaptainB Date: Tue, 3 Jun 2025 15:53:56 +0800 Subject: [PATCH] feat: refactor user_id handling and enhance tool query logic --- apps/knowledge/api/knowledge.py | 7 --- apps/knowledge/serializers/knowledge.py | 32 ++++++------ apps/knowledge/sql/list_knowledge.sql | 66 +++++++++++++++++-------- apps/knowledge/views/knowledge.py | 2 +- apps/tools/serializers/tool.py | 65 +++++++++++++++--------- apps/tools/sql/list_tool.sql | 25 ++++++++++ 6 files changed, 130 insertions(+), 67 deletions(-) create mode 100644 apps/tools/sql/list_tool.sql diff --git a/apps/knowledge/api/knowledge.py b/apps/knowledge/api/knowledge.py index 1fa2c07d6..d194e9690 100644 --- a/apps/knowledge/api/knowledge.py +++ b/apps/knowledge/api/knowledge.py @@ -185,13 +185,6 @@ class KnowledgePageAPI(KnowledgeReadAPI): location='query', required=True, ), - OpenApiParameter( - name="user_id", - description="用户id", - type=OpenApiTypes.STR, - location='query', - required=False, - ), OpenApiParameter( name="name", description="名称", diff --git a/apps/knowledge/serializers/knowledge.py b/apps/knowledge/serializers/knowledge.py index 62c316516..5190d5e99 100644 --- a/apps/knowledge/serializers/knowledge.py +++ b/apps/knowledge/serializers/knowledge.py @@ -28,7 +28,6 @@ from knowledge.models import Knowledge, KnowledgeScope, KnowledgeType, Document, from knowledge.serializers.common import ProblemParagraphManage, get_embedding_model_id_by_knowledge_id, MetaSerializer, \ GenerateRelatedSerializer, get_embedding_model_by_knowledge_id, list_paragraph from knowledge.serializers.document import DocumentSerializers -from knowledge.serializers.knowledge_folder import KnowledgeFolderFlatSerializer from knowledge.task.embedding import embedding_by_knowledge, delete_embedding_by_knowledge from knowledge.task.generate import generate_related_by_knowledge_id from knowledge.task.sync import sync_web_knowledge, sync_replace_web_knowledge @@ -117,23 +116,30 @@ class KnowledgeSerializer(serializers.Serializer): 'temp.folder_id': models.CharField(), 'temp.id': models.CharField() })) + folder_query_set = QuerySet(KnowledgeFolder) + if "desc" in self.data and self.data.get('desc') is not None: query_set = query_set.filter(**{'temp.desc__icontains': self.data.get("desc")}) + folder_query_set = folder_query_set.filter(**{'desc__icontains': self.data.get("desc")}) if "name" in self.data and self.data.get('name') is not None: query_set = query_set.filter(**{'temp.name__icontains': self.data.get("name")}) + folder_query_set = folder_query_set.filter(**{'name__icontains': self.data.get("name")}) if "user_id" in self.data and self.data.get('user_id') is not None: query_set = query_set.filter(**{'temp.user_id': self.data.get("user_id")}) + folder_query_set = folder_query_set.filter(**{'user_id': self.data.get("user_id")}) if "workspace_id" in self.data and self.data.get('workspace_id') is not None: query_set = query_set.filter(**{'temp.workspace_id': self.data.get("workspace_id")}) + folder_query_set = folder_query_set.filter(**{'workspace_id': self.data.get("workspace_id")}) if "folder_id" in self.data and self.data.get('folder_id') is not None: query_set = query_set.filter(**{'temp.folder_id': self.data.get("folder_id")}) + folder_query_set = folder_query_set.filter(**{'parent_id': self.data.get("folder_id")}) query_set = query_set.order_by("-temp.create_time", "temp.id") query_set_dict['default_sql'] = query_set query_set_dict['knowledge_custom_sql'] = QuerySet(model=get_dynamics_model({ 'knowledge.workspace_id': models.CharField(), })).filter(**{'knowledge.workspace_id': workspace_id}) - + query_set_dict['folder_query_set'] = folder_query_set return query_set_dict def page(self, current_page: int, page_size: int): @@ -143,22 +149,16 @@ class KnowledgeSerializer(serializers.Serializer): root = KnowledgeFolder.objects.filter(id=folder_id).first() if not root: raise serializers.ValidationError(_('Folder not found')) - # 获取当前文件夹下的直接子文件夹 - child_folders = KnowledgeFolder.objects.filter(parent=root) - folders_data = KnowledgeFolderFlatSerializer(child_folders, many=True).data - return { - 'knowledge': native_page_search( - current_page, - page_size, - self.get_query_set(), - select_string=get_file_content( - os.path.join(PROJECT_DIR, "apps", "knowledge", 'sql', 'list_knowledge.sql') - ), - post_records_handler=lambda r: r + return native_page_search( + current_page, + page_size, + self.get_query_set(), + select_string=get_file_content( + os.path.join(PROJECT_DIR, "apps", "knowledge", 'sql', 'list_knowledge.sql') ), - 'folders': folders_data - } + post_records_handler=lambda r: r + ) def list(self): self.is_valid(raise_exception=True) diff --git a/apps/knowledge/sql/list_knowledge.sql b/apps/knowledge/sql/list_knowledge.sql index d30b81849..146644f91 100644 --- a/apps/knowledge/sql/list_knowledge.sql +++ b/apps/knowledge/sql/list_knowledge.sql @@ -1,21 +1,45 @@ -SELECT - *, - to_json(meta) as meta -FROM - ( - SELECT - "temp_knowledge".*, - "document_temp"."char_length", - CASE - WHEN - "app_knowledge_temp"."count" IS NULL THEN 0 ELSE "app_knowledge_temp"."count" END AS application_mapping_count, - "document_temp".document_count FROM ( - SELECT knowledge.* - FROM - knowledge knowledge - ${knowledge_custom_sql} - ) temp_knowledge - LEFT JOIN ( SELECT "count" ( "id" ) AS document_count, "sum" ( "char_length" ) "char_length", knowledge_id FROM "document" GROUP BY knowledge_id ) "document_temp" ON temp_knowledge."id" = "document_temp".knowledge_id - LEFT JOIN (SELECT "count"("id"),knowledge_id FROM application_knowledge_mapping GROUP BY knowledge_id) app_knowledge_temp ON temp_knowledge."id" = "app_knowledge_temp".knowledge_id - ) temp - ${default_sql} \ No newline at end of file +SELECT * +FROM (SELECT "temp_knowledge".id::text, "temp_knowledge".name, + "temp_knowledge".desc, + "temp_knowledge".type, + 'knowledge' as resource_type, + "temp_knowledge".workspace_id, + "temp_knowledge".folder_id, + "temp_knowledge".user_id, + "temp_knowledge".create_time, + "temp_knowledge".update_time, + "document_temp"."char_length", + CASE + WHEN + "app_knowledge_temp"."count" IS NULL THEN 0 + ELSE "app_knowledge_temp"."count" END AS application_mapping_count, + "document_temp".document_count + FROM (SELECT knowledge.* + FROM knowledge knowledge ${knowledge_custom_sql} + AND id in (SELECT target + FROM workspace_user_resource_permission + WHERE auth_target_type = 'KNOWLEDGE' + AND 'VIEW' = any (permission_list))) temp_knowledge + LEFT JOIN (SELECT "count"("id") AS document_count, "sum"("char_length") "char_length", knowledge_id + FROM "document" + GROUP BY knowledge_id) "document_temp" ON temp_knowledge."id" = "document_temp".knowledge_id + LEFT JOIN (SELECT "count"("id"), knowledge_id + FROM application_knowledge_mapping + GROUP BY knowledge_id) app_knowledge_temp + ON temp_knowledge."id" = "app_knowledge_temp".knowledge_id + UNION + SELECT "id", + "name", + "desc", + 0 as "type", + 'folder' as "resource_type", + "workspace_id", + "parent_id" as "folder_id", + "user_id", + "create_time", + "update_time", + 0 as char_length, + 0 as application_mapping_count, + 0 as document_count + from knowledge_folder ${folder_query_set}) temp + ${default_sql} \ No newline at end of file diff --git a/apps/knowledge/views/knowledge.py b/apps/knowledge/views/knowledge.py index 49a1f5411..8b60294d7 100644 --- a/apps/knowledge/views/knowledge.py +++ b/apps/knowledge/views/knowledge.py @@ -108,7 +108,7 @@ class KnowledgeView(APIView): 'folder_id': request.query_params.get('folder_id'), 'name': request.query_params.get('name'), 'desc': request.query_params.get("desc"), - 'user_id': request.query_params.get('user_id') + 'user_id': request.user.id } ).page(current_page, page_size)) diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index 9ce56de0a..1e36ecf86 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -15,9 +15,10 @@ from pylint.lint import Run from pylint.reporters import JSON2Reporter from rest_framework import serializers, status -from common.db.search import page_search +from common.db.search import page_search, native_page_search from common.exception.app_exception import AppApiException from common.result import result +from common.utils.common import get_file_content from common.utils.tool_code import ToolExecutor from maxkb.const import CONFIG, PROJECT_DIR from tools.models import Tool, ToolScope, ToolFolder @@ -418,27 +419,47 @@ class ToolTreeSerializer(serializers.Serializer): Q(tool_type=self.data.get('tool_type'))) return page_search(current_page, page_size, tools, lambda record: ToolModelSerializer(record).data) + def get_query_set(self): + tool_query_set = QuerySet(Tool) + tool_type_query_set = QuerySet(Tool) + folder_query_set = QuerySet(ToolFolder) + workspace_id = self.data.get('workspace_id') + user_id = self.data.get('user_id') + tool_type = self.data.get('tool_type') + desc = self.data.get('desc') + name = self.data.get('name') + folder_id = self.data.get('folder_id') + + if workspace_id is not None: + folder_query_set = folder_query_set.filter(workspace_id=workspace_id) + tool_query_set = tool_query_set.filter(workspace_id=workspace_id) + if user_id is not None: + folder_query_set = folder_query_set.filter(user_id=user_id) + tool_query_set = tool_query_set.filter(user_id=user_id) + if folder_id is not None: + folder_query_set = folder_query_set.filter(parent=folder_id) + tool_query_set = tool_query_set.filter(folder_id=folder_id) + if name is not None: + folder_query_set = folder_query_set.filter(name__contains=name) + tool_query_set = tool_query_set.filter(name__contains=name) + if desc is not None: + folder_query_set = folder_query_set.filter(desc__contains=desc) + tool_query_set = tool_query_set.filter(desc__contains=desc) + tool_query_set = tool_query_set.order_by("-update_time") + + if tool_type is not None: + tool_type_query_set = tool_type_query_set.filter(tool_type=tool_type) + + return { + 'folder_query_set': folder_query_set, + 'tool_query_set': tool_query_set, + 'tool_type_query_set': tool_type_query_set + } + def page_tool_with_folders(self, current_page: int, page_size: int): self.is_valid(raise_exception=True) - folder_id = self.data.get('folder_id', 'root') - root = ToolFolder.objects.filter(id=folder_id).first() - if not root: - raise serializers.ValidationError(_('Folder not found')) - # 获取当前文件夹下的直接子文件夹 - child_folders = ToolFolder.objects.filter(parent=root) - folders_data = ToolFolderFlatSerializer(child_folders, many=True).data - - if self.data.get('name'): - tools = QuerySet(Tool).filter(Q(workspace_id=self.data.get('workspace_id')) & - Q(folder_id=root) & - Q(tool_type=self.data.get('tool_type')) & - Q(name__contains=self.data.get('name'))) - else: - tools = QuerySet(Tool).filter(Q(workspace_id=self.data.get('workspace_id')) & - Q(folder_id=root) & - Q(tool_type=self.data.get('tool_type'))) - return { - 'tools': page_search(current_page, page_size, tools, lambda record: ToolModelSerializer(record).data), - 'folders': folders_data - } + return native_page_search( + current_page, page_size, self.get_query_set(), + get_file_content(os.path.join(PROJECT_DIR, "apps", "tools", 'sql', 'list_tool.sql')) + ) diff --git a/apps/tools/sql/list_tool.sql b/apps/tools/sql/list_tool.sql new file mode 100644 index 000000000..8b7160323 --- /dev/null +++ b/apps/tools/sql/list_tool.sql @@ -0,0 +1,25 @@ +select * +from (select "id"::text, + "name", + "desc", + "tool_type", + 'tool' as "resource_type", + "workspace_id", + "folder_id", + "user_id", + "create_time", + "update_time" + from tool ${tool_type_query_set} + UNION + select "id", + "name", + "desc", + 'folder' as "tool_type", + 'folder' as "resource_type", + "workspace_id", + "parent_id" as "folder_id", + "user_id", + "create_time", + "update_time" + from tool_folder ${folder_query_set}) temp + ${tool_query_set} \ No newline at end of file