API/models/schemas.py

from pydantic import BaseModel
from typing import Dict

class PredictionResponse(BaseModel):
    predictions: Dict[str, float]
    message: str = "Success"

class BatchPredictionResponse(BaseModel):
    results: Dict[str, Dict[str, float]]
    message: str = "Success"

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