Skip to content

perf: improve knowledge base retrieval quality #9455

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lxfight wants to merge 2 commits into
base: master
Choose a base branch
from
+178 −43
Open

perf: improve knowledge base retrieval quality #9455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
21fb800
perf: improve knowledge base score fusion
lxfight Jul 30, 2026
5b72981
perf: diversify knowledge base source documents
lxfight Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
140 changes: 102 additions & 38 deletions astrbot/core/knowledge_base/retrieval/rank_fusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""检索结果融合器
"""检索结果融合器。

使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果
使用相对分数融合稠密检索和稀疏检索结果，并以 RRF 处理同分结果。
"""

import json
Expand Down Expand Up @@ -28,30 +28,43 @@ 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,
dense_results: list[Result],
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: 稠密检索结果
Expand All @@ -62,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)
Expand All @@ -76,6 +92,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:
Expand All @@ -87,63 +104,110 @@ 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。

@sourcery-ai sourcery-ai Bot Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (complexity): Consider extracting normalization and candidate-building helpers so fuse focuses on orchestration rather than detailed score and dict management.

You can keep all the new behavior but make fuse much easier to follow by extracting a couple of helpers and removing duplicated logic.

1. Extract score normalization into a helper

The dense/sparse normalization loops are identical. Factor them out so fuse only expresses what is normalized, not how:

from collections.abc import Iterable

def _normalize_groups(
    self,
    groups: dict[str, list[tuple[str, float]]],
) -> dict[str, float]:
    normalized: dict[str, float] = {}
    for group in groups.values():
        scores = [score for _, score in group]
        minimum = min(scores)
        score_range = max(scores) - minimum
        for identifier, score in group:
            normalized[identifier] = (
                (score - minimum) / score_range if score_range else 1.0
            )
    return normalized

Use it in fuse:

# 3. 在每个知识库内归一化两路分数
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") or \
            chunk_id_to_sparse.get(identifier, SparseResult("", "", "", "", 0, 0, 0)).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 = self._normalize_groups(dense_groups)
normalized_sparse = self._normalize_groups(sparse_groups)

(You may want a tiny helper for resolving kb_id cleanly instead of the inline or chain.)

2. Introduce a small internal structure to reduce parallel dicts

Right now vec_doc_id_to_dense, chunk_id_to_sparse, dense_metadata, fusion_scores, rrf_scores all key off identifier. A minimal internal “candidate” object lets you collapse these parallel maps:

from dataclasses import dataclass

@dataclass
class _Candidate:
    identifier: str      # chunk_id
    kb_id: str
    doc_id: str
    chunk_index: int | None
    dense_score: float | None
    sparse_score: float | None
    metadata: dict | None  # for dense-only cases

Then, in fuse, you can build candidates in one place:

candidates: dict[str, _Candidate] = {}

# dense
for r in dense_results:
    identifier = r.data["doc_id"]
    md = json.loads(r.data["metadata"])
    cand = candidates.get(identifier)
    if cand is None:
        cand = _Candidate(
            identifier=identifier,
            kb_id=md["kb_id"],
            doc_id=md["kb_doc_id"],
            chunk_index=md["chunk_index"],
            dense_score=r.similarity,
            sparse_score=None,
            metadata=md,
        )
        candidates[identifier] = cand
    else:
        cand.dense_score = r.similarity

# sparse
for r in sparse_results:
    cand = candidates.get(r.chunk_id)
    if cand is None:
        cand = _Candidate(
            identifier=r.chunk_id,
            kb_id=r.kb_id,
            doc_id=r.doc_id,
            chunk_index=r.chunk_index,
            dense_score=None,
            sparse_score=r.score,
            metadata=None,
        )
        candidates[r.chunk_id] = cand
    else:
        cand.sparse_score = r.score

After that:

  • dense_groups / sparse_groups can be built from candidates.values() instead of separate dicts.
  • fusion_scores and rrf_scores can be stored directly on _Candidate or in small dicts keyed by identifier.
  • Building FusedResult no longer needs to look back into multiple maps; everything is on the candidate.

This reduces the cognitive load of keeping many coordinated dictionaries in sync while preserving all current behavior.

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")),
cid,
),
)[:top_k]
)

# 5. 构建融合结果
# 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=rrf_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 = json.loads(vec_result.data["metadata"])
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=rrf_scores[identifier],
),
chunk_md = dense_metadata[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
81 changes: 76 additions & 5 deletions tests/unit/test_rank_fusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import json

import pytest

from astrbot.core.db.vec_db.base import Result
Expand All @@ -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",
}
),
},
)

Expand All @@ -21,18 +29,36 @@ 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,
rank=rank,
)


@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_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 = [
Expand All @@ -56,11 +82,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),
Expand All @@ -75,11 +101,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
Expand All @@ -106,3 +133,47 @@ 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")


@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"]
Loading