Merge branch 'v2' into pr@v2@feat_parameter_node

Author: shaohuzhang1

apps/application/flow/step_node/__init__.py

@@ -24,6 +24,7 @@
 from .parameter_extraction_node import BaseParameterExtractionNode
 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 *
@@ -35,7 +36,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/image_understand_step_node/impl/base_image_understand_node.py

@@ -77,8 +77,6 @@ def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, hist
                 image,
                 **kwargs) -> NodeResult:
         # 处理不正确的参数
-        if image is None or not isinstance(image, list):
-            image = []
         workspace_id = self.workflow_manage.get_body().get('workspace_id')
         image_model = get_model_instance_by_model_workspace_id(model_id, workspace_id,
                                                                **model_params_setting)
@@ -91,7 +89,7 @@ def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, hist
         message_list = self.generate_message_list(image_model, system, prompt,
                                                   self.get_history_message(history_chat_record, dialogue_number), image)
         self.context['message_list'] = message_list
-        self.context['image_list'] = image
+        self.generate_context_image(image)
         self.context['dialogue_type'] = dialogue_type
         if stream:
             r = image_model.stream(message_list)
@@ -104,6 +102,12 @@ def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, hist
                                'history_message': history_message, 'question': question.content}, {},
                               _write_context=write_context)
 
+    def generate_context_image(self, image):
+        if isinstance(image, str) and image.startswith('http'):
+            self.context['image_list'] = [{'url': image}]
+        elif image is not None and len(image) > 0:
+            self.context['image_list'] = image
+
     def get_history_message_for_details(self, history_chat_record, dialogue_number):
         start_index = len(history_chat_record) - dialogue_number
         history_message = reduce(lambda x, y: [*x, *y], [
@@ -164,28 +168,32 @@ def generate_history_human_message(self, chat_record):
     def generate_prompt_question(self, prompt):
         return HumanMessage(self.workflow_manage.generate_prompt(prompt))
 
-    def generate_message_list(self, image_model, system: str, prompt: str, history_message, image):
-        if image is not None and len(image) > 0:
-            # 处理多张图片
-            images = []
+    def _process_images(self, image):
+        """
+        处理图像数据，转换为模型可识别的格式
+        """
+        images = []
+        if isinstance(image, str) and image.startswith('http'):
+            images.append({'type': 'image_url', 'image_url': {'url': image}})
+        elif image is not None and len(image) > 0:
             for img in image:
-                if isinstance(img, str) and img.startswith('http'):
-                    images.append({'type': 'image_url', 'image_url': {'url': img}})
-                else:
-                    file_id = img['file_id']
-                    file = QuerySet(File).filter(id=file_id).first()
-                    image_bytes = file.get_bytes()
-                    base64_image = base64.b64encode(image_bytes).decode("utf-8")
-                    image_format = what(None, image_bytes)
-                    images.append(
-                        {'type': 'image_url', 'image_url': {'url': f'data:image/{image_format};base64,{base64_image}'}})
-            messages = [HumanMessage(
-                content=[
-                    {'type': 'text', 'text': self.workflow_manage.generate_prompt(prompt)},
-                    *images
-                ])]
+                file_id = img['file_id']
+                file = QuerySet(File).filter(id=file_id).first()
+                image_bytes = file.get_bytes()
+                base64_image = base64.b64encode(image_bytes).decode("utf-8")
+                image_format = what(None, image_bytes)
+                images.append(
+                    {'type': 'image_url', 'image_url': {'url': f'data:image/{image_format};base64,{base64_image}'}})
+        return images
+
+    def generate_message_list(self, image_model, system: str, prompt: str, history_message, image):
+        prompt_text = self.workflow_manage.generate_prompt(prompt)
+        images = self._process_images(image)
+
+        if images:
+            messages = [HumanMessage(content=[{'type': 'text', 'text': prompt_text}, *images])]
         else:
-            messages = [HumanMessage(self.workflow_manage.generate_prompt(prompt))]
+            messages = [HumanMessage(prompt_text)]
 
         if system is not None and len(system) > 0:
             return [

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,178 @@
+# 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 common.constants.permission_constants import RoleConstants
+from common.database_model_manage.database_model_manage import DatabaseModelManage
+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)
+
+        # 权限过滤
+        get_knowledge_list_of_authorized = DatabaseModelManage.get_model('get_knowledge_list_of_authorized')
+        chat_user_type = self.workflow_manage.get_body().get('chat_user_type')
+
+        if get_knowledge_list_of_authorized is not None and RoleConstants.CHAT_USER.value.name == chat_user_type:
+            # 获取授权的知识库ID列表
+            authorized_knowledge_ids = get_knowledge_list_of_authorized(
+                self.workflow_manage.get_body().get('chat_user_id'),
+                knowledge_id_list
+            )
+
+            # 过滤出授权知识库下的文档
+            document_id_list = QuerySet(Document).filter(
+                id__in=document_id_list,
+                knowledge_id__in=authorized_knowledge_ids
+            ).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/application/flow/step_node/search_knowledge_node/i_search_knowledge_node.py

@@ -10,13 +10,12 @@
 from typing import Type
 
 from django.core import validators
+from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
 
 from application.flow.i_step_node import INode, NodeResult
 from common.utils.common import flat_map
 
-from django.utils.translation import gettext_lazy as _
-
 
 class DatasetSettingSerializer(serializers.Serializer):
     # 需要查询的条数
@@ -43,6 +42,17 @@ class SearchDatasetStepNodeSerializer(serializers.Serializer):
 
     show_knowledge = serializers.BooleanField(required=True,
                                               label=_("The results are displayed in the knowledge sources"))
+    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
+    )
 
     def is_valid(self, *, raise_exception=False):
         super().is_valid(raise_exception=True)
@@ -76,7 +86,9 @@ def _run(self):
         return self.execute(**self.node_params_serializer.data, question=str(question),
                             exclude_paragraph_id_list=exclude_paragraph_id_list)
 
-    def execute(self, dataset_id_list, dataset_setting, question, show_knowledge,
+    def execute(self, dataset_id_list, dataset_setting, question, show_knowledge, search_scope_type,
+                search_scope_source,
+                search_scope_reference,
                 exclude_paragraph_id_list=None,
                 **kwargs) -> NodeResult:
         pass