The refactor was done for version 1.x (1.4.x and 1.5.x). It was not possible to easily cherry-pick the changes to 2.x and 2.0.x. This issue is to port changes from https://git.drupalcode.org/project/ai/-/merge_requests/1536 to 2.x and 2.0.x branches.