endpoints/rerank.py

#!/usr/bin/env python3
"""Rerank endpoint."""
from logger import setup_logger
from rerank.schemas import RerankRequest, RerankResponse
from rerank.service import rerank as rerank_service
from security.validation import validate_json

logger = setup_logger("endpoints.rerank")


async def post_rerank(request: dict) -> RerankResponse:
    req = await validate_json(RerankRequest, request)
    try:
        results = rerank_service(
            query=req.query,
            articles=req.articles,
            top_k=req.top_k,
            use_cross_encoder=bool(req.use_cross_encoder),
        )
        return RerankResponse(results=results)
    except Exception as exc:
        logger.error("Rerank error: %s", exc)
        raise