app/rag/reranker.py

import logging
import re
from functools import lru_cache
from typing import Any

logger = logging.getLogger(__name__)


@lru_cache
def get_flashrank_ranker():
    from flashrank import Ranker

    return Ranker(max_length=256)


class FlashRankReranker:
    def __init__(self) -> None:
        self._ranker = None

    @property
    def ranker(self):
        if self._ranker is None:
            self._ranker = get_flashrank_ranker()
        return self._ranker

    def rerank(self, query: str, docs: list[dict[str, Any]], top_k: int) -> list[dict[str, Any]]:
        if not docs:
            return []
        try:
            from flashrank import RerankRequest

            passages = [
                {
                    "id": doc["id"],
                    "text": doc["text"],
                    "meta": {**doc.get("metadata", {}), "source_name": doc.get("source_name")},
                }
                for doc in docs
            ]
            result = self.ranker.rerank(RerankRequest(query=query, passages=passages))
            by_id = {doc["id"]: doc for doc in docs}
            reranked = []
            seen = set()
            for item in result:
                doc = by_id[str(item["id"])]
                fingerprint = self._fingerprint(doc.get("text", ""))
                if fingerprint in seen:
                    continue
                seen.add(fingerprint)
                reranked.append(
                    {
                        **doc,
                        "score": float(item.get("score", doc.get("score", 0.0))),
                        "metadata": {**doc.get("metadata", {}), "reranker": "flashrank"},
                    }
                )
                if len(reranked) >= top_k:
                    break
            return reranked
        except Exception as exc:  # pragma: no cover - fallback for missing model cache
            logger.warning("FlashRank fallback used: %s", exc)
            return self._dedupe(
                sorted(docs, key=lambda d: d.get("score", d.get("rrf_score", 0.0)), reverse=True),
                top_k=top_k,
            )

    def _dedupe(self, docs: list[dict[str, Any]], top_k: int) -> list[dict[str, Any]]:
        selected = []
        seen = set()
        for doc in docs:
            fingerprint = self._fingerprint(doc.get("text", ""))
            if fingerprint in seen:
                continue
            seen.add(fingerprint)
            selected.append(doc)
            if len(selected) >= top_k:
                break
        return selected

    def _fingerprint(self, text: str) -> str:
        normalized = re.sub(r"\s+", " ", text.lower()).strip()
        return normalized[:500]