import time
from abc import ABC, abstractmethod
from typing import Optional

from pydantic import BaseModel


class RerankDocument(BaseModel):
    """
    Model class for rerank document.
    """

    index: int
    text: str
    score: float


class RerankResult(BaseModel):
    """
    Model class for rerank result.
    """

    docs: list[RerankDocument]


class RerankModel(ABC):
    """
    Base class for rerank models.
    """

    @abstractmethod
    def rerank(
        self,
        query: str,
        docs: list[str],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
    ) -> RerankResult:
        """
        Invoke rerank model

        :param model: model name
        :param credentials: model credentials
        :param query: search query
        :param docs: docs for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :param user: unique user id
        :return: rerank result
        """
        self.started_at = time.perf_counter()

        try:
            return self._invoke(query, docs, score_threshold, top_n)
        except Exception as e:
            raise self._transform_invoke_error(e)
