diff --git a/apps/application/workflow/nodes/reranker_node/__init__.py b/apps/application/workflow/nodes/reranker_node/__init__.py new file mode 100644 index 00000000000..e6adc8ee183 --- /dev/null +++ b/apps/application/workflow/nodes/reranker_node/__init__.py @@ -0,0 +1,7 @@ +# coding=utf-8 +""" + @project: MaxKB + @file： __init__.py + @desc: +""" +from .reranker_node import RerankerNode diff --git a/apps/application/workflow/nodes/reranker_node/reranker_node.py b/apps/application/workflow/nodes/reranker_node/reranker_node.py new file mode 100644 index 00000000000..549f28569f9 --- /dev/null +++ b/apps/application/workflow/nodes/reranker_node/reranker_node.py @@ -0,0 +1,157 @@ +# coding=utf-8 +""" + @project: MaxKB + @file： reranker_node.py + @desc: +""" +from typing import List + +from django.utils.translation import gettext_lazy as _ +from langchain_core.documents import Document +from rest_framework import serializers + +from application.workflow.common import WorkflowType +from application.workflow.i_node import INode +from models_provider.tools import get_model_instance_by_model_workspace_id + + +class RerankerSettingSerializer(serializers.Serializer): + top_n = serializers.IntegerField(required=True, label=_("Reference segment number")) + similarity = serializers.FloatField(required=True, max_value=2, min_value=0, label=_("Reference segment number")) + max_paragraph_char_number = serializers.IntegerField(required=True, + label=_("Maximum number of words in a quoted segment")) + + +class RerankerNodeSerializer(serializers.Serializer): + reranker_setting = RerankerSettingSerializer(required=True) + question_reference_address = serializers.ListField(required=True) + reranker_model_id = serializers.CharField(required=False, allow_blank=True, allow_null=True) + reranker_model_id_type = serializers.CharField(required=False, default='custom') + reranker_model_id_reference = serializers.ListField(required=False, child=serializers.CharField(), allow_empty=True) + reranker_reference_list = serializers.ListField(required=True, child=serializers.ListField(required=True)) + show_knowledge = serializers.BooleanField(required=True, + label=_("The results are displayed in the knowledge sources")) + + +def _merge_reranker_list(reranker_list, result=None): + if result is None: + result = [] + for document in reranker_list: + if isinstance(document, list): + _merge_reranker_list(document, result) + elif isinstance(document, dict): + content = document.get('title', '') + document.get('content', '') + title = document.get("title") + result.append( + Document(page_content=str(document) if len(content) == 0 else content, + metadata={'title': title, **document})) + else: + result.append(Document(page_content=str(document), metadata={})) + return result + + +def _filter_result(document_list: List[Document], max_paragraph_char_number, top_n, similarity): + use_len = 0 + result = [] + for index in range(len(document_list)): + document = document_list[index] + if use_len >= max_paragraph_char_number or index >= top_n or document.metadata.get( + 'relevance_score') < similarity: + break + content = document.page_content[0:max_paragraph_char_number - use_len] + use_len = use_len + len(content) + result.append({'page_content': content, 'metadata': document.metadata}) + return result + + +def _reset_result_list(result_list: List[Document], document_list: List[Document]): + r = [] + document_list = document_list.copy() + for result in result_list: + filter_result_list = [document for document in document_list if document.page_content == result.page_content] + if len(filter_result_list) > 0: + item = filter_result_list[0] + document_list.remove(item) + r.append(Document(page_content=item.page_content, + metadata={**item.metadata, 'relevance_score': result.metadata.get('relevance_score')})) + else: + r.append(result) + return r + + +def _reset_metadata(metadata): + meta = metadata.get('meta') + if isinstance(metadata.get('meta'), dict): + if not meta.get('allow_download', False): + metadata['meta'] = {'allow_download': False} + return metadata + + +class RerankerNode(INode): + serializer_class = RerankerNodeSerializer + supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.TOOL] + type = 'reranker-node' + + def execute(self): + node_params = self.get_parameters() + workflow_params = self.get_workflow_parameters() + + question_ref = node_params.get('question_reference_address') + reranker_reference_list = node_params.get('reranker_reference_list') + reranker_setting = node_params.get('reranker_setting') + reranker_model_id = node_params.get('reranker_model_id') + reranker_model_id_type = node_params.get('reranker_model_id_type', 'custom') + reranker_model_id_reference = node_params.get('reranker_model_id_reference') + show_knowledge = node_params.get('show_knowledge', False) + + question = self.workflow_manage.get_reference_field(question_ref[0], question_ref[1:]) + question = str(question) + + reranker_list = [self.workflow_manage.get_reference_field(ref[0], ref[1:]) + for ref in reranker_reference_list] + + if reranker_model_id_type == 'reference' and reranker_model_id_reference: + reference_data = self.workflow_manage.get_reference_field( + reranker_model_id_reference[0], reranker_model_id_reference[1:], + ) + if reference_data and isinstance(reference_data, dict): + reranker_model_id = reference_data.get('reranker_model_id', + reference_data.get('model_id', reranker_model_id)) + + if not reranker_model_id: + raise Exception(_('Model is not allowed to be empty')) + + self.write_context('show_knowledge', show_knowledge) + + documents = _merge_reranker_list(reranker_list) + documents = [d for d in documents if d.page_content and len(d.page_content) > 0] + + if len(documents) == 0: + self.write_context('document_list', []) + self.write_context('question', question) + self.write_context('result_list', []) + self.write_context('result', '') + return + + top_n = reranker_setting.get('top_n', 3) + self.write_context('document_list', [ + {'page_content': document.page_content, 'metadata': _reset_metadata(document.metadata)} + for document in documents + ]) + self.write_context('question', question) + + workspace_id = workflow_params.get('workspace_id') + reranker_model = get_model_instance_by_model_workspace_id(reranker_model_id, workspace_id, top_n=top_n) + + result = reranker_model.compress_documents(documents, question) + + similarity = reranker_setting.get('similarity', 0.6) + max_paragraph_char_number = reranker_setting.get('max_paragraph_char_number', 5000) + + result = _reset_result_list(result, documents) + r = _filter_result(result, max_paragraph_char_number, top_n, similarity) + + self.write_context('result_list', r) + self.write_context('result', ''.join([item.get('page_content') for item in r])) + self.write_context('is_hit_handling_method_list', + [row for row in r if row.get('metadata').get('is_hit_handling_method')])