diff --git a/api/core/tool/dataset_multi_retriever_tool.py b/api/core/tool/dataset_multi_retriever_tool.py index 11c32503e..07174b1d7 100644 --- a/api/core/tool/dataset_multi_retriever_tool.py +++ b/api/core/tool/dataset_multi_retriever_tool.py @@ -80,6 +80,9 @@ class DatasetMultiRetrieverTool(BaseTool): hit_callback = DatasetIndexToolCallbackHandler(self.conversation_message_task) hit_callback.on_tool_end(all_documents) + document_score_list = {} + for item in all_documents: + document_score_list[item.metadata['doc_id']] = item.metadata['score'] document_context_list = [] index_node_ids = [document.metadata['doc_id'] for document in all_documents] @@ -120,8 +123,10 @@ class DatasetMultiRetrieverTool(BaseTool): 'document_name': document.name, 'data_source_type': document.data_source_type, 'segment_id': segment.id, - 'retriever_from': self.retriever_from + 'retriever_from': self.retriever_from, + 'score': document_score_list.get(segment.index_node_id, None) } + if self.retriever_from == 'dev': source['hit_count'] = segment.hit_count source['word_count'] = segment.word_count diff --git a/api/core/tool/dataset_retriever_tool.py b/api/core/tool/dataset_retriever_tool.py index de19ccb6b..cc8b8e138 100644 --- a/api/core/tool/dataset_retriever_tool.py +++ b/api/core/tool/dataset_retriever_tool.py @@ -196,10 +196,10 @@ class DatasetRetrieverTool(BaseTool): 'document_name': document.name, 'data_source_type': document.data_source_type, 'segment_id': segment.id, - 'retriever_from': self.retriever_from + 'retriever_from': self.retriever_from, + 'score': document_score_list.get(segment.index_node_id, None) + } - if dataset.indexing_technique != "economy": - source['score'] = document_score_list.get(segment.index_node_id) if self.retriever_from == 'dev': source['hit_count'] = segment.hit_count source['word_count'] = segment.word_count