unified config parsing + tests

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+env

autoload/vim_ai.vim

@@ -4,6 +4,7 @@ let s:plugin_root = expand('<sfile>:p:h:h')
 let s:complete_py = s:plugin_root . "/py/complete.py"
 let s:chat_py = s:plugin_root . "/py/chat.py"
 let s:roles_py = s:plugin_root . "/py/roles.py"
+let s:config_py = s:plugin_root . "/py/config.py"
 
 " remembers last command parameters to be used in AIRedoRun
 let s:last_is_selection = 0
@@ -45,7 +46,7 @@ function! s:OpenChatWindow(open_conf, force_new) abort
   execute l:open_cmd
 
   " reuse chat in keep-open mode
-  let l:keep_open = g:vim_ai_chat['ui']['scratch_buffer_keep_open']
+  let l:keep_open = g:vim_ai_chat['ui']['scratch_buffer_keep_open'] == '1'
   let l:last_scratch_buffer_name = s:GetLastScratchBufferName()
   if l:keep_open && bufexists(l:last_scratch_buffer_name) && !a:force_new
     let l:current_buffer = bufnr('%')
@@ -102,7 +103,7 @@ endfunction
 let s:is_handling_paste_mode = 0
 
 function! s:set_paste(config)
-  if !&paste && a:config['ui']['paste_mode']
+  if !&paste && a:config['ui']['paste_mode'] == '1'
     let s:is_handling_paste_mode = 1
     setlocal paste
   endif
@@ -151,8 +152,18 @@ endfunction
 " - config       - function scoped vim_ai_complete config
 " - a:1          - optional instruction prompt
 function! vim_ai#AIRun(uses_range, config, ...) range abort
-  let l:config = vim_ai_config#ExtendDeep(g:vim_ai_complete, a:config)
   let l:instruction = a:0 > 0 ? a:1 : ""
+  let l:config_input = {
+  \  "config_default": g:vim_ai_edit,
+  \  "config_extension": a:config,
+  \  "instruction": l:instruction,
+  \  "command_type": 'complete',
+  \}
+  execute "py3file " . s:config_py
+  execute "py3 make_config('l:config_input', 'l:config_output')"
+  let l:config = l:config_output['config']
+  let l:role_prompt = l:config_output['role_prompt']
+
   " l:is_selection used in Python script
   let l:is_selection = a:uses_range && a:firstline == line("'<") && a:lastline == line("'>")
   let l:selection = s:GetSelectionOrRange(l:is_selection, a:uses_range, a:firstline, a:lastline)
@@ -185,8 +196,18 @@ endfunction
 " - config       - function scoped vim_ai_edit config
 " - a:1          - optional instruction prompt
 function! vim_ai#AIEditRun(uses_range, config, ...) range abort
-  let l:config = vim_ai_config#ExtendDeep(g:vim_ai_edit, a:config)
   let l:instruction = a:0 > 0 ? a:1 : ""
+  let l:config_input = {
+  \  "config_default": g:vim_ai_edit,
+  \  "config_extension": a:config,
+  \  "instruction": l:instruction,
+  \  "command_type": 'complete',
+  \}
+  execute "py3file " . s:config_py
+  execute "py3 make_config('l:config_input', 'l:config_output')"
+  let l:config = l:config_output['config']
+  let l:role_prompt = l:config_output['role_prompt']
+
   " l:is_selection used in Python script
   let l:is_selection = a:uses_range && a:firstline == line("'<") && a:lastline == line("'>")
   let l:selection = s:GetSelectionOrRange(l:is_selection, a:uses_range, a:firstline, a:lastline)
@@ -248,8 +269,18 @@ endfunction
 " - config       - function scoped vim_ai_chat config
 " - a:1          - optional instruction prompt
 function! vim_ai#AIChatRun(uses_range, config, ...) range abort
-  let l:config = vim_ai_config#ExtendDeep(g:vim_ai_chat, a:config)
-  let l:instruction = ""
+  let l:instruction = a:0 > 0 ? a:1 : ""
+  let l:config_input = {
+  \  "config_default": g:vim_ai_chat,
+  \  "config_extension": a:config,
+  \  "instruction": l:instruction,
+  \  "command_type": 'chat',
+  \}
+  execute "py3file " . s:config_py
+  execute "py3 make_config('l:config_input', 'l:config_output')"
+  let l:config = l:config_output['config']
+  let l:role_prompt = l:config_output['role_prompt']
+
   " l:is_selection used in Python script
   let l:is_selection = a:uses_range && a:firstline == line("'<") && a:lastline == line("'>")
   let l:selection = s:GetSelectionOrRange(l:is_selection, a:uses_range, a:firstline, a:lastline)
