new role syntax

9 files changed, 223 insertions, 81 deletions

diff --git a/README.md b/README.md
index c1cac6b..d8df479 100644
--- a/README.md
+++ b/README.md
@@ -196,20 +196,14 @@ let g:vim_ai_roles_config_file = '/path/to/my/roles.ini'
 
 [grammar]
 prompt = fix spelling and grammar
-
-[grammar.options]
-temperature = 0.4
-
-[grammar]
-prompt = fix spelling and grammar
+config.options.temperature = 0.4
 
 [o1-mini]
-[o1-mini.options]
-stream = 0
-model = o1-mini
-max_completion_tokens = 25000
-temperature = 1
-initial_prompt =
+config.options.stream = 0
+config.options.model = o1-mini
+config.options.max_completion_tokens = 25000
+config.options.temperature = 1
+config.options.initial_prompt =
 ```
 
 Now you can select text and run it with command `:AIEdit /grammar`.
@@ -429,22 +423,19 @@ Then you set up a custom role that points to the OpenRouter endpoint:
 
 ```ini
 [gemini]
-[gemini.options]
-token_file_path = ~/.config/vim-ai-openrouter.token
-endpoint_url = https://openrouter.ai/api/v1/chat/completions
-model = google/gemini-exp-1121:free
+config.options.token_file_path = ~/.config/vim-ai-openrouter.token
+config.options.endpoint_url = https://openrouter.ai/api/v1/chat/completions
+config.options.model = google/gemini-exp-1121:free
 
 [llama]
-[llama.options]
-token_file_path = ~/.config/vim-ai-openrouter.token
-endpoint_url = https://openrouter.ai/api/v1/chat/completions
-model = meta-llama/llama-3.3-70b-instruct
+config.options.token_file_path = ~/.config/vim-ai-openrouter.token
+config.options.endpoint_url = https://openrouter.ai/api/v1/chat/completions
+config.options.model = meta-llama/llama-3.3-70b-instruct
 
 [claude]
