diff options
Diffstat (limited to '')
| -rw-r--r-- | autoload/vim_ai.vim | 79 | ||||
| -rw-r--r-- | py/chat.py | 2 | ||||
| -rw-r--r-- | py/complete.py | 2 | ||||
| -rw-r--r-- | py/config.py | 66 | ||||
| -rw-r--r-- | py/utils.py | 24 | ||||
| -rw-r--r-- | tests/config_test.py | 106 | ||||
| -rw-r--r-- | tests/mocks/vim.py | 12 |
7 files changed, 137 insertions, 154 deletions
diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim index ec3e8e1..477df99 100644 --- a/autoload/vim_ai.vim +++ b/autoload/vim_ai.vim @@ -77,29 +77,6 @@ function! s:OpenChatWindow(open_conf, force_new) abort endif endfunction -function! s:MakeSelectionPrompt(selection, instruction, config) - let l:selection = "" - if a:instruction == "" - let l:selection = a:selection - elseif !empty(a:selection) - let l:boundary = a:config['options']['selection_boundary'] - if l:boundary != "" && match(a:selection, l:boundary) == -1 - " NOTE: surround selection with boundary (e.g. #####) in order to eliminate empty responses - let l:selection = l:boundary . "\n" . a:selection . "\n" . l:boundary - else - let l:selection = a:selection - endif - endif - return l:selection -endfunction - -function! s:MakePrompt(selection, instruction, config) - let l:instruction = trim(a:instruction) - let l:delimiter = l:instruction != "" && a:selection != "" ? ":\n" : "" - let l:selection = s:MakeSelectionPrompt(a:selection, l:instruction, a:config) - return join([l:instruction, l:delimiter, l:selection], "") -endfunction - let s:is_handling_paste_mode = 0 function! s:set_paste(config) @@ -153,21 +130,21 @@ endfunction " - a:1 - optional instruction prompt function! vim_ai#AIRun(uses_range, config, ...) range abort let l:instruction = a:0 > 0 ? a:1 : "" + " 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) + let l:config_input = { \ "config_default": g:vim_ai_edit, \ "config_extension": a:config, - \ "instruction": l:instruction, + \ "user_instruction": l:instruction, + \ "user_selection": l:selection, \ "command_type": 'complete', \} execute "py3file " . s:config_py - execute "py3 make_config('l:config_input', 'l:config_output')" + let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") 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) - let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config) + let l:prompt = l:config_output['prompt'] let s:last_command = "complete" let s:last_config = a:config @@ -197,21 +174,21 @@ endfunction " - a:1 - optional instruction prompt function! vim_ai#AIEditRun(uses_range, config, ...) range abort let l:instruction = a:0 > 0 ? a:1 : "" + " 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) + let l:config_input = { \ "config_default": g:vim_ai_edit, \ "config_extension": a:config, - \ "instruction": l:instruction, + \ "user_instruction": l:instruction, + \ "user_selection": l:selection, \ "command_type": 'complete', \} execute "py3file " . s:config_py - execute "py3 make_config('l:config_input', 'l:config_output')" + let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") 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) - let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config) + let l:prompt = l:config_output['prompt'] let s:last_command = "edit" let s:last_config = a:config @@ -270,30 +247,30 @@ endfunction " - a:1 - optional instruction prompt function! vim_ai#AIChatRun(uses_range, config, ...) range abort let l:instruction = a:0 > 0 ? a:1 : "" + " 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) + let l:config_input = { \ "config_default": g:vim_ai_chat, \ "config_extension": a:config, - \ "instruction": l:instruction, + \ "user_instruction": l:instruction, + \ "user_selection": l:selection, \ "command_type": 'chat', \} execute "py3file " . s:config_py - execute "py3 make_config('l:config_input', 'l:config_output')" + let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") let l:config = l:config_output['config'] - let l:role_prompt = l:config_output['role_prompt'] + let l:prompt = "" + if a:0 > 0 || a:uses_range + let l:prompt = l:config_output['prompt'] + endif + - " 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) try call s:set_paste(l:config) - call s:ReuseOrCreateChatWindow(l:config) - let l:prompt = "" - if a:0 > 0 || a:uses_range - let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config) - endif - let s:last_command = "chat" let s:last_config = a:config @@ -4,7 +4,7 @@ import vim plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt")) +prompt = vim.eval("l:prompt") config = make_config(vim.eval("l:config")) config_options = config['options'] config_ui = config['ui'] diff --git a/py/complete.py b/py/complete.py index 9b63b0b..a7b9569 100644 --- a/py/complete.py +++ b/py/complete.py @@ -4,7 +4,7 @@ import vim plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt")) +prompt = vim.eval("l:prompt") config = make_config(vim.eval("l:config")) config_options = config['options'] config_ui = config['ui'] diff --git a/py/config.py b/py/config.py index 03eb3ca..7739b27 100644 --- a/py/config.py +++ b/py/config.py @@ -3,6 +3,9 @@ import re import os import configparser +def unwrap(input_var): + return vim.eval(input_var) + def merge_deep_recursive(target, source = {}): source = source.copy() for key, value in source.items(): @@ -74,32 +77,55 @@ def parse_role_names(prompt): 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) +def parse_prompt_and_role_config(user_instruction, command_type): + user_instruction = user_instruction.strip() + roles = parse_role_names(user_instruction) if not roles: # does not require role - return ('', {}) + return (user_instruction, '', {}) 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', '') - 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) - + return user_prompt, role_prompt, config + +def make_selection_prompt(user_selection, user_prompt, role_prompt, selection_boundary): + if not user_prompt and not role_prompt: + return user_selection + elif user_selection: + if selection_boundary and selection_boundary not in user_selection: + return f"{selection_boundary}\n{user_selection}\n{selection_boundary}" + else: + return user_selection + return '' + +def make_prompt(role_prompt, user_prompt, user_selection, selection_boundary): + user_prompt = user_prompt.strip() + delimiter = ":\n" if user_prompt and user_selection else "" + user_selection = make_selection_prompt(user_selection, user_prompt, role_prompt, selection_boundary) + prompt = f"{user_prompt}{delimiter}{user_selection}" + if not role_prompt: + return prompt + delimiter = '' if prompt.startswith(':') else ':\n' + prompt = f"{role_prompt}{delimiter}{prompt}" + return prompt + +def make_config_and_prompt(params): + config_default = params['config_default'] + config_extension = params['config_extension'] + user_instruction = params['user_instruction'] + user_selection = params['user_selection'] + command_type = params['command_type'] + + user_prompt, role_prompt, role_config = parse_prompt_and_role_config(user_instruction, command_type) final_config = merge_deep([config_default, config_extension, role_config]) + selection_boundary = final_config['options']['selection_boundary'] + prompt = make_prompt(role_prompt, user_prompt, user_selection, selection_boundary) - output = {} - output['config'] = final_config - output['role_prompt'] = role_prompt - vim.command(f'let {output_var}={output}') - return output + return { + 'config': final_config, + 'prompt': prompt, + } diff --git a/py/utils.py b/py/utils.py index fc888ab..ae32bc2 100644 --- a/py/utils.py +++ b/py/utils.py @@ -44,30 +44,6 @@ def load_api_key(config_token_file_path): return (api_key, org_id) -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 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 diff --git a/tests/config_test.py b/tests/config_test.py index 7b8d5f1..a13d89f 100644 --- a/tests/config_test.py +++ b/tests/config_test.py @@ -1,17 +1,5 @@ import vim -import os -from config import make_config - -dirname = os.path.dirname(__file__) - -def default_eval_mock(cmd): - match cmd: - case 'g:vim_ai_debug_log_file': - return '/tmp/vim_ai_debug.log' - case 'g:vim_ai_roles_config_file': - return dirname + '/resources/roles.ini' - case _: - return None +from config import make_config_and_prompt, make_prompt default_config = { "options": { @@ -36,98 +24,104 @@ default_config = { }, } -def make_input_mock(mocker, input_options): - def eval_mock(cmd): - if cmd == 'l:input': - return input_options - return default_eval_mock(cmd) - mocker.patch('vim.eval', eval_mock) - - -def test_default_config(mocker): - make_input_mock(mocker, { +def test_default_config(): + actual_output = make_config_and_prompt({ 'config_default': default_config, 'config_extension': {}, - 'instruction': 'hello', + 'user_instruction': 'translate to Slovak', + 'user_selection': 'Hello world!', 'command_type': 'chat', }) - command_spy = mocker.spy(vim, "command") - actual_output = make_config('l:input', 'l:output') expected_output = { 'config': default_config, - 'role_prompt': '', + 'prompt': 'translate to Slovak:\nHello world!', } - command_spy.assert_called_once_with(f"let l:output={expected_output}") assert expected_output == actual_output -def test_param_config(mocker): - make_input_mock(mocker, { +def test_param_config(): + actual_config = make_config_and_prompt({ 'config_default': default_config, 'config_extension': { 'options': { 'max_tokens': '1000', }, }, - 'instruction': 'hello', + 'user_instruction': 'hello', + 'user_selection': '', 'command_type': 'chat', - }) - actual_config = make_config('l:input', 'l:output')['config'] + })['config'] assert '1000' == actual_config['options']['max_tokens'] assert 'gpt-4o' == actual_config['options']['model'] -def test_role_config(mocker): - make_input_mock(mocker, { +def test_role_config(): + config = make_config_and_prompt({ 'config_default': default_config, 'config_extension': {}, - 'instruction': '/test-role-simple', + 'user_instruction': '/test-role-simple user instruction', + 'user_selection': 'selected text', 'command_type': 'chat', }) - config = make_config('l:input', 'l:output') actual_config = config['config'] - actual_role_prompt = config['role_prompt'] + actual_prompt = config['prompt'] assert 'o1-preview' == actual_config['options']['model'] - assert 'simple role prompt' == actual_role_prompt + assert 'simple role prompt:\nuser instruction:\nselected text' == actual_prompt -def test_role_config_different_commands(mocker): - make_input_mock(mocker, { +def test_role_config_different_commands(): + config = make_config_and_prompt({ 'config_default': default_config, 'config_extension': {}, - 'instruction': '/test-role hello', + 'user_instruction': '/test-role hello', + 'user_selection': '', 'command_type': 'chat', }) - config = make_config('l:input', 'l:output') actual_config = config['config'] - actual_role_prompt = config['role_prompt'] + actual_prompt = config['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 '' == actual_role_prompt + assert 'hello' == actual_prompt - make_input_mock(mocker, { + config = make_config_and_prompt({ 'config_default': default_config, 'config_extension': {}, - 'instruction': '/test-role hello', + 'user_instruction': '/test-role hello', + 'user_selection': '', 'command_type': 'complete', }) - config = make_config('l:input', 'l:output') actual_config = config['config'] - actual_role_prompt = config['role_prompt'] + actual_prompt = config['prompt'] assert 'model-common' == actual_config['options']['model'] assert 'https://localhost/complete' == actual_config['options']['endpoint_url'] assert '0' == actual_config['ui']['paste_mode'] - assert '' == actual_role_prompt + assert 'hello' == actual_prompt -def test_multiple_role_configs(mocker): - make_input_mock(mocker, { +def test_multiple_role_configs(): + config = make_config_and_prompt({ 'config_default': default_config, 'config_extension': {}, - 'instruction': '/test-role /test-role-simple hello', + 'user_instruction': '/test-role /test-role-simple hello', + 'user_selection': '', 'command_type': 'chat', }) - config = make_config('l:input', 'l:output') actual_config = config['config'] - actual_role_prompt = config['role_prompt'] + actual_prompt = config['prompt'] assert 'o1-preview' == actual_config['options']['model'] assert 'https://localhost/chat' == actual_config['options']['endpoint_url'] - assert 'simple role prompt' == actual_role_prompt + assert 'simple role prompt:\nhello' == actual_prompt + +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', '') + +def test_role_prompt(): + assert 'fix grammar:\nhelo word' == make_prompt( 'fix grammar', 'helo word', '', '') + assert 'fix grammar:\nhelo word' == make_prompt( 'fix grammar', '', 'helo word', '') + assert 'fix grammar:\nand spelling:\nhelo word' == make_prompt( 'fix grammar', 'and spelling', 'helo word', '') + +def test_selection_prompt(): + assert 'fix grammar:\nhelo word' == make_prompt( '', '', 'fix grammar:\nhelo word', '') + +def test_selection_boundary(): + assert 'fix grammar:\n###\nhelo word\n###' == make_prompt( '', 'fix grammar', 'helo word', '###') + assert 'fix grammar:\n###\nhelo word\n###' == make_prompt( 'fix grammar', '', 'helo word', '###') diff --git a/tests/mocks/vim.py b/tests/mocks/vim.py index be14b4b..c0e01e9 100644 --- a/tests/mocks/vim.py +++ b/tests/mocks/vim.py @@ -1,5 +1,15 @@ +import os + +dirname = os.path.dirname(__file__) + def eval(cmd): - pass + match cmd: + case 'g:vim_ai_debug_log_file': + return '/tmp/vim_ai_debug.log' + case 'g:vim_ai_roles_config_file': + return os.path.join(dirname, '..', 'resources/roles.ini') + case _: + return None def command(cmd): pass |