from enum import Enum

from pydantic import BaseModel


class ModelType(Enum):
    gpt_4o = "gpt-4o"
    gpt_4o_mini = "gpt-4o-mini"
    o1_mini = "o1-mini"
    o1_preview = "o1-preview"

class EntityData(BaseModel):
    age: int | None = None
    treatmentArea: str | None = None
    treatmentMethod: str | None = None
    location: str | None = None
    postalCode: str | None = None