Small architecture refactor