diff --git a/astrbot/core/knowledge_base/retrieval/manager.py b/astrbot/core/knowledge_base/retrieval/manager.py index 5b414b171e..1cce536856 100644 --- a/astrbot/core/knowledge_base/retrieval/manager.py +++ b/astrbot/core/knowledge_base/retrieval/manager.py @@ -139,6 +139,7 @@ async def retrieve( dense_results=dense_results, sparse_results=sparse_results, top_k=top_k_fusion, + query=query, ) time_end = time.time() logger.debug( diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py index 0fb3ab8997..ff10979642 100644 --- a/astrbot/core/knowledge_base/retrieval/rank_fusion.py +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -60,6 +60,7 @@ async def fuse( dense_results: list[Result], sparse_results: list[SparseResult], top_k: int = 20, + query: str | None = None, ) -> list[FusedResult]: """融合稠密和稀疏检索结果。 @@ -72,6 +73,9 @@ async def fuse( dense_results: 稠密检索结果 sparse_results: 稀疏检索结果 top_k: 返回结果数量 + query: 原始查询文本。提供时，对 Dense 完全未召回但 Sparse 命中 + 且内容包含查询词（大小写不敏感）的候选给予保底提升，避免 + Dense embedding 的大小写敏感导致精确词面匹配被挤出 top_k。 Returns: List[FusedResult]: 融合后的结果列表 @@ -153,10 +157,32 @@ async def fuse( rrf_scores[identifier] = rrf_score - # 5. 排序 + # 保护 Dense 漏召回的精确词面匹配：Dense embedding 对大小写敏感， + # 当查询词以大小写变体出现时（如查询 "oni" 而文档中是 "Oni"）， + # 目标 chunk 可能完全不被 Dense 召回。此时若 Sparse 命中了该 chunk + # 且其内容仅以大小写变体形式包含查询词，则将其显式排到所有普通 + # 候选之前，确保它不会被大量纯 Dense 候选挤出 top_k。 + # 仅保护"大小写变体"匹配（原样查询词未出现、小写形式出现）， + # 避免改变普通精确命中的既有排序。 + protected_ids: set[str] = set() + if query is not None and query.strip(): + raw_query = query.strip() + lowered_query = raw_query.lower() + for identifier in all_chunk_ids: + if identifier in vec_doc_id_to_dense: + continue + sparse_result = chunk_id_to_sparse.get(identifier) + if not sparse_result: + continue + content = sparse_result.content or "" + if raw_query not in content and lowered_query in content.lower(): + protected_ids.add(identifier) + + # 5. 排序。受保护的大小写变体精确匹配优先于所有普通候选。 sorted_ids = sorted( fusion_scores, key=lambda cid: ( + cid not in protected_ids, -fusion_scores[cid], -rrf_scores[cid], dense_ranks.get(cid, float("inf")), diff --git a/tests/unit/test_rank_fusion.py b/tests/unit/test_rank_fusion.py index 94e5217ecd..f48a20ce8a 100644 --- a/tests/unit/test_rank_fusion.py +++ b/tests/unit/test_rank_fusion.py @@ -256,3 +256,130 @@ async def test_rank_fusion_does_not_promote_a_single_low_scoring_kb_result(): "weak", ] assert results[-1].score == pytest.approx(0.1) + + +@pytest.mark.asyncio +async def test_rank_fusion_protects_case_insensitive_sparse_exact_match(): + # Dense 因 embedding 对大小写敏感而完全没有召回目标 chunk（查询 "oni"， + # 文档中是 "Oni"）。Sparse (FTS5) 大小写不敏感，命中了目标 chunk。 + # 大量高分的纯 Dense 候选会把目标挤出 top_k，除非融合阶段对 + # 大小写不敏感的词面匹配给予保护。 + dense_results = [ + make_dense_result(f"dense-{rank}", 0.95 - rank / 100) for rank in range(1, 21) + ] + sparse_results = [ + make_sparse_result( + "target-oni", + "kb", + 30.0, + 1, + content="#### 恶鬼\n**恶鬼** **Oni** 是日式奇幻中的经典怪物。", + ), + *[ + make_sparse_result(f"sparse-{rank}", "kb", 20.0 - rank, rank) + for rank in range(2, 11) + ], + ] + + # 不传 query：目标 chunk 仅靠稀疏侧得分，被大量 dense 候选挤出。 + without_query = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=5, + ) + assert "target-oni" not in [r.chunk_id for r in without_query] + + # 传入 query：目标 chunk 因大小写不敏感的词面匹配被保底提升，进入 top_k。 + with_query = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=5, + query="oni", + ) + assert "target-oni" in [r.chunk_id for r in with_query] + assert with_query[0].chunk_id == "target-oni" + + +@pytest.mark.asyncio +async def test_rank_fusion_query_protection_ignores_dense_recalled_chunks(): + # 已同时被 Dense 召回的目标 chunk 不应因 query 保护而再次提升排序， + # 即保护只作用于 Dense 完全漏召回的候选。 + dense_results = [ + make_dense_result("target-oni", 0.99), + make_dense_result("other", 0.95), + ] + sparse_results = [ + make_sparse_result( + "target-oni", + "kb", + 30.0, + 1, + content="#### 恶鬼\n**恶鬼** **Oni**", + ), + ] + + with_query = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=5, + query="oni", + ) + + assert [r.chunk_id for r in with_query] == ["target-oni", "other"] + assert with_query[0].score == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_rank_fusion_query_protection_survives_dense_weight_boundary(): + # dense_weight=1.0 是合法边界值。保护不依赖分数保底，而是显式排序， + # 因此即使 dense 权重为 1.0，大小写变体精确匹配仍必须排在第一位。 + dense_results = [ + make_dense_result(f"dense-{rank}", 0.95 - rank / 100) for rank in range(1, 11) + ] + sparse_results = [ + make_sparse_result( + "target-oni", + "kb", + 30.0, + 1, + content="#### 恶鬼\n**恶鬼** **Oni** 是日式奇幻中的经典怪物。", + ), + ] + + results = await RankFusion(kb_db=None, dense_weight=1.0).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=5, + query="oni", + ) + + assert results[0].chunk_id == "target-oni" + + +@pytest.mark.asyncio +async def test_rank_fusion_query_protection_skips_exact_case_matches(): + # 只有当查询词以大小写变体出现时才保护（原样词未出现、小写形式出现）。 + # 如果 chunk 内容本身就包含原样的查询词，说明不是大小写变体问题， + # 应保持既有融合排序，不触发保护。 + dense_results = [ + make_dense_result(f"dense-{rank}", 0.95 - rank / 100) for rank in range(1, 11) + ] + sparse_results = [ + make_sparse_result( + "target-oni", + "kb", + 30.0, + 1, + content="#### 恶鬼\n**恶鬼** **oni** 是日式奇幻中的经典怪物。", + ), + ] + + results = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=5, + query="oni", + ) + + # 内容包含原样查询词 "oni"，不是大小写变体，不触发保护。 + assert results[0].chunk_id != "target-oni"