import logging.config from enum import Enum from typing import List, Union from pydantic import BaseModel, validator logger = logging.getLogger(__name__) class PredictionResponse(BaseModel): label: str score: float class BedType(str, Enum): none = "" single = "single" double = "double" queen = "queen" king = "king" california_king = "california_king" bunk = "bunk" sofa = "sofa" rollaway = "rollaway" futon = "futon" class BedData(BaseModel): type: Union[BedType, None] = None count: Union[int, None] = None @validator("count", pre=True, always=True) def validate_count(cls, v): return v if v and v > 0 else None @validator("type", pre=True, always=True) def validate_type(cls, v): try: return BedType(v.lower()) if v else None except ValueError: return None @classmethod def create_valid_bed(cls, type: Union[str, None], count: Union[int, None]): if type and count and count > 0: try: type_enum = BedType(type.lower()) return cls(type=type_enum, count=count) except ValueError: logger.error(f"Invalid bed type: {type}") return None def __init__(self, **data): variations = { "individual": BedType.single, "camaindividual": BedType.single, "doble": BedType.double, "camadoble": BedType.double, "reina": BedType.queen, "queenbed": BedType.queen, "rey": BedType.king, "californiaking": BedType.california_king, "litera": BedType.bunk, "sofá": BedType.sofa, "sofacama": BedType.sofa, "plegable": BedType.rollaway, "futón": BedType.futon, "twin": BedType.single, "twinbed": BedType.single, "singlebed": BedType.single, "doublebed": BedType.double, "largedouble": BedType.queen, "extralargedouble": BedType.king, "bunkbed": BedType.bunk, "couch": BedType.sofa, "airmattress": BedType.futon, "floormattress": BedType.futon, } bed_type = data.get("type") if bed_type: normalized_bed_type = bed_type.replace(" ", "").lower() data["type"] = variations.get(normalized_bed_type, bed_type) super().__init__(**data) class RoomData(BaseModel): room_description: str beds: Union[BedData, None] class Predictions(BaseModel): type: PredictionResponse category: PredictionResponse environment: List[PredictionResponse] feature: List[PredictionResponse] view: List[PredictionResponse] language_detected: str beds: List[BedData] class AllPredictionsResponse(BaseModel): predictions: Predictions