perf: improve knowledge base retrieval quality

[lxfight]

Equal-weight RRF over-promotes mediocre candidates that appear in both dense and sparse result lists, and multiple chunks from one source document can consume the limited final result slots. This PR addresses both issues with relative-score fusion and source-document diversification.

Modifications / 改动点

• Normalize dense similarity and BM25 scores per knowledge base before combining them.
• Use a dense weight of 0.9 selected on the development split.
• Keep deterministic RRF ordering only as the tie breaker.
• Return at most one highest-scoring chunk per source document in the final fused results.
• This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

Local benchmark:

• Public dataset: mteb/DuRetrieval
• Dataset revision: 8d9b3d6a6a62b7fd2d4df0821d612366dc52d14e
• Fixed split seed: astrbot-kb-v1
• Corpus: 100,001 documents / 119,965 chunks
• Queries: 2,000
• Embedding model: baai/bge-m3
• Reranking disabled to isolate retrieval changes

Strategy	Recall@5	HitRate@5	nDCG@5
Production equal-weight RRF	0.6118	0.9230	0.7087
Relative-score fusion	0.7070	0.9610	0.8038
Relative-score fusion + source deduplication	0.7181	0.9615	0.8142

The baseline also produced duplicate source documents in the top five for 249/2,000 queries. Source deduplication increased the mean number of unique documents without changing the stored chunks. Benchmark data and generated result files remain local and are not included in this PR.

Verification steps:

uv run pytest -q tests/unit/test_rank_fusion.py
uv run ruff check astrbot/core/knowledge_base/retrieval/rank_fusion.py tests/unit/test_rank_fusion.py

Results:

• 8 tests passed.
• Ruff checks passed.

---

Checklist / 检查清单

• 😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
  / 如果 PR 中有新加入的功能，已经通过 Issue / 邮件等方式和作者讨论过。
• 👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
  / 我的更改经过了良好的测试，并已在上方提供了“验证步骤”和“运行截图”。
• 🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
  / 我确保没有引入新依赖库，或者引入了新依赖库的同时将其添加到 requirements.txt 和 pyproject.toml 文件相应位置。
• 😮 My changes do not introduce malicious code.
  / 我的更改没有引入恶意代码。

Summary by Sourcery

Improve knowledge base retrieval quality by fusing normalized dense and sparse scores with document-level diversification.

New Features:

• Introduce relative-score fusion combining normalized dense similarity and BM25 scores per knowledge base.
• Add support for configuring dense retrieval weight in the fusion process.
• Ensure fused retrieval results return at most one highest-scoring chunk per source document.

Bug Fixes:

• Prevent low-quality overlapping candidates from being overvalued when they appear in both dense and sparse result lists.
• Avoid duplicate source documents in top-k fused retrieval results.

Enhancements:

• Use RRF only as a deterministic tie-breaker on top of fused scores instead of as the primary fusion method.
• Normalize dense and sparse scores independently within each knowledge base to avoid cross-index score bias.
• Handle non-positive top_k values gracefully by returning an empty result set.
• Refine test coverage around fusion behavior, dense-vs-sparse preference, and document-level deduplication.

[dosubot]

📄 Knowledge review

Dosu skipped reviewing this PR because your organization has used its 200 included credits for the month. Your usage will reset on 2026-09-01. To have Dosu review this PR before then, ask your organization admin to upgrade to a pro account.

---

[sourcery-ai]

Hey - I've found 1 issue, and left some high level feedback:

• In the dense-score normalization, results with missing kb_id are grouped under an empty-string key, which may unintentionally mix scores across knowledge bases when metadata is incomplete; consider using None or a per-index grouping key to keep normalization aligned with the actual source index.
• The fusion logic now sorts all candidate IDs and only applies top_k when building fused_results, which can be wasteful for large result sets; you can restore the previous [:top_k] slice on sorted_ids to avoid unnecessary processing while preserving the new diversification behavior.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In the dense-score normalization, results with missing `kb_id` are grouped under an empty-string key, which may unintentionally mix scores across knowledge bases when metadata is incomplete; consider using `None` or a per-index grouping key to keep normalization aligned with the actual source index.
- The fusion logic now sorts all candidate IDs and only applies `top_k` when building `fused_results`, which can be wasteful for large result sets; you can restore the previous `[:top_k]` slice on `sorted_ids` to avoid unnecessary processing while preserving the new diversification behavior.

## Individual Comments

### Comment 1
<location path="astrbot/core/knowledge_base/retrieval/rank_fusion.py" line_range="109" />
<code_context>
             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():
</code_context>
<issue_to_address>
**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*:

```python
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`:

```python
# 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:

```python
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:

```python
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.
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

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.