models/model.py

from typing import List, Optional
from pydantic import BaseModel, Field


class RerankRequest(BaseModel):
    """
    Request model for document reranking.
    
    Attributes:
        query: The search query
        documents: List of documents to rerank
        model_id: Identifier of the reranking model to use
        instruction: Optional instruction for instruction-based models
        top_k: Maximum number of documents to return (optional)
    """
    query: str = Field(..., description="Search query text")
    documents: List[str] = Field(..., min_items=1, description="List of documents to rerank")
    model_id: Optional[str] = Field(..., description="Model identifier for reranking")
    instruction: Optional[str] = Field(None, description="Optional instruction for reranking task")
    top_k: Optional[int] = Field(None, description="Maximum number of results to return")

class RerankResult(BaseModel):
    """
    Single reranking result.
    
    Attributes:
        text: The document text
        score: Relevance score from the reranking model
        index: Original index of the document in input list
    """
    text: str
    score: float
    index: int

class RerankResponse(BaseModel):
    """
    Response model for document reranking.
    
    Attributes:
        results: List of reranked documents with scores
        query: The original search query
        model_id: Identifier of the model used
        processing_time: Time taken to process the request
        total_documents: Total number of input documents
        returned_documents: Number of documents returned
    """
    results: List[RerankResult]
    query: str
    model_id: str
    processing_time: float
    total_documents: int
    returned_documents: int