From 21fb80058f0ea3a55bf8e83d7d05f8bb53fd05ea Mon Sep 17 00:00:00 2001 From: lxfight <1686540385@qq.com> Date: Thu, 30 Jul 2026 09:32:33 +0800 Subject: [PATCH 1/2] perf: improve knowledge base score fusion --- .../knowledge_base/retrieval/rank_fusion.py | 96 +++++++++++++++---- tests/unit/test_rank_fusion.py | 45 ++++++++- 2 files changed, 116 insertions(+), 25 deletions(-) diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py index dad49ceac6..8922cd4996 100644 --- a/astrbot/core/knowledge_base/retrieval/rank_fusion.py +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -1,6 +1,6 @@ -"""检索结果融合器 +"""检索结果融合器。 -使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果 +使用相对分数融合稠密检索和稀疏检索结果，并以 RRF 处理同分结果。 """ import json @@ -28,19 +28,32 @@ class RankFusion: 职责: - 融合稠密检索和稀疏检索的结果 - - 使用 Reciprocal Rank Fusion (RRF) 算法 + - 在每个知识库内归一化不可直接比较的检索分数 + - 使用 RRF 作为确定性的同分排序依据 """ - def __init__(self, kb_db: KBSQLiteDatabase, k: int = 60) -> None: + def __init__( + self, + kb_db: KBSQLiteDatabase, + k: int = 60, + dense_weight: float = 0.9, + ) -> None: """初始化结果融合器 Args: kb_db: 知识库数据库实例 - k: RRF 参数,用于平滑排名 + k: RRF 同分排序的平滑参数 + dense_weight: 相对分数融合中的稠密检索权重 + + Raises: + ValueError: If dense_weight is outside the inclusive range [0, 1]. """ + if not 0 <= dense_weight <= 1: + raise ValueError("dense_weight must be between 0 and 1") self.kb_db = kb_db self.k = k + self.dense_weight = dense_weight async def fuse( self, @@ -48,10 +61,10 @@ async def fuse( sparse_results: list[SparseResult], top_k: int = 20, ) -> list[FusedResult]: - """融合稠密和稀疏检索结果 + """融合稠密和稀疏检索结果。 - RRF 公式: - score(doc) = sum(1 / (k + rank_i)) + 每个知识库内分别对稠密相似度和 BM25 分数做 min-max 归一化， + 再按权重合并。RRF 分数仅用于融合分数相同时的稳定排序。 Args: dense_results: 稠密检索结果 @@ -76,6 +89,7 @@ async def fuse( all_chunk_ids = set() vec_doc_id_to_dense: dict[str, Result] = {} # vec_doc_id -> Result chunk_id_to_sparse: dict[str, SparseResult] = {} # chunk_id -> SparseResult + dense_metadata: dict[str, dict] = {} # 处理稀疏检索结果 for r in sparse_results: @@ -87,27 +101,67 @@ async def fuse( vec_doc_id = r.data["doc_id"] all_chunk_ids.add(vec_doc_id) vec_doc_id_to_dense[vec_doc_id] = r + dense_metadata[vec_doc_id] = json.loads(r.data["metadata"]) + + # 3. 在每个知识库内归一化两路分数，避免跨索引直接比较 BM25。 + dense_groups: dict[str, list[tuple[str, float]]] = {} + for identifier, result in vec_doc_id_to_dense.items(): + kb_id = dense_metadata[identifier].get("kb_id") + if not kb_id and identifier in chunk_id_to_sparse: + kb_id = chunk_id_to_sparse[identifier].kb_id + dense_groups.setdefault(kb_id or "", []).append( + (identifier, result.similarity) + ) + + sparse_groups: dict[str, list[tuple[str, float]]] = {} + for identifier, result in chunk_id_to_sparse.items(): + sparse_groups.setdefault(result.kb_id, []).append( + (identifier, result.score) + ) + + normalized_dense: dict[str, float] = {} + for group in dense_groups.values(): + scores = [score for _, score in group] + minimum = min(scores) + score_range = max(scores) - minimum + for identifier, score in group: + normalized_dense[identifier] = ( + (score - minimum) / score_range if score_range else 1.0 + ) + + normalized_sparse: dict[str, float] = {} + for group in sparse_groups.values(): + scores = [score for _, score in group] + minimum = min(scores) + score_range = max(scores) - minimum + for identifier, score in group: + normalized_sparse[identifier] = ( + (score - minimum) / score_range if score_range else 1.0 + ) - # 3. 计算 RRF 分数 + # 4. 计算融合分数和 RRF 同分分数 + fusion_scores: dict[str, float] = {} rrf_scores: dict[str, float] = {} for identifier in all_chunk_ids: - score = 0.0 + fusion_scores[identifier] = self.dense_weight * normalized_dense.get( + identifier, 0.0 + ) + (1 - self.dense_weight) * normalized_sparse.get(identifier, 0.0) - # 来自稠密检索的贡献 + rrf_score = 0.0 if identifier in dense_ranks: - score += 1.0 / (self.k + dense_ranks[identifier]) + rrf_score += 1.0 / (self.k + dense_ranks[identifier]) - # 来自稀疏检索的贡献 if identifier in sparse_ranks: - score += 1.0 / (self.k + sparse_ranks[identifier]) + rrf_score += 1.0 / (self.k + sparse_ranks[identifier]) - rrf_scores[identifier] = score + rrf_scores[identifier] = rrf_score - # 4. 排序 + # 5. 排序 sorted_ids = sorted( - rrf_scores, + fusion_scores, key=lambda cid: ( + -fusion_scores[cid], -rrf_scores[cid], dense_ranks.get(cid, float("inf")), sparse_ranks.get(cid, float("inf")), @@ -115,7 +169,7 @@ async def fuse( ), )[:top_k] - # 5. 构建融合结果 + # 6. 构建融合结果 fused_results = [] for identifier in sorted_ids: # 优先从稀疏检索获取完整信息 @@ -128,13 +182,13 @@ async def fuse( doc_id=sr.doc_id, kb_id=sr.kb_id, content=sr.content, - score=rrf_scores[identifier], + score=fusion_scores[identifier], ), ) elif identifier in vec_doc_id_to_dense: # 从向量检索获取信息,需要从数据库获取块的详细信息 vec_result = vec_doc_id_to_dense[identifier] - chunk_md = json.loads(vec_result.data["metadata"]) + chunk_md = dense_metadata[identifier] fused_results.append( FusedResult( chunk_id=identifier, @@ -142,7 +196,7 @@ async def fuse( doc_id=chunk_md["kb_doc_id"], kb_id=chunk_md["kb_id"], content=vec_result.data["text"], - score=rrf_scores[identifier], + score=fusion_scores[identifier], ), ) diff --git a/tests/unit/test_rank_fusion.py b/tests/unit/test_rank_fusion.py index b1e0e88713..73d3cb914e 100644 --- a/tests/unit/test_rank_fusion.py +++ b/tests/unit/test_rank_fusion.py @@ -1,3 +1,5 @@ +import json + import pytest from astrbot.core.db.vec_db.base import Result @@ -11,7 +13,13 @@ def make_dense_result(chunk_id: str, similarity: float) -> Result: data={ "doc_id": chunk_id, "text": chunk_id, - "metadata": "{}", + "metadata": json.dumps( + { + "chunk_index": 0, + "kb_doc_id": f"doc-{chunk_id}", + "kb_id": "kb", + } + ), }, ) @@ -33,6 +41,12 @@ def make_sparse_result( ) +@pytest.mark.parametrize("dense_weight", [-0.1, 1.1]) +def test_rank_fusion_rejects_invalid_dense_weight(dense_weight): + with pytest.raises(ValueError, match="dense_weight"): + RankFusion(kb_db=None, dense_weight=dense_weight) + + @pytest.mark.asyncio async def test_rank_fusion_uses_source_rank_for_independent_sparse_indexes(): dense_results = [ @@ -56,11 +70,11 @@ async def test_rank_fusion_uses_source_rank_for_independent_sparse_indexes(): "large-1", "large-2", ] - assert results[0].score == pytest.approx(2 / 61) + assert results[0].score == pytest.approx(1.0) @pytest.mark.asyncio -async def test_rank_fusion_prefers_dense_rank_when_scores_are_equal(): +async def test_rank_fusion_prefers_dense_signal_when_sources_disagree(): dense_results = [ make_dense_result("dense-first", 0.99), make_dense_result("sparse-first", 0.98), @@ -75,11 +89,12 @@ async def test_rank_fusion_prefers_dense_rank_when_scores_are_equal(): sparse_results=sparse_results, ) - assert results[0].score == pytest.approx(results[1].score) assert [result.chunk_id for result in results] == [ "dense-first", "sparse-first", ] + assert results[0].score == pytest.approx(0.9) + assert results[1].score == pytest.approx(0.1) @pytest.mark.asyncio @@ -106,3 +121,25 @@ async def test_rank_fusion_uses_chunk_id_as_stable_final_tiebreaker(): "chunk-a", "chunk-b", ] + + +@pytest.mark.asyncio +async def test_rank_fusion_does_not_overvalue_low_rank_source_overlap(): + dense_results = [make_dense_result("dense-best", 0.99)] + [ + make_dense_result(f"dense-{rank}", 0.9 - rank / 100) + for rank in range(2, 51) + ] + sparse_results = [ + make_sparse_result(f"sparse-{rank}", "kb", 51 - rank, rank) + for rank in range(1, 50) + ] + [make_sparse_result("dense-50", "kb", 1.0, 50)] + + results = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=100, + ) + result_ids = [result.chunk_id for result in results] + + assert result_ids[0] == "dense-best" + assert result_ids.index("dense-best") < result_ids.index("dense-50") From 5b72981a72fb885dd88208810379bb8f525e7c75 Mon Sep 17 00:00:00 2001 From: lxfight <1686540385@qq.com> Date: Thu, 30 Jul 2026 09:45:22 +0800 Subject: [PATCH 2/2] perf: diversify knowledge base source documents --- .../knowledge_base/retrieval/rank_fusion.py | 50 +++++++++++-------- tests/unit/test_rank_fusion.py | 40 +++++++++++++-- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py index 8922cd4996..2bc8618072 100644 --- a/astrbot/core/knowledge_base/retrieval/rank_fusion.py +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -75,6 +75,9 @@ async def fuse( List[FusedResult]: 融合后的结果列表 """ + if top_k <= 0: + return [] + # 1. 构建排名映射 dense_ranks = { r.data["doc_id"]: (idx + 1) for idx, r in enumerate(dense_results) @@ -167,37 +170,44 @@ async def fuse( sparse_ranks.get(cid, float("inf")), cid, ), - )[:top_k] + ) - # 6. 构建融合结果 + # 6. 构建融合结果，每篇来源文档只保留得分最高的块。 fused_results = [] + seen_documents: set[tuple[str, str]] = set() for identifier in sorted_ids: # 优先从稀疏检索获取完整信息 if identifier in chunk_id_to_sparse: sr = chunk_id_to_sparse[identifier] - fused_results.append( - FusedResult( - chunk_id=sr.chunk_id, - chunk_index=sr.chunk_index, - doc_id=sr.doc_id, - kb_id=sr.kb_id, - content=sr.content, - score=fusion_scores[identifier], - ), + fused_result = FusedResult( + chunk_id=sr.chunk_id, + chunk_index=sr.chunk_index, + doc_id=sr.doc_id, + kb_id=sr.kb_id, + content=sr.content, + score=fusion_scores[identifier], ) elif identifier in vec_doc_id_to_dense: # 从向量检索获取信息,需要从数据库获取块的详细信息 vec_result = vec_doc_id_to_dense[identifier] chunk_md = dense_metadata[identifier] - fused_results.append( - FusedResult( - chunk_id=identifier, - chunk_index=chunk_md["chunk_index"], - doc_id=chunk_md["kb_doc_id"], - kb_id=chunk_md["kb_id"], - content=vec_result.data["text"], - score=fusion_scores[identifier], - ), + fused_result = FusedResult( + chunk_id=identifier, + chunk_index=chunk_md["chunk_index"], + doc_id=chunk_md["kb_doc_id"], + kb_id=chunk_md["kb_id"], + content=vec_result.data["text"], + score=fusion_scores[identifier], ) + else: + continue + + document_key = (fused_result.kb_id, fused_result.doc_id) + if document_key in seen_documents: + continue + seen_documents.add(document_key) + fused_results.append(fused_result) + if len(fused_results) >= top_k: + break return fused_results diff --git a/tests/unit/test_rank_fusion.py b/tests/unit/test_rank_fusion.py index 73d3cb914e..472f3426a9 100644 --- a/tests/unit/test_rank_fusion.py +++ b/tests/unit/test_rank_fusion.py @@ -29,11 +29,12 @@ def make_sparse_result( kb_id: str, score: float, rank: int, + doc_id: str | None = None, ) -> SparseResult: return SparseResult( chunk_index=0, chunk_id=chunk_id, - doc_id=f"doc-{chunk_id}", + doc_id=doc_id or f"doc-{chunk_id}", kb_id=kb_id, content=chunk_id, score=score, @@ -47,6 +48,17 @@ def test_rank_fusion_rejects_invalid_dense_weight(dense_weight): RankFusion(kb_db=None, dense_weight=dense_weight) +@pytest.mark.asyncio +async def test_rank_fusion_returns_empty_for_non_positive_top_k(): + results = await RankFusion(kb_db=None).fuse( + dense_results=[make_dense_result("chunk", 0.99)], + sparse_results=[], + top_k=0, + ) + + assert results == [] + + @pytest.mark.asyncio async def test_rank_fusion_uses_source_rank_for_independent_sparse_indexes(): dense_results = [ @@ -126,8 +138,7 @@ async def test_rank_fusion_uses_chunk_id_as_stable_final_tiebreaker(): @pytest.mark.asyncio async def test_rank_fusion_does_not_overvalue_low_rank_source_overlap(): dense_results = [make_dense_result("dense-best", 0.99)] + [ - make_dense_result(f"dense-{rank}", 0.9 - rank / 100) - for rank in range(2, 51) + make_dense_result(f"dense-{rank}", 0.9 - rank / 100) for rank in range(2, 51) ] sparse_results = [ make_sparse_result(f"sparse-{rank}", "kb", 51 - rank, rank) @@ -143,3 +154,26 @@ async def test_rank_fusion_does_not_overvalue_low_rank_source_overlap(): assert result_ids[0] == "dense-best" assert result_ids.index("dense-best") < result_ids.index("dense-50") + + +@pytest.mark.asyncio +async def test_rank_fusion_keeps_only_the_best_chunk_per_document(): + dense_results = [ + make_dense_result("doc-a-best", 0.99), + make_dense_result("doc-a-second", 0.98), + make_dense_result("doc-b", 0.97), + ] + sparse_results = [ + make_sparse_result("doc-a-best", "kb", 10.0, 1, doc_id="doc-a"), + make_sparse_result("doc-a-second", "kb", 9.0, 2, doc_id="doc-a"), + make_sparse_result("doc-b", "kb", 8.0, 3, doc_id="doc-b"), + ] + + results = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=2, + ) + + assert [result.chunk_id for result in results] == ["doc-a-best", "doc-b"] + assert [result.doc_id for result in results] == ["doc-a", "doc-b"]