diff options
Diffstat (limited to 'py')
| -rw-r--r-- | py/context.py | 3 | ||||
| -rw-r--r-- | py/roles.py | 10 |
2 files changed, 9 insertions, 4 deletions
diff --git a/py/context.py b/py/context.py index f1eacd6..87c3a14 100644 --- a/py/context.py +++ b/py/context.py @@ -92,7 +92,8 @@ def load_role_config(role): enhance_roles_with_custom_function(roles) - if not role in roles: + postfixes = ["", ".complete", ".edit", ".chat"] + if not any([f"{role}{postfix}" in roles for postfix in postfixes]): raise Exception(f"Role `{role}` not found") if is_deprecated_role_syntax(roles, role): 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) |