Skip to content

fix: correct knowledge base ranking behavior - #9454

Open
lxfight wants to merge 2 commits into
AstrBotDevs:masterfrom
lxfight:fix/kb-ranking-correctness
Open

fix: correct knowledge base ranking behavior#9454
lxfight wants to merge 2 commits into
AstrBotDevs:masterfrom
lxfight:fix/kb-ranking-correctness

Conversation

@lxfight

@lxfight lxfight commented Jul 30, 2026

Copy link
Copy Markdown
Member

Knowledge base rank fusion could produce unstable ordering for tied candidates, while the vLLM rerank provider could silently treat malformed or failed HTTP responses as empty results. This PR makes both failure modes explicit and deterministic.

Modifications / 改动点

  • Preserve deterministic candidate ordering when fused scores are tied.
  • Validate vLLM rerank HTTP status and response shape before consuming results.
  • Add focused regression coverage for rank ties, malformed responses, and provider errors.
  • This is NOT a breaking change. / 这不是一个破坏性变更。

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

Verification steps:

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

Results:

  • 13 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.txtpyproject.toml 文件相应位置。
  • 😮 My changes do not introduce malicious code.
    / 我的更改没有引入恶意代码。

Summary by Sourcery

Clarify and harden knowledge base ranking by making rank fusion tie-breaking deterministic and enforcing strict validation of vLLM rerank API responses.

Bug Fixes:

  • Ensure rank-fused knowledge base results use deterministic tie-breaking based on dense rank, sparse rank, and chunk ID.
  • Prevent vLLM rerank provider from silently treating HTTP or schema errors as empty results by surfacing invalid responses and HTTP failures as exceptions.

Tests:

  • Add regression tests covering rank fusion tie scenarios to guarantee stable ordering across input permutations.
  • Add tests for vLLM rerank provider to verify correct handling of successful responses, HTTP error statuses, malformed payloads, empty document lists, and session termination.

@dosubot dosubot Bot added size:M This PR changes 30-99 lines, ignoring generated files. area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. feature:knowledge-base The bug / feature is about knowledge base labels Jul 30, 2026
@dosubot

dosubot Bot commented Jul 30, 2026

Copy link
Copy Markdown

📄 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-08-01. To have Dosu review this PR before then, ask your organization admin to upgrade to a pro account.


Leave Feedback Ask Dosu about AstrBot Add Dosu to your team

@sourcery-ai sourcery-ai Bot left a comment

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.

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

  • The rerank error handling now raises generic ValueError with fixed messages; consider including key details from the raw response (e.g., status code or a truncated body) so callers and logs can more easily diagnose integration issues.
  • In rank_fusion.py, the final sort key uses cid as a tiebreaker; if chunk_id can be non-string types or mixed formats, you may want to normalize it explicitly (e.g., cast to string) or document the expected type to avoid subtle ordering changes.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- The rerank error handling now raises generic `ValueError` with fixed messages; consider including key details from the raw response (e.g., status code or a truncated body) so callers and logs can more easily diagnose integration issues.
- In `rank_fusion.py`, the final sort key uses `cid` as a tiebreaker; if `chunk_id` can be non-string types or mixed formats, you may want to normalize it explicitly (e.g., cast to string) or document the expected type to avoid subtle ordering changes.

## Individual Comments

### Comment 1
<location path="tests/unit/test_rank_fusion.py" line_range="63" />
<code_context>
+
+
+@pytest.mark.asyncio
+async def test_rank_fusion_prefers_dense_rank_when_scores_are_equal():
+    dense_results = [
+        make_dense_result("dense-first", 0.99),
</code_context>
<issue_to_address>
**suggestion (testing):** Add coverage for cases where dense ranks are tied and sparse ranks act as the secondary tiebreaker.

The new tests cover dense-rank preference and `chunk_id` as the final tiebreaker. To fully exercise `(-score, dense_rank, sparse_rank, cid)`, please add a test where multiple candidates share the same fused score and dense rank (or have no dense rank), but differ in sparse rank. Assert that the ordering follows `sparse_rank` and remains stable regardless of input order, so we explicitly verify `sparse_rank` is used as the secondary tiebreaker when dense ranks don’t distinguish candidates.

Suggested implementation:

```python
    results = await RankFusion(kb_db=None).fuse(
        dense_results=dense_results,
        sparse_results=sparse_results,

@pytest.mark.asyncio
async def test_rank_fusion_uses_sparse_rank_as_secondary_tiebreaker_when_dense_ranks_tied():
    # Two candidates with identical scores and (effectively) identical dense rank,
    # but different sparse ranks; ordering should follow sparse_rank and be stable
    # regardless of input order.
    sparse_a = make_sparse_result("sparse-a", "kb", 10.0, 1)
    sparse_b = make_sparse_result("sparse-b", "kb", 10.0, 2)

    # Explicitly set sparse_rank to ensure we can control it even if the helper
    # doesn't derive it from position.
    # Adjust attribute name here if your sparse result objects use a different field.
    sparse_a.sparse_rank = 0
    sparse_b.sparse_rank = 1

    # No dense results: fused ordering must use sparse_rank as the secondary
    # tiebreaker for candidates with equal fused score.
    fusion = RankFusion(kb_db=None)

    # First, fuse in one order
    results_order_1 = await fusion.fuse(
        dense_results=[],
        sparse_results=[sparse_b, sparse_a],
    )
    # Then, fuse in the opposite order
    results_order_2 = await fusion.fuse(
        dense_results=[],
        sparse_results=[sparse_a, sparse_b],
    )

    # In both cases, sparse_a (with better sparse_rank) should come first, and
    # the ordering should be stable regardless of input order.
    assert [r.chunk_id for r in results_order_1] == [sparse_a.chunk_id, sparse_b.chunk_id]
    assert [r.chunk_id for r in results_order_2] == [sparse_a.chunk_id, sparse_b.chunk_id]

```

1. Ensure that `sparse_rank` is indeed the attribute name used by your sparse result objects; if it's named differently (e.g. `rank` or `sparse_position`), update the two assignments accordingly.
2. Place the new test so that it is *outside* the body of `test_rank_fusion_prefers_dense_rank_when_scores_are_equal` (i.e. after that function's closing parenthesis and any assertions). The SEARCH/REPLACE block above may need slight adjustment to match the exact code around the end of that test.
3. If `make_sparse_result` already sets `sparse_rank` based on an argument or its position, you can omit the explicit `sparse_a.sparse_rank` / `sparse_b.sparse_rank` assignments and instead construct the objects so that they have distinct `sparse_rank` values.
4. If your fused result objects don't expose `chunk_id` as `r.chunk_id`, update the two assertions to use whatever identifier you normally use to distinguish candidates (e.g. `r.doc_id`), while keeping the ordering checks the same.
</issue_to_address>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Comment thread tests/unit/test_rank_fusion.py
@lxfight

lxfight commented Jul 30, 2026

Copy link
Copy Markdown
Member Author

Reviewed the remaining high-level suggestions:

  • HTTP failures already retain diagnostic status through response.raise_for_status() / aiohttp.ClientResponseError.status; regression tests cover 400 and 429. I intentionally did not append raw response bodies because rerank services may echo user content or other sensitive payload data.
  • Chunk identifiers are consistently typed as str in SparseResult, FusedResult, and document storage, so mixed-type sorting is outside the supported contract.

No production change is needed for these two points.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. feature:knowledge-base The bug / feature is about knowledge base size:M This PR changes 30-99 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant