command-type only roles

autoload/vim_ai.vim

@@ -299,9 +299,21 @@ function! vim_ai#AIRedoRun() abort
   endif
 endfunction
 
-function! vim_ai#RoleCompletion(A,L,P) abort
+function! s:RoleCompletion(A, command_type) abort
   call s:ImportPythonModules()
-  let l:role_list = py3eval("load_ai_role_names()")
+  let l:role_list = py3eval("load_ai_role_names(unwrap('a:command_type'))")
   call map(l:role_list, '"/" . v:val')
   return filter(l:role_list, 'v:val =~ "^' . a:A . '"')
 endfunction
+
+function! vim_ai#RoleCompletionComplete(A,L,P) abort
+  return s:RoleCompletion(a:A, 'complete')
+endfunction
+
+function! vim_ai#RoleCompletionEdit(A,L,P) abort
+  return s:RoleCompletion(a:A, 'edit')
+endfunction
+
+function! vim_ai#RoleCompletionChat(A,L,P) abort
+  return s:RoleCompletion(a:A, 'chat')
+endfunction

plugin/vim-ai.vim

@@ -4,8 +4,8 @@ if !has('python3')
   finish
 endif
 
-command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletion AI <line1>,<line2>call vim_ai#AIRun(<range>, {}, <q-args>)
-command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletion AIEdit <line1>,<line2>call vim_ai#AIEditRun(<range>, {}, <q-args>)
-command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletion AIChat <line1>,<line2>call vim_ai#AIChatRun(<range>, {}, <q-args>)
+command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletionComplete AI <line1>,<line2>call vim_ai#AIRun(<range>, {}, <q-args>)
+command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletionEdit AIEdit <line1>,<line2>call vim_ai#AIEditRun(<range>, {}, <q-args>)
+command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletionChat AIChat <line1>,<line2>call vim_ai#AIChatRun(<range>, {}, <q-args>)
 command! -nargs=? AINewChat call vim_ai#AINewChatRun(<f-args>)
 command! AIRedo call vim_ai#AIRedoRun()

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):

py/roles.py

@@ -6,7 +6,7 @@
 
 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)

tests/context_test.py

@@ -115,6 +115,17 @@ def test_multiple_role_configs():
     assert 'https://localhost/chat' == actual_config['options']['endpoint_url']
     assert 'simple role prompt:\nhello' == actual_prompt
 
+def test_chat_only_role():
+    context = make_ai_context({
+        'config_default': default_config,
+        'config_extension': {},
+        'user_instruction': '/chat-only-role',
+        'user_selection': '',
+        'command_type': 'chat',
+    })
+    actual_config = context['config']
+    assert 'preset_tab' == actual_config['options']['open_chat_command']
+
 def test_user_prompt():
     assert 'fix grammar: helo word' == make_prompt( '', 'fix grammar: helo word', '', '')
     assert 'fix grammar:\nhelo word' == make_prompt( '', 'fix grammar', 'helo word', '')

tests/resources/roles.ini

@@ -15,6 +15,9 @@ config.options.endpoint_url = https://localhost/complete
 config.engine = complete
 config.options.endpoint_url = https://localhost/edit
 
+[chat-only-role.chat]
+config.options.open_chat_command = preset_tab
+
 [deprecated-test-role-simple]
 prompt = simple role prompt
 [deprecated-test-role-simple.options]

tests/roles_test.py

@@ -1,10 +1,20 @@
 from roles import load_ai_role_names
 
 def test_role_completion():
-    role_names = load_ai_role_names()
-    assert role_names == [
+    role_names = load_ai_role_names('complete')
+    assert set(role_names) == {
         'test-role-simple',
         'test-role',
         'deprecated-test-role-simple',
         'deprecated-test-role',
-    ]
+    }
+
+def test_role_chat_only():
+    role_names = load_ai_role_names('chat')
+    assert set(role_names) == {
+        'test-role-simple',
+        'test-role',
+        'chat-only-role',
+        'deprecated-test-role-simple',
+        'deprecated-test-role',
+    }