50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import os
|
|
|
|
import httpx
|
|
|
|
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class JinaReranker(RerankerBase):
|
|
"""Jina AI reranker using the Jina Reranker API."""
|
|
|
|
def __init__(self, model: str = "jina-reranker-v3"):
|
|
self._model = model
|
|
self._api_key = os.environ.get("JINA_API_KEY")
|
|
if not self._api_key:
|
|
raise ValueError("JINA_API_KEY environment variable required")
|
|
|
|
async def rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
if not chunks:
|
|
return []
|
|
|
|
documents = [chunk.content for chunk in chunks]
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.post(
|
|
"https://api.jina.ai/v1/rerank",
|
|
json={
|
|
"model": self._model,
|
|
"query": query,
|
|
"documents": documents,
|
|
"top_n": top_n,
|
|
},
|
|
headers={
|
|
"Authorization": f"Bearer {self._api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
|
|
scored_chunks = []
|
|
for item in result.get("results", []):
|
|
index = item["index"]
|
|
score = item["relevance_score"]
|
|
scored_chunks.append((chunks[index], score))
|
|
|
|
return scored_chunks
|