from langchain_core.documents import Document


class RerankService:
    """Cross-Encoder 기반 재순위(Reranker) 서비스."""

    def __init__(self, model_id: str = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"):
        from sentence_transformers import CrossEncoder
        print(f"Reranker 로딩 중: {model_id}")
        self._model = CrossEncoder(model_id)
        print("Reranker 로딩 완료")

    def rerank(self, query: str, docs: list[Document], top_k: int) -> list[Document]:
        if not docs:
            return docs
        pairs = [(query, doc.page_content) for doc in docs]
        scores = self._model.predict(pairs)
        ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
        return [doc for _, doc in ranked[:top_k]]