diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py index 65356e100b..030f59dd57 100644 --- a/astrbot/core/provider/sources/bailian_rerank_source.py +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -1,5 +1,6 @@ import os from typing import Any +from urllib.parse import urlsplit import aiohttp @@ -35,6 +36,10 @@ class BailianRerankProvider(RerankProvider): """阿里云百炼文本重排序适配器.""" QWEN3_RERANK_MODEL = "qwen3-rerank" + COMPATIBLE_API_PATH_SUFFIXES = ( + "/compatible-api/v1/reranks", + "/compatible-mode/v1/reranks", + ) def __init__(self, provider_config: dict, provider_settings: dict) -> None: super().__init__(provider_config, provider_settings) @@ -73,6 +78,10 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: logger.info(f"AstrBot 百炼 Rerank 初始化完成。模型: {self.model}") + def _uses_compatible_api(self) -> bool: + base_url_path = urlsplit(self.base_url).path.rstrip("/") + return base_url_path.endswith(self.COMPATIBLE_API_PATH_SUFFIXES) + def _build_payload( self, query: str, documents: list[str], top_n: int | None ) -> dict: @@ -88,8 +97,9 @@ def _build_payload( """ normalized_model = self.model.strip().lower() normalized_top_n = top_n if top_n is not None and top_n > 0 else None + is_compatible_api = self._uses_compatible_api() - if normalized_model == self.QWEN3_RERANK_MODEL: + if normalized_model == self.QWEN3_RERANK_MODEL and is_compatible_api: payload = { "model": self.model, "query": query, @@ -112,6 +122,12 @@ def _build_payload( for k, v in [ ("top_n", normalized_top_n), ("return_documents", True if self.return_documents else None), + ( + "instruct", + self.instruct + if self.instruct and normalized_model == self.QWEN3_RERANK_MODEL + else None, + ), ] if v is not None } @@ -135,7 +151,7 @@ def _parse_results(self, data: dict) -> list[RerankResult]: BailianAPIError: API返回错误 KeyError: 结果缺少必要字段 """ - is_compatible_api = "compatible-api" in self.base_url + is_compatible_api = self._uses_compatible_api() if is_compatible_api: code = data.get("code") diff --git a/tests/test_bailian_rerank_source.py b/tests/test_bailian_rerank_source.py new file mode 100644 index 0000000000..880b6e2eb3 --- /dev/null +++ b/tests/test_bailian_rerank_source.py @@ -0,0 +1,144 @@ +import pytest + +import astrbot.core.provider.sources.bailian_rerank_source as bailian_rerank_module +from astrbot.core.config.default import CONFIG_METADATA_2 +from astrbot.core.provider.sources.bailian_rerank_source import ( + BailianRerankProvider, +) + +CHINA_COMPATIBLE_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" +SINGAPORE_COMPATIBLE_URL = ( + "https://example.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1/reranks" +) +NATIVE_URL = ( + "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank" +) + + +@pytest.fixture +def provider() -> BailianRerankProvider: + instance = BailianRerankProvider.__new__(BailianRerankProvider) + instance.model = "qwen3-rerank" + instance.return_documents = False + instance.instruct = "" + return instance + + +def test_bailian_rerank_provider_preserves_native_default_endpoint(monkeypatch): + monkeypatch.setattr( + bailian_rerank_module.aiohttp, + "ClientSession", + lambda **_kwargs: object(), + ) + + provider = BailianRerankProvider( + provider_config={"rerank_api_key": "test-key"}, + provider_settings={}, + ) + + assert provider.base_url == NATIVE_URL + + +def test_bailian_rerank_config_template_preserves_native_default_endpoint(): + templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ + "config_template" + ] + + assert templates["阿里云百炼重排序"]["rerank_api_base"] == NATIVE_URL + + +def test_bailian_rerank_provider_preserves_explicit_endpoint(monkeypatch): + monkeypatch.setattr( + bailian_rerank_module.aiohttp, + "ClientSession", + lambda **_kwargs: object(), + ) + custom_url = "https://rerank.example.test/custom" + + provider = BailianRerankProvider( + provider_config={ + "rerank_api_key": "test-key", + "rerank_api_base": custom_url, + }, + provider_settings={}, + ) + + assert provider.base_url == custom_url + + +@pytest.mark.parametrize( + "base_url", + [CHINA_COMPATIBLE_URL, SINGAPORE_COMPATIBLE_URL], +) +def test_qwen3_compatible_endpoints_use_flat_payload(provider, base_url): + provider.base_url = base_url + + assert provider._build_payload("query", ["document"], top_n=1) == { + "model": "qwen3-rerank", + "query": "query", + "documents": ["document"], + "top_n": 1, + } + + +def test_qwen3_native_endpoint_uses_wrapped_payload(provider): + provider.base_url = NATIVE_URL + provider.instruct = "Focus on technical relevance." + + assert provider._build_payload("query", ["document"], top_n=1) == { + "model": "qwen3-rerank", + "input": {"query": "query", "documents": ["document"]}, + "parameters": { + "top_n": 1, + "instruct": "Focus on technical relevance.", + }, + } + + +def test_protocol_detection_ignores_compatible_text_outside_url_path(provider): + provider.base_url = f"{NATIVE_URL}?redirect=/compatible-api/v1/reranks" + + assert provider._build_payload("query", ["document"], top_n=1) == { + "model": "qwen3-rerank", + "input": {"query": "query", "documents": ["document"]}, + "parameters": {"top_n": 1}, + } + + +def test_protocol_detection_accepts_compatible_endpoint_suffix(provider): + provider.base_url = f"{CHINA_COMPATIBLE_URL}/?workspace=test" + + assert provider._build_payload("query", ["document"], top_n=1) == { + "model": "qwen3-rerank", + "query": "query", + "documents": ["document"], + "top_n": 1, + } + + +@pytest.mark.parametrize( + "base_url", + [CHINA_COMPATIBLE_URL, SINGAPORE_COMPATIBLE_URL], +) +def test_compatible_endpoints_parse_top_level_results(provider, base_url): + provider.base_url = base_url + + results = provider._parse_results( + {"results": [{"index": 0, "relevance_score": 0.75}]} + ) + + assert len(results) == 1 + assert results[0].index == 0 + assert results[0].relevance_score == 0.75 + + +def test_native_endpoint_parses_nested_results(provider): + provider.base_url = NATIVE_URL + + results = provider._parse_results( + {"output": {"results": [{"index": 0, "relevance_score": 0.75}]}} + ) + + assert len(results) == 1 + assert results[0].index == 0 + assert results[0].relevance_score == 0.75