cbrkit.retrieval.rerank.http

 1"""Re-ranking via an HTTP ``/rerank`` endpoint."""
 2
 3from collections.abc import Iterable, Mapping
 4from dataclasses import dataclass, field
 5from typing import override
 6
 7import httpx
 8from pydantic import BaseModel
 9
10from ._common import RerankFunc
11
12
13class _RerankRequest(BaseModel):
14    model: str
15    query: str
16    documents: list[str]
17    top_n: int | None = None
18
19
20class _RerankResult(BaseModel):
21    index: int
22    relevance_score: float
23
24
25class _RerankResponse(BaseModel):
26    results: list[_RerankResult]
27
28
29@dataclass(slots=True)
30class http[K](RerankFunc[K]):
31    """Re-ranking via an HTTP ``/rerank`` endpoint.
32
33    The endpoint receives ``{"model", "query", "documents"}`` and returns
34    ``{"results": [{"index", "relevance_score"}]}``.
35
36    Args:
37        model: Name of the reranker model served by the endpoint.
38        url: Absolute URL, or a path resolved against the client's base URL.
39        client: HTTP client used for transport, auth, and timeouts.
40        api_key: Optional bearer token sent as an ``Authorization`` header.
41            Kept separate from ``client`` so a shared transport can stay
42            auth-neutral while each reranker carries its own token. Excluded
43            from ``repr`` so the secret is not leaked when printing instances.
44        extra_headers: Optional extra headers sent with every request.
45            ``api_key`` takes precedence over an ``Authorization`` entry here.
46        top_n: Optional cap on the number of results the server returns.
47    """
48
49    model: str
50    url: str
51    client: httpx.AsyncClient = field(
52        default_factory=httpx.AsyncClient, repr=False, compare=False
53    )
54    api_key: str | None = field(default=None, repr=False)
55    extra_headers: Mapping[str, str] | None = None
56    top_n: int | None = None
57
58    @property
59    def _request_headers(self) -> dict[str, str]:
60        headers = dict(self.extra_headers or {})
61        if self.api_key:
62            headers["Authorization"] = f"Bearer {self.api_key}"
63        return headers
64
65    @override
66    async def _rerank(
67        self, query: str, documents: list[str]
68    ) -> Iterable[tuple[int, float]]:
69        body = _RerankRequest(
70            model=self.model, query=query, documents=documents, top_n=self.top_n
71        )
72        response = await self.client.post(
73            self.url,
74            json=body.model_dump(exclude_none=True),
75            headers=self._request_headers,
76        )
77        response.raise_for_status()
78        data = _RerankResponse.model_validate_json(response.content)
79        return ((result.index, result.relevance_score) for result in data.results)
80
81
82__all__ = ["http"]