-[claude.options]
-token_file_path = ~/.config/vim-ai-openrouter.token
-endpoint_url = https://openrouter.ai/api/v1/chat/completions
-model = anthropic/claude-3.5-haiku
+config.options.token_file_path = ~/.config/vim-ai-openrouter.token
+config.options.endpoint_url = https://openrouter.ai/api/v1/chat/completions
+config.options.model = anthropic/claude-3.5-haiku
 ```
 
 Now you can use the role:
diff --git a/doc/vim-ai.txt b/doc/vim-ai.txt
index 241e33b..03c7e4d 100644
--- a/doc/vim-ai.txt
+++ b/doc/vim-ai.txt
@@ -198,10 +198,7 @@ Example of a role: >
 
     [grammar]
     prompt = fix spelling and grammar
-
-    [grammar.options]
-    temperature = 0.4
-
+    config.options.temperature = 0.4
 
 Now you can select text and run it with command `:AIEdit /grammar`.
 See roles-example.ini for more examples.
diff --git a/py/context.py b/py/context.py
index 15c5e8e..f1eacd6 100644
--- a/py/context.py
+++ b/py/context.py
@@ -24,6 +24,63 @@ def merge_deep(objects):
         merge_deep_recursive(result, o)
     return result
 
+def is_deprecated_role_syntax(roles, role):
+    deprecated_sections = [
+        'options', 'options-complete', 'options-edit', 'options-chat',
+        'ui', 'ui-complete', 'ui-edit', 'ui-chat',
+    ]
+    for section in deprecated_sections:
+        if f"{role}.{section}" in roles:
+            return True
+    return False
+
+def load_roles_with_deprecated_syntax(roles, role):
+    prompt = dict(roles[role]).get('prompt', '')
+    return {
+        'role_default': {
+            'prompt': prompt,
+            'config': {
+                'options': dict(roles.get(f"{role}.options", {})),
+                'ui': dict(roles.get(f"{role}.ui", {})),
+            },
+        },
+        'role_complete': {
+            'prompt': prompt,
+            'config': {
+                'options': dict(roles.get(f"{role}.options-complete", {})),
+                'ui': dict(roles.get(f"{role}.ui-complete", {})),
+            },
+        },
+        'role_edit': {
+            'prompt': prompt,
+            'config': {
+                'options': dict(roles.get(f"{role}.options-edit", {})),
+                'ui': dict(roles.get(f"{role}.ui-edit", {})),
+            },
+        },
+        'role_chat': {
+            'prompt': prompt,
+            'config': {
+                'options': dict(roles.get(f"{role}.options-chat", {})),
+                'ui': dict(roles.get(f"{role}.ui-chat", {})),
+            },
+        },
+    }
+
+def parse_role_section(role):
+    result = {}
+    for key in role.keys():
+        parts = key.split('.')
+        structure = parts[:-1]
+        primitive = parts[-1]
+        obj = result
+        for path in structure:
+            if not path in obj:
+                obj[path] = {}
+            obj = obj[path]
+        obj[primitive] = role.get(key)
+    return result
+
 def load_role_config(role):
     roles_config_path = os.path.expanduser(vim.eval("g:vim_ai_roles_config_file"))
     if not os.path.exists(roles_config_path):
@@ -38,34 +95,14 @@ def load_role_config(role):
     if not role in roles:
         raise Exception(f"Role `{role}` not found")
 
-    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", {})
+    if is_deprecated_role_syntax(roles, role):
+        return load_roles_with_deprecated_syntax(roles, role)
 
     return {
-        'role': dict(roles[role]),
-        'config_default': {
-            'options': dict(options),
-            'ui': dict(ui),
-        },
-        'config_complete': {
-            '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),
-        },
+        'role_default': parse_role_section(roles.get(role, {})),
+        'role_complete': parse_role_section(roles.get(f"{role}.complete", {})),
+        'role_edit': parse_role_section(roles.get(f"{role}.edit", {})),
+        'role_chat': parse_role_section(roles.get(f"{role}.chat", {})),
     }
 
 def parse_role_names(prompt):
@@ -87,9 +124,11 @@ def parse_prompt_and_role_config(user_instruction, command_type):
     last_role = roles[-1]
     user_prompt = user_instruction[user_instruction.index(last_role) + len(last_role):].strip() # strip roles
 
-    role_configs = merge_deep([load_role_config(role) for role in roles])
-    config = merge_deep([role_configs['config_default'], role_configs['config_' + command_type]])
-    role_prompt = role_configs['role'].get('prompt', '')
+    parsed_role = merge_deep([load_role_config(role) for role in roles])
+    role_default = parsed_role['role_default']
+    role_command = parsed_role['role_' + command_type]
+    config = merge_deep([role_default.get('config', {}), role_command.get('config', {})])
+    role_prompt = role_default.get('prompt') or role_command.get('prompt', '')
     return user_prompt, role_prompt, config
 
 def make_selection_prompt(user_selection, user_prompt, role_prompt, selection_boundary):
diff --git a/py/roles.py b/py/roles.py
index 16aa4e9..37e5b4d 100644
--- a/py/roles.py
+++ b/py/roles.py
@@ -1,4 +1,8 @@
 import vim
+import os
+
+if "PYTEST_VERSION" in os.environ:
+    from utils import *
 
 roles_py_imported = True
 
diff --git a/roles-example.ini b/roles-example.ini
index 9b13cd4..6134aca 100644
--- a/roles-example.ini
+++ b/roles-example.ini
@@ -7,27 +7,25 @@
 [grammar]
 prompt = fix spelling and grammar
 
+# common options for all commands (complete, edit, chat)
 [refactor]
 prompt =
   You are a Clean Code expert, I have the following code,
   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 commands (complete, edit, chat)
-[refactor.options]
-temperature = 0.4
+config.options.temperature = 0.4
 # command specific options:
-[refactor.options-chat]
-model = gpt-4o
-[refactor.options-complete]
-model = gpt-4
-[refactor.options-edit]
-model = gpt-4
+[refactor.chat]
+config.options.model = gpt-4o
+[refactor.complete]
+config.options.model = gpt-4
+[refactor.edit]
+config.options.model = gpt-4
 
 [o1-mini]
-[o1-mini.options]
-stream = 0
-model = o1-mini
-max_completion_tokens = 25000
-temperature = 1
-initial_prompt =
+config.options.stream = 0
+config.options.model = o1-mini
+config.options.max_completion_tokens = 25000
+config.options.temperature = 1
+config.options.initial_prompt =
diff --git a/tests/context_test.py b/tests/context_test.py
index 869215d..1e179f1 100644
--- a/tests/context_test.py
+++ b/tests/context_test.py
@@ -1,7 +1,7 @@
-import vim
 from context import make_ai_context, make_prompt
 
 default_config = {
+  "engine": "chat",
   "options": {
     "model": "gpt-4o",
     "endpoint_url": "https://api.openai.com/v1/chat/completions",
@@ -81,6 +81,7 @@ def test_role_config_different_commands():
     assert 'preset_tab' == actual_config['ui']['open_chat_command']
     assert 'hello' == actual_prompt
     assert 'https://localhost/chat' == actual_config['options']['endpoint_url']
+    assert 'chat' == actual_config['engine']
 
     context  = make_ai_context({ **base, 'command_type': 'complete' })
     actual_config = context['config']
@@ -89,6 +90,7 @@ def test_role_config_different_commands():
     assert '0' == actual_config['ui']['paste_mode']
     assert 'hello' == actual_prompt
     assert 'https://localhost/complete' == actual_config['options']['endpoint_url']
+    assert 'complete' == actual_config['engine']
 
     context  = make_ai_context({ **base, 'command_type': 'edit' })
     actual_config = context['config']
@@ -97,6 +99,7 @@ def test_role_config_different_commands():
     assert '0' == actual_config['ui']['paste_mode']
     assert 'hello' == actual_prompt
     assert 'https://localhost/edit' == actual_config['options']['endpoint_url']
+    assert 'complete' == actual_config['engine']
 
 def test_multiple_role_configs():
     context = make_ai_context({
diff --git a/tests/deprecated_role_syntax_test.py b/tests/deprecated_role_syntax_test.py
new file mode 100644
index 0000000..19744b9
--- /dev/null
+++ b/tests/deprecated_role_syntax_test.py
@@ -0,0 +1,83 @@
+from context import make_ai_context, make_prompt
+
+default_config = {
+  "options": {
+    "model": "gpt-4o",
+    "endpoint_url": "https://api.openai.com/v1/chat/completions",
+    "max_tokens": "0",
+    "max_completion_tokens": "0",
+    "temperature": "1",
+    "request_timeout": "20",
+    "stream": "1",
+    "enable_auth": "1",
+    "token_file_path": "",
+    "selection_boundary": "",
+    "initial_prompt": "You are a general assistant.",
+  },
+  "ui": {
+    "open_chat_command": "preset_below",
+    "scratch_buffer_keep_open": "0",
+    "populate_options": "0",
+    "code_syntax_enabled": "1",
+    "paste_mode": "1",
+  },
+}
+
+def test_role_config():
+    context = make_ai_context({
+        'config_default': default_config,
+        'config_extension': {},
+        'user_instruction': '/deprecated-test-role-simple user instruction',
+        'user_selection': 'selected text',
+        'command_type': 'chat',
+    })
+    actual_config = context['config']
+    actual_prompt = context['prompt']
+    assert 'o1-preview' == actual_config['options']['model']
+    assert 'simple role prompt:\nuser instruction:\nselected text' == actual_prompt
+
+def test_role_config_different_commands():
+    base = {
+        'config_default': default_config,
+        'config_extension': {},
+        'user_instruction': '/deprecated-test-role hello',
+        'user_selection': '',
+    }
+    context  = make_ai_context({ **base, 'command_type': 'chat' })
+    actual_config = context['config']
+    actual_prompt = context['prompt']
+    assert 'model-common' == actual_config['options']['model']
+    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({ **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({
+        'config_default': default_config,
+        'config_extension': {},
+        'user_instruction': '/deprecated-test-role /deprecated-test-role-simple hello',
+        'user_selection': '',
+        'command_type': 'chat',
+    })
+    actual_config = context['config']
+    actual_prompt = context['prompt']
+    assert 'o1-preview' == actual_config['options']['model']
+    assert 'https://localhost/chat' == actual_config['options']['endpoint_url']
+    assert 'simple role prompt:\nhello' == actual_prompt
diff --git a/tests/resources/roles.ini b/tests/resources/roles.ini
index 66b77ca..16335d3 100644
--- a/tests/resources/roles.ini
+++ b/tests/resources/roles.ini
@@ -1,18 +1,35 @@
 [test-role-simple]
 prompt = simple role prompt
-[test-role-simple.options]
-model = o1-preview
+config.options.model = o1-preview
 
 [test-role]
-[test-role.options]
+config.options.model = model-common
+config.ui.paste_mode = 0
+[test-role.chat]
+config.options.endpoint_url = https://localhost/chat
+config.ui.open_chat_command = preset_tab
+[test-role.complete]
+config.engine = complete
+config.options.endpoint_url = https://localhost/complete
+[test-role.edit]
+config.engine = complete
+config.options.endpoint_url = https://localhost/edit
+
+[deprecated-test-role-simple]
+prompt = simple role prompt
+[deprecated-test-role-simple.options]
+model = o1-preview
+
+[deprecated-test-role]
+[deprecated-test-role.options]
 model = model-common
-[test-role.options-chat]
+[deprecated-test-role.options-chat]
 endpoint_url = https://localhost/chat
-[test-role.options-complete]
+[deprecated-test-role.options-complete]
 endpoint_url = https://localhost/complete
-[test-role.options-edit]
+[deprecated-test-role.options-edit]
 endpoint_url = https://localhost/edit
-[test-role.ui]
+[deprecated-test-role.ui]
 paste_mode = 0
-[test-role.ui-chat]
+[deprecated-test-role.ui-chat]
 open_chat_command = preset_tab
diff --git a/tests/roles_test.py b/tests/roles_test.py
new file mode 100644
index 0000000..3230329
--- /dev/null
+++ b/tests/roles_test.py
@@ -0,0 +1,10 @@
+from roles import load_ai_role_names
+
+def test_role_completion():
+    role_names = load_ai_role_names()
+    assert role_names == [
+        'test-role-simple',
+        'test-role',
+        'deprecated-test-role-simple',
+        'deprecated-test-role',
+    ]