Fix import breakage after aa refactoring