@@ -260,7 +291,6 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort
 
     let l:prompt = ""
     if a:0 > 0 || a:uses_range
-      let l:instruction = a:0 > 0 ? a:1 : ""
       let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
     endif
 

py/chat.py

@@ -4,7 +4,8 @@
 plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
-prompt, config = load_config_and_prompt('chat')
+prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt"))
+config = make_config(vim.eval("l:config"))
 config_options = config['options']
 config_ui = config['ui']
 

py/complete.py

@@ -4,7 +4,8 @@
 plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
-prompt, config = load_config_and_prompt('complete')
+prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt"))
+config = make_config(vim.eval("l:config"))
 config_options = config['options']
 config_ui = config['ui']
 

py/config.py

@@ -0,0 +1,105 @@
+import vim
+import re
+import os
+import configparser
+
+def merge_deep_recursive(target, source = {}):
+    source = source.copy()
+    for key, value in source.items():
+        if isinstance(value, dict):
+            target_child = target.setdefault(key, {})
+            merge_deep_recursive(target_child, value)
+        else:
+            target[key] = value
+    return target
+
+def merge_deep(objects):
+    result = {}
+    for o in objects:
+        merge_deep_recursive(result, o)
+    return result
+
+def enhance_roles_with_custom_function(roles):
+    if vim.eval("exists('g:vim_ai_roles_config_function')") == '1':
+        roles_config_function = vim.eval("g:vim_ai_roles_config_function")
+        if not vim.eval("exists('*" + roles_config_function + "')"):
+            raise Exception(f"Role config function does not exist: {roles_config_function}")
+        else:
+            roles.update(vim.eval(roles_config_function + "()"))
+
+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):
+        raise Exception(f"Role config file does not exist: {roles_config_path}")
+
+    roles = configparser.ConfigParser()
+    roles.read(roles_config_path)
+    roles = dict(roles)
+
+    enhance_roles_with_custom_function(roles)
+
+    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_chat = roles.get(f"{role}.options-chat", {})
+
+    ui = roles.get(f"{role}.ui", {})
+    ui_complete = roles.get(f"{role}.ui-complete", {})
+    ui_chat = roles.get(f"{role}.ui-chat", {})
+
+    return {
+        'role': dict(roles[role]),
+        'config_default': {
+            'options': dict(options),
+            'ui': dict(ui),
+        },
+        'config_complete': {
+            'options': dict(options_complete),
+            'ui': dict(ui_complete),
+        },
+        'config_chat': {
+            'options': dict(options_chat),
+            'ui': dict(ui_chat),
+        },
+    }
+
+def parse_role_names(prompt):
+    chunks = re.split(r'[ :]+', prompt)
+    roles = []
+    for chunk in chunks:
+        if not chunk.startswith("/"):
+            break
+        roles.append(chunk)
+    return [raw_role[1:] for raw_role in roles]
+
+def parse_prompt_and_role_config(instruction, command_type):
+    instruction = instruction.strip()
+    roles = parse_role_names(instruction)
+    if not roles:
+        # does not require role
+        return ('', {})
+
+    last_role = roles[-1]
+    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', '')
+    return role_prompt, config
+
+def make_config(input_var, output_var):
+    input_options = vim.eval(input_var)
+    config_default = input_options['config_default']
+    config_extension = input_options['config_extension']
+    instruction = input_options['instruction']
+    command_type = input_options['command_type']
+
+    role_prompt, role_config = parse_prompt_and_role_config(instruction, command_type)
+
+    final_config = merge_deep([config_default, config_extension, role_config])
+
+    output = {}
+    output['config'] = final_config
+    output['role_prompt'] = role_prompt
+    vim.command(f'let {output_var}={output}')
+    return output

py/utils.py

@@ -44,21 +44,36 @@ def load_api_key(config_token_file_path):
 
     return (api_key, org_id)
 
