43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
import httpx
|
|
|
|
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class VLLMReranker(RerankerBase): # pragma: no cover
|
|
def __init__(self, model: str, base_url: str):
|
|
self._model = model
|
|
self._base_url = base_url
|
|
|
|
async def rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
if not chunks:
|
|
return []
|
|
|
|
# Prepare documents for reranking
|
|
documents = [chunk.content for chunk in chunks]
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.post(
|
|
f"{self._base_url}/v1/rerank",
|
|
json={"model": self._model, "query": query, "documents": documents},
|
|
headers={
|
|
"accept": "application/json",
|
|
"Content-Type": "application/json",
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
|
|
# Extract scores and pair with chunks
|
|
scored_chunks = []
|
|
for item in result.get("results", []):
|
|
index = item["index"]
|
|
score = item["relevance_score"]
|
|
scored_chunks.append((chunks[index], score))
|
|
|
|
# Sort by score (descending) and return top_n
|
|
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
|
return scored_chunks[:top_n]
|