Refactoring to add new Model class