api/inference/schema.py

"""
Schema for the inference endpoint
"""

from pydantic import BaseModel, validator
from enum import Enum

class Sex(str, Enum):
    male = "m"
    female = "f"


class Patient(BaseModel):
    age: int
    sex: Sex

    @validator('age')
    def patient_age_must_be_legit(cls, v):
        if v < 18:
            raise ValueError('age must be at least 18')
        if v > 100:
            raise ValueError('age must be at most 100')
        return v
    
class InferenceRequest(BaseModel):

    model_name: str
    patient: Patient