20 lines
681 B
Python
20 lines
681 B
Python
from haiku.rag.config import Config
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class RerankerBase:
|
|
_model: str | None = Config.reranking.model.name if Config.reranking.model else None
|
|
|
|
async def rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
if not chunks:
|
|
return []
|
|
return await self._rerank(query, chunks, top_n)
|
|
|
|
async def _rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
raise NotImplementedError(
|
|
"Reranker is an abstract class. Please implement the _rerank method in a subclass."
|
|
)
|