feat: add search document node functionality and related configurations

Author: liuruibin

apps/application/flow/step_node/__init__.py

@@ -23,6 +23,7 @@
 from .mcp_node import BaseMcpNode
 from .question_node import *
 from .reranker_node import *
+from .search_document_node import BaseSearchDocumentNode
 from .search_knowledge_node import *
 from .speech_to_text_step_node import BaseSpeechToTextNode
 from .start_node import *
@@ -34,7 +35,7 @@
 from .variable_splitting_node import BaseVariableSplittingNode
 from .video_understand_step_node import BaseVideoUnderstandNode
 
-node_list = [BaseStartStepNode, BaseChatNode, BaseSearchKnowledgeNode, BaseQuestionNode,
+node_list = [BaseStartStepNode, BaseChatNode, BaseSearchKnowledgeNode, BaseSearchDocumentNode, BaseQuestionNode,
              BaseConditionNode, BaseReplyNode,
              BaseToolNodeNode, BaseToolLibNodeNode, BaseRerankerNode, BaseApplicationNode,
              BaseDocumentExtractNode,

apps/application/flow/step_node/search_document_node/__init__.py

@@ -0,0 +1 @@
+from .impl import *

apps/application/flow/step_node/search_document_node/i_search_document_node.py

@@ -0,0 +1,56 @@
+# coding=utf-8
+from typing import Type, List
+
+from django.utils.translation import gettext_lazy as _
+from rest_framework import serializers
+
+from application.flow.i_step_node import INode, NodeResult
+
+
+class SearchDocumentStepNodeSerializer(serializers.Serializer):
+    knowledge_id_list = serializers.ListField(
+        required=False, child=serializers.UUIDField(required=True),
+        label=_("knowledge id list"), default=list
+    )
+    search_mode = serializers.ChoiceField(
+        required=False, choices=['auto', 'custom'], label=_("search mode"), default='auto'
+    )
+    search_scope_type = serializers.ChoiceField(
+        required=False, choices=['custom', 'referencing'], label=_("search scope type"),
+        allow_null=True, default='custom'
+    )
+    search_scope_source = serializers.ChoiceField(
+        required=False, choices=['document', 'knowledge'],
+        label=_("search scope variable type"), default='knowledge'
+    )
+    search_scope_reference = serializers.ListField(
+        required=False, label=_("search scope variable"), default=list
+    )
+    question_reference = serializers.ListField(
+        required=False, label=_("question reference address"), default=list
+    )
+    search_condition_type = serializers.ChoiceField(
+        required=False, choices=['AND', 'OR'], label=_("search condition type"), default='AND'
+    )
+    search_condition_list = serializers.ListField(
+        required=False, label=_("search condition list"), default=list
+    )
+
+    def is_valid(self, *, raise_exception=False):
+        super().is_valid(raise_exception=True)
+
+
+class ISearchDocumentStepNode(INode):
+    type = 'search-document-node'
+
+    def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
+        return SearchDocumentStepNodeSerializer
+
+    def _run(self):
+        return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data)
+
+    def execute(self, knowledge_id_list: List, search_mode: str, search_scope_type: str, search_scope_source: str,
+                search_scope_reference: List, question_reference: List, search_condition_type: str,
+                search_condition_list: List,
+                **kwargs) -> NodeResult:
+        pass

apps/application/flow/step_node/search_document_node/impl/__init__.py

@@ -0,0 +1 @@
+from .base_search_document_node import BaseSearchDocumentNode

apps/application/flow/step_node/search_document_node/impl/base_search_document_node.py

@@ -0,0 +1,159 @@
+# coding=utf-8
+from typing import List
+
+import jieba
+from django.db.models import Q
+from django.db.models import QuerySet
+
+from application.flow.i_step_node import NodeResult
+from application.flow.step_node.search_document_node.i_search_document_node import ISearchDocumentStepNode
+from knowledge.models import Document, DocumentTag, Knowledge
+
+
+class BaseSearchDocumentNode(ISearchDocumentStepNode):
+    def save_context(self, details, workflow_manage):
+        self.context['document_list'] = details.get('document_list')
+        self.context['knowledge_list'] = details.get('knowledge_list')
+        self.context['document_items'] = details.get('document_items')
+        self.context['knowledge_items'] = details.get('knowledge_items')
+        self.context['question'] = details.get('question')
+        self.context['run_time'] = details.get('run_time')
+
+    def get_reference_content(self, fields: List[str]):
+        return self.workflow_manage.get_reference_field(fields[0], fields[1:])
+
+    def execute(self, knowledge_id_list: List, search_mode: str, search_scope_type: str, search_scope_source: str,
+                search_scope_reference: List, question_reference: List, search_condition_type: str,
+                search_condition_list: List,
+                **kwargs) -> NodeResult:
+
+        if search_scope_type == 'custom':  # 手动选择知识库
+            document_id_list = QuerySet(Document).filter(
+                knowledge_id__in=knowledge_id_list
+            ).values_list('id', flat=True)
+        else:  # 引用上一步知识库/文档
+            if search_scope_source == 'document':  # 文档
+                document_id_list = self.get_reference_content(search_scope_reference)
+            else:  # 知识库
+                document_id_list = QuerySet(Document).filter(
+                    knowledge_id__in=self.get_reference_content(search_scope_reference)
+                ).values_list('id', flat=True)
+
+        if search_mode == 'auto':  # 通过问题自动检索
+            matched_doc_ids = self.handle_auto_tags(document_id_list, question_reference)
+
+            final_document_ids = list(matched_doc_ids)
+        else:  # 自定义检索条件
+            matched_document_ids = self.handle_custom_tags(
+                document_id_list, search_condition_list, search_condition_type
+            )
+
+            final_document_ids = list(matched_document_ids)
+
+        # UUID to str
+        final_document_ids = [str(doc_id) for doc_id in final_document_ids]
+        document_items = QuerySet(Document).filter(id__in=final_document_ids).values()
+        final_knowledge_ids = list(set(str(doc['knowledge_id']) for doc in document_items))
+        knowledge_items = QuerySet(Knowledge).filter(id__in=final_knowledge_ids).values()
+
+        return NodeResult({
+            'document_list': final_document_ids,
+            'document_items': list(document_items),
+            'knowledge_list': final_knowledge_ids,
+            'knowledge_items': list(knowledge_items)
+        }, {})
+
+    def handle_auto_tags(self, document_id_list: list, question_reference: list):
+        question = self.get_reference_content(question_reference)
+
+        # 使用jieba分词
+        keywords = jieba.lcut(question)
+        if not keywords:
+            return set()
+
+        # 构建OR查询,一次性获取所有匹配的文档
+        q_objects = Q()
+        for keyword in keywords:
+            q_objects |= Q(tag__value__icontains=keyword)
+
+        # 单次数据库查询
+        matched_doc_ids = set(
+            QuerySet(DocumentTag)
+            .filter(document_id__in=document_id_list)
+            .filter(q_objects)
+            .values_list('document_id', flat=True)
+            .distinct()
+        )
+
+        return matched_doc_ids
+
+    def handle_custom_tags(self, document_id_list: List, search_condition_list: list, search_condition_type: str):
+
+        if not search_condition_list:
+            return set(document_id_list)
+
+        if search_condition_type == 'AND':
+            # AND逻辑:使用子查询和聚合
+            matched_doc_ids = set(document_id_list)
+
+            for condition in search_condition_list:
+                tag_key = condition['key']
+                field_value = self.workflow_manage.generate_prompt(condition['value'])
+                compare_type = condition['compare']
+
+                # 构建查询条件
+                if compare_type == 'contain':
+                    q_filter = Q(tag__key=tag_key, tag__value__icontains=field_value)
+                elif compare_type == 'eq':
+                    q_filter = Q(tag__key=tag_key, tag__value=field_value)
+                elif compare_type == 'not_contain':
+                    q_filter = ~Q(tag__key=tag_key, tag__value__icontains=field_value)
+                else:
+                    continue
+
+                # 单次查询获取符合条件的文档
+                tag_docs = set(QuerySet(DocumentTag).filter(
+                    document_id__in=matched_doc_ids
+                ).filter(q_filter).values_list('document_id', flat=True).distinct())
+
+                matched_doc_ids = matched_doc_ids.intersection(tag_docs)
+
+            return matched_doc_ids
+
+        else:
+            # OR逻辑:使用一次查询完成
+            q_objects = Q()
+
+            for condition in search_condition_list:
+                tag_key = condition['key']
+                field_value = self.workflow_manage.generate_prompt(condition['value'])
+                compare_type = condition['compare']
+
+                if compare_type == 'contain':
+                    q_objects |= Q(tag__key=tag_key, tag__value__icontains=field_value)
+                elif compare_type == 'eq':
+                    q_objects |= Q(tag__key=tag_key, tag__value=field_value)
+                elif compare_type == 'not_contain':
+                    q_objects |= ~Q(tag__key=tag_key, tag__value__icontains=field_value)
+
+            # 一次查询获取所有匹配的文档
+            matched_docs = set(QuerySet(DocumentTag).filter(
+                document_id__in=document_id_list
+            ).filter(q_objects).values_list('document_id', flat=True).distinct())
+
+            return matched_docs
+
+    def get_details(self, index: int, **kwargs):
+        return {
+            'name': self.node.properties.get('stepName'),
+            'question': self.context.get('question'),
+            "index": index,
+            'run_time': self.context.get('run_time'),
+            'document_list': self.context.get('document_list'),
+            'knowledge_list': self.context.get('knowledge_list'),
+            'document_items': self.context.get('document_items'),
+            'knowledge_items': self.context.get('knowledge_items'),
+            'type': self.node.type,
+            'status': self.status,
+            'err_message': self.err_message
+        }

apps/knowledge/serializers/knowledge.py

@@ -3,6 +3,7 @@
 import os
 import re
 import traceback
+from collections import defaultdict
 from functools import reduce
 from tempfile import TemporaryDirectory
 from typing import Dict, List
@@ -13,6 +14,7 @@
 from django.db import transaction, models
 from django.db.models import QuerySet
 from django.db.models.functions import Reverse, Substr
+from django.db.models.query_utils import Q
 from django.http import HttpResponse
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
@@ -29,8 +31,9 @@
 from common.utils.logger import maxkb_logger
 from common.utils.split_model import get_split_model
 from knowledge.models import Knowledge, KnowledgeScope, KnowledgeType, Document, Paragraph, Problem, \
-    ProblemParagraphMapping, TaskType, State, SearchMode, KnowledgeFolder, File
-from knowledge.serializers.common import ProblemParagraphManage, drop_knowledge_index, get_embedding_model_id_by_knowledge_id, MetaSerializer, \
+    ProblemParagraphMapping, TaskType, State, SearchMode, KnowledgeFolder, File, Tag
+from knowledge.serializers.common import ProblemParagraphManage, drop_knowledge_index, \
+    get_embedding_model_id_by_knowledge_id, MetaSerializer, \
     GenerateRelatedSerializer, get_embedding_model_by_knowledge_id, list_paragraph, write_image, zip_dir
 from knowledge.serializers.document import DocumentSerializers
 from knowledge.task.embedding import embedding_by_knowledge, delete_embedding_by_knowledge
@@ -148,7 +151,8 @@ def get_query_set(self, workspace_manage, is_x_pack_ee):
             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 and self.data.get('workspace_id') != self.data.get('folder_id'):
+            if "folder_id" in self.data and self.data.get('folder_id') is not None and self.data.get(
+                    'workspace_id') != self.data.get('folder_id'):
                 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")})
             if "scope" in self.data and self.data.get('scope') is not None:
@@ -764,3 +768,50 @@ def hit_test(self):
                     'comprehensive_score': hit_dict.get(p.get('id')).get('comprehensive_score')
                 } for p in p_list
             ]
