diff --git a/api/core/rag/rerank/rerank_model.py b/api/core/rag/rerank/rerank_model.py index 40ebf0befd9189..fc82b2080b2b3c 100644 --- a/api/core/rag/rerank/rerank_model.py +++ b/api/core/rag/rerank/rerank_model.py @@ -27,18 +27,17 @@ def run( :return: """ docs = [] - doc_id = [] + doc_id = set() unique_documents = [] - dify_documents = [item for item in documents if item.provider == "dify"] - external_documents = [item for item in documents if item.provider == "external"] - for document in dify_documents: - if document.metadata["doc_id"] not in doc_id: - doc_id.append(document.metadata["doc_id"]) + for document in documents: + if document.provider == "dify" and document.metadata["doc_id"] not in doc_id: + doc_id.add(document.metadata["doc_id"]) docs.append(document.page_content) unique_documents.append(document) - for document in external_documents: - docs.append(document.page_content) - unique_documents.append(document) + elif document.provider == "external": + if document not in unique_documents: + docs.append(document.page_content) + unique_documents.append(document) documents = unique_documents