1 files changed, 7 insertions, 3 deletions
diff --git a/py/roles.py b/py/roles.py
index 37e5b4d..692ac6c 100644
--- a/py/roles.py
+++ b/py/roles.py
@@ -6,7 +6,7 @@ if "PYTEST_VERSION" in os.environ:
 
 roles_py_imported = True
 
-def load_ai_role_names():
+def load_ai_role_names(command_type):
     roles_config_path = os.path.expanduser(vim.eval("g:vim_ai_roles_config_file"))
     if not os.path.exists(roles_config_path):
         raise Exception(f"Role config file does not exist: {roles_config_path}")
@@ -16,6 +16,10 @@ def load_ai_role_names():
 
     enhance_roles_with_custom_function(roles)
 
-    role_names = [name for name in roles.sections() if not '.' in name]
+    role_names = set()
+    for name in roles.sections():
+        parts = name.split('.')
+        if len(parts) == 1 or parts[-1] == command_type:
+            role_names.add(parts[0])
 
-    return role_names
+    return list(role_names)