AI changes and handling improvements