Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
20 changes: 18 additions & 2 deletions astrbot/core/provider/sources/bailian_rerank_source.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
from typing import Any
from urllib.parse import urlsplit

import aiohttp

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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
}
Expand All @@ -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")
Expand Down
144 changes: 144 additions & 0 deletions tests/test_bailian_rerank_source.py
Original file line number Diff line number Diff line change
@@ -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
Loading