+
+    class Tags(serializers.Serializer):
+        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
+        user_id = serializers.UUIDField(required=True, label=_('user id'))
+        knowledge_ids = serializers.ListField(
+            required=True, label=_('knowledge ids'),
+            child=serializers.UUIDField(required=True, label=_('id'))
+        )
+
+        def list(self):
+            self.is_valid(raise_exception=True)
+            if self.data.get('name'):
+                name = self.data.get('name')
+                tags = QuerySet(Tag).filter(
+                    knowledge_id__in=self.data.get('knowledge_ids')
+                ).filter(
+                    Q(key__icontains=name) | Q(value__icontains=name)
+                ).values('key', 'value', 'id', 'create_time', 'update_time').order_by('create_time', 'key', 'value')
+            else:
+                # 获取所有标签，按创建时间排序保持稳定顺序
+                tags = QuerySet(Tag).filter(
+                    knowledge_id__in=self.data.get('knowledge_ids')
+                ).values('key', 'value', 'id', 'create_time', 'update_time').order_by('create_time', 'key', 'value')
+
+            # 按key分组
+            grouped_tags = defaultdict(list)
+            for tag in tags:
+                grouped_tags[tag['key']].append({
+                    'id': tag['id'],
+                    'value': tag['value'],
+                    'create_time': tag['create_time'],
+                    'update_time': tag['update_time']
+                })
+
+            # 转换为期望的格式，保持key的顺序
+            result = []
+            # 按key排序以确保结果顺序一致
+            for key in sorted(grouped_tags.keys()):
+                values = grouped_tags[key]
+                # 按创建时间对values进行排序
+                values.sort(key=lambda x: x['create_time'])
+                result.append({
+                    'key': key,
+                    'values': values,
+                })
+
+            return result

