Updated class with refactored package