Switch to use Pydantic models