apps/knowledge/urls.py

@@ -12,6 +12,7 @@
     path('workspace/<str:workspace_id>/knowledge/web', views.KnowledgeWebView.as_view()),
     path('workspace/<str:workspace_id>/knowledge/model', views.KnowledgeView.Model.as_view()),
     path('workspace/<str:workspace_id>/knowledge/embedding_model', views.KnowledgeView.EmbeddingModel.as_view()),
+    path('workspace/<str:workspace_id>/knowledge/tags', views.KnowledgeView.Tags.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>', views.KnowledgeView.Operate.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/sync', views.KnowledgeView.SyncWeb.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/generate_related', views.KnowledgeView.GenerateRelated.as_view()),

apps/knowledge/views/knowledge.py

@@ -382,6 +382,29 @@ def get(self, request: Request, workspace_id: str):
                 }
             ).list(workspace_id, True))
 
+    class Tags(APIView):
+        authentication_classes = [TokenAuth]
+
+        @extend_schema(
+            methods=['GET'],
+            description=_('Get all tags of knowledge base'),
+            summary=_('Get all tags of knowledge base'),
+            operation_id=_('Get all tags of knowledge base'),  # type: ignore
+            parameters=KnowledgeReadAPI.get_parameters(),
+            responses=KnowledgeReadAPI.get_response(),
+            tags=[_('Knowledge Base')]  # type: ignore
+        )
+        @has_permissions(
+            PermissionConstants.KNOWLEDGE_READ.get_workspace_permission(),
+            RoleConstants.WORKSPACE_MANAGE.get_workspace_role(), RoleConstants.USER.get_workspace_role()
+        )
+        def get(self, request: Request, workspace_id: str):
+            return result.success(KnowledgeSerializer.Tags(data={
+                'user_id': request.user.id,
+                'workspace_id': workspace_id,
+                'knowledge_ids': request.query_params.getlist('knowledge_ids[]')
+            }).list())
+
 
 class KnowledgeBaseView(APIView):
     authentication_classes = [TokenAuth]

ui/src/api/knowledge/knowledge.ts

@@ -255,6 +255,12 @@ const putLarkKnowledge: (
   return put(`${prefix.value}/lark/${knowledge_id}`, data, undefined, loading)
 }
 
+const getAllTags: (params: any, loading?: Ref<boolean>) => Promise<Result<any>> = (
+  params,
+  loading,
+) => {
+  return get(`${prefix.value}/tags`, params, loading)
+}
 
 const getTags: (knowledge_id: string, params: any, loading?: Ref<boolean>) => Promise<Result<any>> = (
   knowledge_id,
@@ -315,6 +321,7 @@ export default {
   postWebKnowledge,
   postLarkKnowledge,
   putLarkKnowledge,
+  getAllTags,
   getTags,
   postTags,
   putTag,