Added more code for the entity wrapper