from pydantic import BaseModel
from typing import List, Optional, Dict, Any


class PredictionResponse(BaseModel):
    prediction: int
    probability: Optional[List[float]] = None


class BatchPredictionResponse(BaseModel):
    predictions: List[int]
    probabilities: Optional[List[List[float]]] = None
    num_samples: int


class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    message: str


class TrainingResponse(BaseModel):
    status: str
    message: str
    metrics: Optional[Dict[str, Any]] = None
    model_path: Optional[str] = None