136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
"""重排序服务 — Reranker 适配器外观。
|
|
|
|
用法::
|
|
|
|
from app.ai_provider.reranker import RerankerService
|
|
|
|
service = RerankerService(manager)
|
|
result = await service.rerank("查询", ["文档1", "文档2"])
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from app.ai_provider.reranker.base import (
|
|
RerankerProtocol,
|
|
RerankerServiceError,
|
|
RerankRequest,
|
|
RerankResponse,
|
|
RerankResult,
|
|
)
|
|
from app.ai_provider.types.errors import AIProviderError
|
|
|
|
|
|
class RerankerService:
|
|
"""外观类，负责解析 Reranker 模型并委托给 ProviderManager。
|
|
|
|
接受 ProviderManager 或实现了以下接口的适配器：
|
|
- ``.model_registry`` → 具有 ``list_by_capability(capability)`` 方法
|
|
- ``.create_reranker(name)`` → 返回 RerankerProtocol 实例
|
|
"""
|
|
|
|
def __init__(self, manager) -> None:
|
|
self._manager = manager
|
|
self._adapter: RerankerProtocol | None = None
|
|
|
|
# ── Registry 兼容适配 ─────────────────────────────────────────
|
|
|
|
@property
|
|
def _registry(self):
|
|
mgr = self._manager
|
|
if hasattr(mgr, "model_registry"):
|
|
return mgr.model_registry
|
|
return mgr
|
|
|
|
@property
|
|
def _factory(self):
|
|
mgr = self._manager
|
|
if hasattr(mgr, "create_reranker"):
|
|
return mgr
|
|
return None
|
|
|
|
# ── Reranker ───────────────────────────────────────────────────
|
|
|
|
@property
|
|
def adapter(self) -> RerankerProtocol:
|
|
"""解析并缓存 Reranker 适配器。"""
|
|
if self._adapter is None:
|
|
self._adapter = self._resolve()
|
|
return self._adapter
|
|
|
|
async def rerank(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
model_hint: str | None = None,
|
|
top_n: int | None = None,
|
|
) -> RerankResponse:
|
|
"""快捷方法：构造请求并调用 Reranker 适配器。"""
|
|
model = self._resolve_model_name(model_hint)
|
|
request = RerankRequest(
|
|
query=query,
|
|
documents=documents,
|
|
model=model,
|
|
top_n=top_n,
|
|
)
|
|
return await self.adapter.rerank(request)
|
|
|
|
# ── 解析 ───────────────────────────────────────────────────────
|
|
|
|
def _resolve(self) -> RerankerProtocol:
|
|
"""查找第一个活跃的 Reranker 模型并通过工厂创建。"""
|
|
models = self._list_capability_models("reranker")
|
|
if not models:
|
|
raise AIProviderError(
|
|
status_code=503,
|
|
code="reranker_no_model",
|
|
message="No reranker model configured in the model registry.",
|
|
)
|
|
factory = self._factory
|
|
if factory is None:
|
|
raise AIProviderError(
|
|
status_code=501,
|
|
code="reranker_no_factory",
|
|
message="Reranker adapter factory not available (provider manager missing create_reranker).",
|
|
)
|
|
return factory.create_reranker(models[0].name)
|
|
|
|
def _resolve_model_name(self, model_hint: str | None) -> str:
|
|
models = self._list_capability_models("reranker")
|
|
if model_hint:
|
|
return model_hint
|
|
if models:
|
|
return models[0].name
|
|
return ""
|
|
|
|
def _list_capability_models(self, tag: str):
|
|
reg = self._registry
|
|
if hasattr(reg, "list_capability_models"):
|
|
return reg.list_capability_models(tag)
|
|
if hasattr(reg, "list_by_capability"):
|
|
return reg.list_by_capability(tag)
|
|
return []
|
|
|
|
# ── 生命周期 ────────────────────────────────────────────────────
|
|
|
|
async def close(self) -> None:
|
|
if self._adapter is not None:
|
|
close = getattr(self._adapter, "close", None)
|
|
if close is not None:
|
|
await close()
|
|
|
|
async def __aenter__(self) -> RerankerService:
|
|
return self
|
|
|
|
async def __aexit__(self, *exc_info: object) -> None:
|
|
await self.close()
|
|
|
|
|
|
__all__ = [
|
|
"RerankerProtocol",
|
|
"RerankRequest",
|
|
"RerankResponse",
|
|
"RerankResult",
|
|
"RerankerService",
|
|
"RerankerServiceError",
|
|
]
|