Fix/economical knowledge retrieval (langgenius#9396)

AceDataCloud · Nov 9, 2024 · 82dff1d · 82dff1d
1 parent c964d3b
commit 82dff1d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py
@@ -217,7 +217,7 @@ def retrieve(
                                 "data_source_type": document.data_source_type,
                                 "segment_id": segment.id,
                                 "retriever_from": invoke_from.to_source(),
-                                "score": document_score_list.get(segment.index_node_id, None),
+                                "score": document_score_list.get(segment.index_node_id, 0.0),
                             }
 
                             if invoke_from.to_source() == "dev":
@@ -231,12 +231,12 @@ def retrieve(
                                 source["content"] = segment.content
                             retrieval_resource_list.append(source)
         if hit_callback and retrieval_resource_list:
-            retrieval_resource_list = sorted(retrieval_resource_list, key=lambda x: x.get("score"), reverse=True)
+            retrieval_resource_list = sorted(retrieval_resource_list, key=lambda x: x.get("score") or 0.0, reverse=True)
             for position, item in enumerate(retrieval_resource_list, start=1):
                 item["position"] = position
             hit_callback.return_retriever_resource_info(retrieval_resource_list)
         if document_context_list:
-            document_context_list = sorted(document_context_list, key=lambda x: x.score, reverse=True)
+            document_context_list = sorted(document_context_list, key=lambda x: x.score or 0.0, reverse=True)
             return str("\n".join([document_context.content for document_context in document_context_list]))
         return ""
 

diff --git a/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py
@@ -234,7 +234,7 @@ def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query:
                         retrieval_resource_list.append(source)
         if retrieval_resource_list:
             retrieval_resource_list = sorted(
-                retrieval_resource_list, key=lambda x: x.get("metadata").get("score"), reverse=True
+                retrieval_resource_list, key=lambda x: x.get("metadata").get("score") or 0.0, reverse=True
             )
             position = 1
             for item in retrieval_resource_list: