from typing import Any, Dict, List

from fastrag.rankers import QuantizedBiEncoderRanker


class RankerHandler:
    def __init__(self, path=""):
        model_id = "Intel/bge-large-en-v1.5-rag-int8-static"
        self.ranker = QuantizedBiEncoderRanker(model_name_or_path=model_id)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        query = data.get("query")
        documents = data.get("documents")

        top_k = data.get("top_k", None)

        assert isinstance(query, str), "Expected query to be a string"
        assert isinstance(documents, list), "Expected documents to be a list"
        assert all(isinstance(document, str) for document in documents), "Expected documents to be a list of strings"

        return self.ranker.predict(query=query, documents=documents, top_k=top_k)