This commit is contained in:
36
backend/app/rag/fusion.py
Normal file
36
backend/app/rag/fusion.py
Normal file
@@ -0,0 +1,36 @@
|
||||
# rag/fusion.py
|
||||
|
||||
from typing import List, Dict
|
||||
from langchain_core.documents import Document
|
||||
|
||||
def reciprocal_rank_fusion(
|
||||
doc_lists: List[List[Document]],
|
||||
k: int = 60
|
||||
) -> List[Document]:
|
||||
"""
|
||||
对多个检索结果列表进行 RRF 融合。
|
||||
|
||||
Args:
|
||||
doc_lists: 多个检索结果列表,每个列表来自一个查询
|
||||
k: RRF 常数,通常设为 60
|
||||
|
||||
Returns:
|
||||
融合后按 RRF 得分降序排列的文档列表
|
||||
"""
|
||||
# 使用文档内容作为唯一标识(如果内容相同但 metadata 不同,视为同一文档)
|
||||
# 更好的做法是用 docstore 的 ID,这里简化处理:用内容 hash
|
||||
doc_to_score: Dict[str, float] = {}
|
||||
doc_map: Dict[str, Document] = {}
|
||||
|
||||
for docs in doc_lists:
|
||||
for rank, doc in enumerate(docs, start=1):
|
||||
# 生成唯一标识符(内容+来源组合,避免不同文件相同内容混淆)
|
||||
doc_id = f"{doc.page_content[:200]}_{doc.metadata.get('source', '')}"
|
||||
if doc_id not in doc_map:
|
||||
doc_map[doc_id] = doc
|
||||
score = doc_to_score.get(doc_id, 0.0) + 1.0 / (k + rank)
|
||||
doc_to_score[doc_id] = score
|
||||
|
||||
# 按得分排序
|
||||
sorted_ids = sorted(doc_to_score.keys(), key=lambda x: doc_to_score[x], reverse=True)
|
||||
return [doc_map[doc_id] for doc_id in sorted_ids]
|
||||
Reference in New Issue
Block a user