148 lines
4.4 KiB
Python
148 lines
4.4 KiB
Python
"""重排序服务类型 — 请求/响应模型和协议定义。
|
|
|
|
Reranker 适配器通过 HTTP 调用重排序端点，经由 ``RerankerService``
|
|
或直接从 ``ProviderManager`` 消费。
|
|
|
|
用法::
|
|
|
|
service = RerankerService(manager)
|
|
result = await service.rerank("查询", ["文档1", "文档2"])
|
|
|
|
# 或通过工厂直接调用:
|
|
adapter = manager.create_reranker("rerank-model")
|
|
result = await adapter.rerank(RerankRequest(query="...", documents=["..."]))
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
|
|
|
from app.ai_provider.types.errors import AIProviderError
|
|
|
|
if TYPE_CHECKING:
|
|
from app.ai_provider.config.models import ModelConfig, ModelEndpointConfig
|
|
|
|
|
|
class RerankerServiceError(AIProviderError):
|
|
"""Reranker API 调用失败的基础错误。"""
|
|
|
|
|
|
# ── 请求 / 响应 ────────────────────────────────────────────────────────
|
|
|
|
|
|
@dataclass
|
|
class RerankRequest:
|
|
"""重排序请求。
|
|
|
|
属性:
|
|
query: 查询文本。
|
|
documents: 待重排序的文档列表。
|
|
model: 模型标识符。
|
|
top_n: 返回的最相关文档数量（nil 表示全部）。
|
|
max_chunks_per_doc: 每个文档最多切分为多少块。
|
|
return_documents: 返回结果中是否包含文档文本。
|
|
"""
|
|
|
|
query: str
|
|
documents: list[str]
|
|
model: str = ""
|
|
top_n: int | None = None
|
|
max_chunks_per_doc: int | None = None
|
|
return_documents: bool = False
|
|
|
|
|
|
@dataclass
|
|
class RerankResult:
|
|
"""单条重排序结果。"""
|
|
|
|
index: int
|
|
relevance_score: float
|
|
text: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class RerankResponse:
|
|
"""重排序 API 的完整响应。"""
|
|
|
|
results: list[RerankResult]
|
|
model: str = ""
|
|
usage: dict[str, int] | None = None
|
|
|
|
|
|
# ── Protocol ────────────────────────────────────────────────────────────
|
|
|
|
|
|
@runtime_checkable
|
|
class RerankerProtocol(Protocol):
|
|
"""Reranker 适配器的协议。"""
|
|
|
|
async def rerank(self, request: RerankRequest) -> RerankResponse:
|
|
"""对文档进行重排序。
|
|
|
|
异常:
|
|
RerankerServiceError: API 调用失败或响应异常时抛出。
|
|
"""
|
|
...
|
|
|
|
async def close(self) -> None:
|
|
"""释放适配器持有的 HTTP 连接。"""
|
|
...
|
|
|
|
|
|
# ── 共享工厂基类 ──────────────────────────────────────────────────────
|
|
|
|
|
|
class RerankerAdapterBase:
|
|
"""Reranker 适配器基类，提供共享的工厂方法。
|
|
|
|
子类应在 ``__init__`` 中接受 ``base_url``、``api_key``、``model``、
|
|
``request_timeout_seconds`` 参数。
|
|
"""
|
|
|
|
_adapter_type: str = "reranker"
|
|
_default_timeout: int = 60
|
|
|
|
@classmethod
|
|
def from_endpoint(
|
|
cls,
|
|
endpoint: ModelEndpointConfig,
|
|
*,
|
|
timeout_seconds: int | None = None,
|
|
) -> Any:
|
|
"""从已解析的端点配置构建适配器。"""
|
|
base_url = endpoint.base_url
|
|
api_key = endpoint.api_key
|
|
if not base_url:
|
|
raise RerankerServiceError(
|
|
status_code=500,
|
|
code="reranker_endpoint_invalid",
|
|
message="Reranker endpoint has no resolved base_url.",
|
|
details={"endpoint": endpoint.name},
|
|
)
|
|
return cls(
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
model=endpoint.model,
|
|
request_timeout_seconds=timeout_seconds
|
|
or endpoint.request_timeout_seconds
|
|
or cls._default_timeout,
|
|
)
|
|
|
|
@classmethod
|
|
def from_model_config(
|
|
cls,
|
|
config: ModelConfig,
|
|
*,
|
|
timeout_seconds: int | None = None,
|
|
) -> Any:
|
|
"""从逻辑模型配置构建适配器（使用第一个活跃端点）。"""
|
|
for ep in config.endpoints:
|
|
if ep.status == "active":
|
|
return cls.from_endpoint(ep, timeout_seconds=timeout_seconds)
|
|
raise RerankerServiceError(
|
|
status_code=500,
|
|
code="reranker_no_endpoint",
|
|
message=f"Reranker model '{config.name}' has no active endpoint.",
|
|
)
|