Refactor AI class