From 53571adfe0b0dbac39ed76b6a61dd2f9552bc397 Mon Sep 17 00:00:00 2001 From: lxfight <1686540385@qq.com> Date: Thu, 30 Jul 2026 09:05:10 +0800 Subject: [PATCH 1/2] fix: make knowledge base rank fusion deterministic --- .../knowledge_base/retrieval/rank_fusion.py | 10 ++-- tests/unit/test_rank_fusion.py | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py index 39f402d99c..dad49ceac6 100644 --- a/astrbot/core/knowledge_base/retrieval/rank_fusion.py +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -106,9 +106,13 @@ async def fuse( # 4. 排序 sorted_ids = sorted( - rrf_scores.keys(), - key=lambda cid: rrf_scores[cid], - reverse=True, + rrf_scores, + key=lambda cid: ( + -rrf_scores[cid], + dense_ranks.get(cid, float("inf")), + sparse_ranks.get(cid, float("inf")), + cid, + ), )[:top_k] # 5. 构建融合结果 diff --git a/tests/unit/test_rank_fusion.py b/tests/unit/test_rank_fusion.py index 534944bad8..b1e0e88713 100644 --- a/tests/unit/test_rank_fusion.py +++ b/tests/unit/test_rank_fusion.py @@ -57,3 +57,52 @@ async def test_rank_fusion_uses_source_rank_for_independent_sparse_indexes(): "large-2", ] assert results[0].score == pytest.approx(2 / 61) + + +@pytest.mark.asyncio +async def test_rank_fusion_prefers_dense_rank_when_scores_are_equal(): + dense_results = [ + make_dense_result("dense-first", 0.99), + make_dense_result("sparse-first", 0.98), + ] + sparse_results = [ + make_sparse_result("sparse-first", "kb", 10.0, 1), + make_sparse_result("dense-first", "kb", 9.0, 2), + ] + + results = await RankFusion(kb_db=None).fuse( + dense_results=dense_results, + 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", + ] + + +@pytest.mark.asyncio +async def test_rank_fusion_uses_chunk_id_as_stable_final_tiebreaker(): + sparse_results = [ + make_sparse_result("chunk-b", "kb", 10.0, 1), + make_sparse_result("chunk-a", "kb", 10.0, 1), + ] + + forward_results = await RankFusion(kb_db=None).fuse( + dense_results=[], + sparse_results=sparse_results, + ) + reverse_results = await RankFusion(kb_db=None).fuse( + dense_results=[], + sparse_results=list(reversed(sparse_results)), + ) + + assert [result.chunk_id for result in forward_results] == [ + "chunk-a", + "chunk-b", + ] + assert [result.chunk_id for result in reverse_results] == [ + "chunk-a", + "chunk-b", + ] From 9f7277e5cc5700de7f05218bbc9433973fc671e6 Mon Sep 17 00:00:00 2001 From: lxfight <1686540385@qq.com> Date: Thu, 30 Jul 2026 09:06:17 +0800 Subject: [PATCH 2/2] fix: surface vllm rerank response failures --- .../provider/sources/vllm_rerank_source.py | 33 +++-- tests/test_vllm_rerank_source.py | 129 ++++++++++++++++++ 2 files changed, 149 insertions(+), 13 deletions(-) create mode 100644 tests/test_vllm_rerank_source.py diff --git a/astrbot/core/provider/sources/vllm_rerank_source.py b/astrbot/core/provider/sources/vllm_rerank_source.py index e5ed791160..6fd44bc589 100644 --- a/astrbot/core/provider/sources/vllm_rerank_source.py +++ b/astrbot/core/provider/sources/vllm_rerank_source.py @@ -1,7 +1,5 @@ import aiohttp -from astrbot import logger - from ..entities import ProviderType, RerankResult from ..provider import RerankProvider from ..register import register_provider_adapter @@ -42,6 +40,9 @@ async def rerank( documents: list[str], top_n: int | None = None, ) -> list[RerankResult]: + if not documents: + return [] + payload = { "query": query, "documents": documents, @@ -55,21 +56,27 @@ async def rerank( rerank_url, json=payload, ) as response: + response.raise_for_status() response_data = await response.json() - results = response_data.get("results", []) + if not isinstance(response_data, dict): + raise ValueError("Rerank API response must be a JSON object") - if not results: - logger.warning( - f"Rerank API 返回了空的列表数据。原始响应: {response_data}", + results = response_data.get("results") + if not isinstance(results, list) or not results: + raise ValueError( + "Rerank API response must contain a non-empty 'results' list" ) - return [ - RerankResult( - index=result["index"], - relevance_score=result["relevance_score"], - ) - for result in results - ] + try: + return [ + RerankResult( + index=result["index"], + relevance_score=result["relevance_score"], + ) + for result in results + ] + except (KeyError, TypeError) as exc: + raise ValueError("Rerank API returned invalid result data") from exc async def terminate(self) -> None: """关闭客户端会话""" diff --git a/tests/test_vllm_rerank_source.py b/tests/test_vllm_rerank_source.py new file mode 100644 index 0000000000..cc86013adc --- /dev/null +++ b/tests/test_vllm_rerank_source.py @@ -0,0 +1,129 @@ +from unittest.mock import Mock + +import aiohttp +import pytest + +from astrbot.core.provider.sources.vllm_rerank_source import VLLMRerankProvider + + +class FakeResponse: + def __init__(self, response_data, status: int = 200) -> None: + self.response_data = response_data + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + def raise_for_status(self) -> None: + if self.status >= 400: + raise aiohttp.ClientResponseError( + request_info=Mock(), + history=(), + status=self.status, + ) + + async def json(self): + return self.response_data + + +class FakeClient: + def __init__(self, response: FakeResponse) -> None: + self.response = response + self.closed = False + self.requests = [] + + def post(self, url: str, json: dict) -> FakeResponse: + self.requests.append((url, json)) + return self.response + + async def close(self) -> None: + self.closed = True + + +@pytest.fixture +def provider() -> VLLMRerankProvider: + instance = VLLMRerankProvider.__new__(VLLMRerankProvider) + instance.base_url = "https://rerank.example.test" + instance.api_suffix = "/v1/rerank" + instance.model = "test-model" + instance.client = None + return instance + + +@pytest.mark.asyncio +async def test_vllm_rerank_maps_successful_response(provider): + provider.client = FakeClient( + FakeResponse( + { + "results": [ + {"index": 1, "relevance_score": 0.9}, + {"index": 0, "relevance_score": 0.7}, + ] + } + ) + ) + + results = await provider.rerank("query", ["first", "second"], top_n=2) + + assert [(result.index, result.relevance_score) for result in results] == [ + (1, 0.9), + (0, 0.7), + ] + assert provider.client.requests == [ + ( + "https://rerank.example.test/v1/rerank", + { + "query": "query", + "documents": ["first", "second"], + "model": "test-model", + "top_n": 2, + }, + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [400, 429]) +async def test_vllm_rerank_raises_for_http_errors(provider, status): + provider.client = FakeClient(FakeResponse({"error": "request failed"}, status)) + + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + await provider.rerank("query", ["document"]) + + assert exc_info.value.status == status + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response_data", + [None, {}, {"results": None}, {"results": []}, {"results": [{}]}], +) +async def test_vllm_rerank_raises_for_invalid_responses(provider, response_data): + provider.client = FakeClient(FakeResponse(response_data)) + + with pytest.raises(ValueError): + await provider.rerank("query", ["document"]) + + +@pytest.mark.asyncio +async def test_vllm_rerank_skips_request_for_empty_documents(provider): + provider.client = FakeClient(FakeResponse({"results": []})) + + results = await provider.rerank("query", []) + + assert results == [] + assert provider.client.requests == [] + + +@pytest.mark.asyncio +async def test_vllm_rerank_terminate_closes_session(provider): + client = FakeClient(FakeResponse({"results": []})) + provider.client = client + + await provider.terminate() + + assert client.closed is True + assert provider.client is None