Skip to content
Closed
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
2 changes: 1 addition & 1 deletion astrbot/core/config/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -1826,7 +1826,7 @@ class ChatProviderTemplate(TypedDict):
"rerank_api_base": {
"description": "重排序模型 API Base URL",
"type": "string",
"hint": "AstrBot 会在请求时在末尾加上 /v1/rerank。",
"hint": "纯域名地址会自动补 /v1/rerank;带路径或完整 rerank 端点的地址保持原样。带路径时请包含 http:// 或 https://。",
},
"rerank_api_key": {
"description": "API Key",
Expand Down
45 changes: 42 additions & 3 deletions astrbot/core/provider/sources/vllm_rerank_source.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from urllib.parse import urlsplit

import aiohttp

from astrbot import logger
Expand All @@ -13,19 +15,43 @@
provider_type=ProviderType.RERANK,
)
class VLLMRerankProvider(RerankProvider):
@staticmethod
def _resolve_rerank_endpoint(base_url: str) -> str:
normalized = base_url.strip().removesuffix("/")
if normalized.endswith("/rerank"):
return normalized
if normalized.endswith("/v1"):
return f"{normalized}/rerank"

has_scheme = "://" in normalized
parsed = urlsplit(normalized if has_scheme else f"//{normalized}")
if not has_scheme:
if parsed.path not in ("", "/"):
raise ValueError(
"VLLM Rerank API Base URL must include a scheme when a path is provided: "
f"{base_url!r}"
)
normalized = f"http://{normalized}"
parsed = urlsplit(normalized)

if not parsed.path or parsed.path == "/":
return f"{normalized}/v1/rerank"
return normalized

def __init__(self, provider_config: dict, provider_settings: dict) -> None:
super().__init__(provider_config, provider_settings)
self.provider_config = provider_config
self.provider_settings = provider_settings
self.auth_key = provider_config.get("rerank_api_key", "")
self.base_url = provider_config.get("rerank_api_base", "http://127.0.0.1:8000")
self.base_url = self.base_url.rstrip("/")
base_url = provider_config.get("rerank_api_base", "http://127.0.0.1:8000")
self.endpoint_url = self._resolve_rerank_endpoint(base_url)
self.timeout = provider_config.get("timeout", 20)
self.model = provider_config.get("rerank_model", "BAAI/bge-reranker-base")

h = {}
if self.auth_key:
h["Authorization"] = f"Bearer {self.auth_key}"
logger.info(f"[vLLM Rerank] Using API URL: {self.endpoint_url}")
self.client = aiohttp.ClientSession(
headers=h,
timeout=aiohttp.ClientTimeout(total=self.timeout),
Expand All @@ -46,10 +72,23 @@ async def rerank(
payload["top_n"] = top_n
assert self.client is not None
async with self.client.post(
f"{self.base_url}/v1/rerank",
self.endpoint_url,
json=payload,
) as response:
response_data = await response.json()
if isinstance(response_data, dict) and "error" in response_data:
error = response_data["error"]
if isinstance(error, dict):
code = error.get("code", "unknown")
message = error.get("message", "Unknown rerank API error")
raise ValueError(f"Rerank API error {code}: {message}")
raise ValueError(f"Rerank API error: {error}")
response.raise_for_status()
if not isinstance(response_data, dict):
raise ValueError(
"Unexpected rerank API response format: "
f"{type(response_data).__name__}"
)
results = response_data.get("results", [])

if not results:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1075,7 +1075,7 @@
},
"rerank_api_base": {
"description": "Rerank Model API Base URL",
"hint": "AstrBot appends /v1/rerank to the request URL."
"hint": "AstrBot appends /v1/rerank only when the API Base URL has no path; URLs with paths or a full rerank endpoint are preserved as-is. Include http:// or https:// when a path is provided."
},
"rerank_api_key": {
"description": "API Key",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1076,7 +1076,7 @@
},
"rerank_api_base": {
"description": "Base URL API модели Rerank",
"hint": "AstrBot добавляет /v1/rerank к URL запроса."
"hint": "AstrBot добавляет /v1/rerank только если API Base URL не содержит пути; URL с путями или полный rerank endpoint сохраняются без изменений. Если указан путь, добавьте http:// или https://."
},
"rerank_api_key": {
"description": "API Key",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1077,7 +1077,7 @@
},
"rerank_api_base": {
"description": "重排序模型 API Base URL",
"hint": "AstrBot 会在请求时在末尾加上 /v1/rerank。"
"hint": "纯域名地址会自动补 /v1/rerank;带路径或完整 rerank 端点的地址保持原样。带路径时请包含 http:// 或 https://。"
},
"rerank_api_key": {
"description": "API Key",
Expand Down
2 changes: 2 additions & 0 deletions docs/en/use/knowledge-base.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ A reranker model can improve the precision of final retrieval results to some ex

Similar to configuring the embedding model, open the service provider page, click "Add Service Provider", and select Reranker. For more information about reranker models, please refer to online resources.

For vLLM and other HTTP rerank services, when `rerank api base` only contains the host, AstrBot automatically appends `/v1/rerank`. If the URL already contains a path or is already a full rerank endpoint, AstrBot preserves it as-is. When a path is provided, include `http://` or `https://`.

## Creating a Knowledge Base

AstrBot supports multiple knowledge base management. During chat, you can **freely specify which knowledge base to use**.
Expand Down
2 changes: 2 additions & 0 deletions docs/zh/use/knowledge-base.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@

和嵌入模型的配置类似,打开服务提供商页面,点击新增服务提供商,选择重排序。有关重排序模型的更多信息请参考网络。

如果使用 vLLM 或其他兼容 HTTP Rerank 接口的服务,当 `rerank api base` 只填写域名时,AstrBot 会自动补上 `/v1/rerank`;如果填写的是带路径的地址或完整的 rerank 端点,AstrBot 会保持原样。填写带路径的地址时,请包含 `http://` 或 `https://`。

## 创建知识库

AstrBot 支持多知识库管理。在聊天时,您可以**自由指定知识库**。
Expand Down
Loading