edit command role

4 files changed, 31 insertions, 18 deletions

diff --git a/py/context.py b/py/context.py
index 254cd3b..15c5e8e 100644
--- a/py/context.py
+++ b/py/context.py
@@ -40,10 +40,12 @@ def load_role_config(role):
 
     options = roles.get(f"{role}.options", {})
     options_complete = roles.get(f"{role}.options-complete", {})
+    options_edit = roles.get(f"{role}.options-edit", {})
     options_chat = roles.get(f"{role}.options-chat", {})
 
     ui = roles.get(f"{role}.ui", {})
     ui_complete = roles.get(f"{role}.ui-complete", {})
+    ui_edit = roles.get(f"{role}.ui-edit", {})
     ui_chat = roles.get(f"{role}.ui-chat", {})
 
     return {
@@ -56,6 +58,10 @@ def load_role_config(role):
             'options': dict(options_complete),
             'ui': dict(ui_complete),
         },
+        'config_edit': {
+            'options': dict(options_edit),
+            'ui': dict(ui_edit),
+        },
         'config_chat': {
             'options': dict(options_chat),
             'ui': dict(ui_chat),
diff --git a/roles-example.ini b/roles-example.ini
index 03e970b..9b13cd4 100644
--- a/roles-example.ini
+++ b/roles-example.ini
@@ -1,5 +1,8 @@
-# .ini file structure:
-# - https://docs.python.org/3/library/configparser.html#supported-ini-file-structure
+# By default vim-ai reads this configuration file but it is highly recommended to
+# set up a custom role location in order to prevent loss of roles while re-installing plugin!
+#
+# - custom config file location: `let g:vim_ai_roles_config_file = '/path/to/my/roles.ini'`
+# - .ini file structure: https://docs.python.org/3/library/configparser.html#supported-ini-file-structure
 
 [grammar]
 prompt = fix spelling and grammar
@@ -10,16 +13,16 @@ prompt =
   please refactor it in a more clean and concise way so that my colleagues
   can maintain the code more easily. Also, explain why you want to refactor
   the code so that I can add the explanation to the Pull Request.
-
-# common options for all engines
+# common options for all commands (complete, edit, chat)
 [refactor.options]
 temperature = 0.4
-
-# engine specific options:
+# command specific options:
 [refactor.options-chat]
 model = gpt-4o
-
 [refactor.options-complete]
+model = gpt-4
+[refactor.options-edit]
+model = gpt-4
 
 [o1-mini]
 [o1-mini.options]
diff --git a/tests/context_test.py b/tests/context_test.py
index 16c6e10..869215d 100644
--- a/tests/context_test.py
+++ b/tests/context_test.py
@@ -67,34 +67,36 @@ def test_role_config():
     assert 'simple role prompt:\nuser instruction:\nselected text' == actual_prompt
 
 def test_role_config_different_commands():
-    context  = make_ai_context({
+    base = {
         'config_default': default_config,
         'config_extension': {},
         'user_instruction': '/test-role hello',
         'user_selection': '',
-        'command_type': 'chat',
-    })
+    }
+    context  = make_ai_context({ **base, 'command_type': 'chat' })
     actual_config = context['config']
     actual_prompt = context['prompt']
     assert 'model-common' == actual_config['options']['model']
-    assert 'https://localhost/chat' == actual_config['options']['endpoint_url']
     assert '0' == actual_config['ui']['paste_mode']
     assert 'preset_tab' == actual_config['ui']['open_chat_command']
     assert 'hello' == actual_prompt
+    assert 'https://localhost/chat' == actual_config['options']['endpoint_url']
 
-    context = make_ai_context({
-        'config_default': default_config,
-        'config_extension': {},
-        'user_instruction': '/test-role hello',
-        'user_selection': '',
-        'command_type': 'complete',
-    })
+    context  = make_ai_context({ **base, 'command_type': 'complete' })
     actual_config = context['config']
     actual_prompt = context['prompt']
     assert 'model-common' == actual_config['options']['model']
+    assert '0' == actual_config['ui']['paste_mode']
+    assert 'hello' == actual_prompt
     assert 'https://localhost/complete' == actual_config['options']['endpoint_url']
+
+    context  = make_ai_context({ **base, 'command_type': 'edit' })
+    actual_config = context['config']
+    actual_prompt = context['prompt']
+    assert 'model-common' == actual_config['options']['model']
     assert '0' == actual_config['ui']['paste_mode']
     assert 'hello' == actual_prompt
+    assert 'https://localhost/edit' == actual_config['options']['endpoint_url']
 
 def test_multiple_role_configs():
     context = make_ai_context({
diff --git a/tests/resources/roles.ini b/tests/resources/roles.ini
index c77d272..66b77ca 100644
--- a/tests/resources/roles.ini
+++ b/tests/resources/roles.ini
@@ -10,6 +10,8 @@ model = model-common
 endpoint_url = https://localhost/chat
 [test-role.options-complete]
 endpoint_url = https://localhost/complete
+[test-role.options-edit]
+endpoint_url = https://localhost/edit
 [test-role.ui]
 paste_mode = 0
 [test-role.ui-chat]