-def load_config_and_prompt(command_type):
-    prompt, role_options = parse_prompt_and_role(vim.eval("l:prompt"))
-    config = vim.eval("l:config")
-    config['options'] = {
-        **normalize_options(config['options']),
-        **normalize_options(role_options['options_default']),
-        **normalize_options(role_options['options_' + command_type]),
-    }
-    return prompt, config
+def strip_roles(prompt):
+    chunks = re.split(r'[ :]+', prompt)
+    roles = []
+    for chunk in chunks:
+        if not chunk.startswith("/"):
+            break
+        roles.append(chunk)
+    if not roles:
+        return prompt
+    last_role = roles[-1]
+    return prompt[prompt.index(last_role) + len(last_role):].strip()
 
-def normalize_options(options):
+def make_prompt(raw_prompt, role_prompt):
+    prompt = raw_prompt.strip()
+    prompt = strip_roles(prompt)
+
+    if not role_prompt:
+        return prompt
+
+    delim = '' if prompt.startswith(':') else ':\n'
+    prompt = role_prompt + delim + prompt
+
+    return prompt
+
+def make_config(config):
+    options = config['options']
     # initial prompt can be both a string and a list of strings, normalize it to list
     if 'initial_prompt' in options and isinstance(options['initial_prompt'], str):
         options['initial_prompt'] = options['initial_prompt'].split('\n')
-    return options
+    return config
 
 def make_openai_options(options):
     max_tokens = int(options['max_tokens'])
@@ -302,77 +317,6 @@ def enhance_roles_with_custom_function(roles):
         else:
             roles.update(vim.eval(roles_config_function + "()"))
 
-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):
-        raise Exception(f"Role config file does not exist: {roles_config_path}")
-
-    roles = configparser.ConfigParser()
-    roles.read(roles_config_path)
-
-    enhance_roles_with_custom_function(roles)
-
-    if not role in roles:
-        raise Exception(f"Role `{role}` not found")
-
-    options = roles[f"{role}.options"] if f"{role}.options" in roles else {}
-    options_complete =roles[f"{role}.options-complete"] if f"{role}.options-complete" in roles else {}
-    options_chat = roles[f"{role}.options-chat"] if f"{role}.options-chat" in roles else {}
-
-    return {
-        'role': dict(roles[role]),
-        'options': {
-            'options_default': dict(options),
-            'options_complete': dict(options_complete),
-            'options_chat': dict(options_chat),
-        },
-    }
-
-empty_role_options = {
-    'options_default': {},
-    'options_complete': {},
-    'options_chat': {},
-}
-
-def parse_roles(prompt):
-    chunks = re.split(r'[ :]+', prompt)
-    roles = []
-    for chunk in chunks:
-        if not chunk.startswith("/"):
-            break
-        roles.append(chunk)
-    return [raw_role[1:] for raw_role in roles]
-
-def merge_role_configs(configs):
-    merged_options = empty_role_options
-    merged_role = {}
-    for config in configs:
-        options = config['options']
-        merged_options = {
-            'options_default': { **merged_options['options_default'], **options['options_default'] },
-            'options_complete': { **merged_options['options_complete'], **options['options_complete'] },
-            'options_chat': { **merged_options['options_chat'], **options['options_chat'] },
-        }
-        merged_role ={ **merged_role, **config['role'] }
-    return { 'role': merged_role, 'options': merged_options }
-
-def parse_prompt_and_role(raw_prompt):
-    prompt = raw_prompt.strip()
-    roles = parse_roles(prompt)
-    if not roles:
-        # does not require role
-        return (prompt, empty_role_options)
-
-    last_role = roles[-1]
-    prompt = prompt[prompt.index(last_role) + len(last_role):].strip()
-
-    role_configs = [load_role_config(role) for role in roles]
-    config = merge_role_configs(role_configs)
-    if 'prompt' in config['role'] and config['role']['prompt']:
-        delim = '' if prompt.startswith(':') else ':\n'
-        prompt = config['role']['prompt'] + delim + prompt
-    return (prompt, config['options'])
-
 def make_chat_text_chunks(messages, config_options):
     openai_options = make_openai_options(config_options)
     http_options = make_http_options(config_options)

pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+pythonpath =
+  ./py
+  ./tests/mocks