Refactor

small changes in API