structure.py

from pydantic import BaseModel, Field
from typing import List

class PredictionEntry(BaseModel):
    model: str = Field(..., description="Name of the model used for prediction")
    predicted_class: int = Field(..., description="Predicted class index")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Prediction confidence (0-1)")

class ImagePredictionResponse(BaseModel):
    predictions: List[PredictionEntry] = Field(..., description="List of predictions with model, class, and confidence")

class TextPredictionRequest(BaseModel):
    text: str = Field(..., example="This is a sample text to classify")


class TextPredictionResponse(BaseModel):
    predicted_class: str = Field(..., description="Predicted label for the input text")
    confidence: float = Field(..., ge=0.0, le=1.0)