diff options
| -rw-r--r-- | autoload/vim_ai.vim | 53 | ||||
| -rw-r--r-- | py/chat.py | 118 | ||||
| -rw-r--r-- | py/complete.py | 98 | ||||
| -rw-r--r-- | py/context.py (renamed from py/config.py) | 16 | ||||
| -rw-r--r-- | py/roles.py | 26 | ||||
| -rw-r--r-- | py/utils.py | 46 | ||||
| -rw-r--r-- | tests/context_test.py (renamed from tests/config_test.py) | 34 |
7 files changed, 183 insertions, 208 deletions
diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim index 477df99..03ac978 100644 --- a/autoload/vim_ai.vim +++ b/autoload/vim_ai.vim @@ -1,10 +1,6 @@ call vim_ai_config#load() 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 @@ -16,6 +12,14 @@ let s:last_config = {} let s:scratch_buffer_name = ">>> AI chat" +function! s:ImportPythonModules() + for py_module in ['utils', 'context', 'chat', 'complete', 'roles'] + if !py3eval("'" . py_module . "_py_imported' in globals()") + execute "py3file " . s:plugin_root . "/py/" . py_module . ".py" + endif + endfor +endfunction + function! s:StartsWith(longer, shorter) abort return a:longer[0:len(a:shorter)-1] ==# a:shorter endfunction @@ -129,8 +133,8 @@ endfunction " - config - function scoped vim_ai_complete config " - a:1 - optional instruction prompt function! vim_ai#AIRun(uses_range, config, ...) range abort + call s:ImportPythonModules() 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) @@ -139,12 +143,11 @@ function! vim_ai#AIRun(uses_range, config, ...) range abort \ "config_extension": a:config, \ "user_instruction": l:instruction, \ "user_selection": l:selection, + \ "is_selection": l:is_selection, \ "command_type": 'complete', \} - execute "py3file " . s:config_py - let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") - let l:config = l:config_output['config'] - let l:prompt = l:config_output['prompt'] + let l:context = py3eval("make_ai_context(unwrap('l:config_input'))") + let l:config = l:context['config'] let s:last_command = "complete" let s:last_config = a:config @@ -161,7 +164,7 @@ function! vim_ai#AIRun(uses_range, config, ...) range abort else execute "normal! " . a:lastline . "Go" endif - execute "py3file " . s:complete_py + py3 run_ai_completition(unwrap('l:context')) execute "normal! " . a:lastline . "G" finally call s:set_nopaste(l:config) @@ -173,8 +176,8 @@ endfunction " - config - function scoped vim_ai_edit config " - a:1 - optional instruction prompt function! vim_ai#AIEditRun(uses_range, config, ...) range abort + call s:ImportPythonModules() 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) @@ -183,12 +186,11 @@ function! vim_ai#AIEditRun(uses_range, config, ...) range abort \ "config_extension": a:config, \ "user_instruction": l:instruction, \ "user_selection": l:selection, + \ "is_selection": l:is_selection, \ "command_type": 'complete', \} - execute "py3file " . s:config_py - let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") - let l:config = l:config_output['config'] - let l:prompt = l:config_output['prompt'] + let l:context = py3eval("make_ai_context(unwrap('l:config_input'))") + let l:config = l:context['config'] let s:last_command = "edit" let s:last_config = a:config @@ -201,7 +203,7 @@ function! vim_ai#AIEditRun(uses_range, config, ...) range abort call s:set_paste(l:config) call s:SelectSelectionOrRange(l:is_selection, a:firstline, a:lastline) execute "normal! c" - execute "py3file " . s:complete_py + py3 run_ai_completition(unwrap('l:context')) finally call s:set_nopaste(l:config) endtry @@ -246,8 +248,8 @@ endfunction " - config - function scoped vim_ai_chat config " - a:1 - optional instruction prompt function! vim_ai#AIChatRun(uses_range, config, ...) range abort + call s:ImportPythonModules() 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) @@ -256,16 +258,12 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort \ "config_extension": a:config, \ "user_instruction": l:instruction, \ "user_selection": l:selection, + \ "is_selection": l:is_selection, \ "command_type": 'chat', \} - execute "py3file " . s:config_py - let l:config_output = py3eval("make_config_and_prompt(unwrap('l:config_input'))") - let l:config = l:config_output['config'] - let l:prompt = "" - if a:0 > 0 || a:uses_range - let l:prompt = l:config_output['prompt'] - endif - + let l:context = py3eval("make_ai_context(unwrap('l:config_input'))") + let l:config = l:context['config'] + let l:context['prompt'] = a:0 > 0 || a:uses_range ? l:context['prompt'] : '' try call s:set_paste(l:config) @@ -274,7 +272,7 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort let s:last_command = "chat" let s:last_config = a:config - execute "py3file " . s:chat_py + py3 run_ai_chat(unwrap('l:context')) finally call s:set_nopaste(l:config) endtry @@ -302,7 +300,8 @@ function! vim_ai#AIRedoRun() abort endfunction function! vim_ai#RoleCompletion(A,L,P) abort - execute "py3file " . s:roles_py + call s:ImportPythonModules() + let l:role_list = py3eval("load_ai_role_names()") call map(l:role_list, '"/" . v:val') return filter(l:role_list, 'v:val =~ "^' . a:A . '"') endfunction @@ -1,76 +1,74 @@ import vim -# import utils -plugin_root = vim.eval("s:plugin_root") -vim.command(f"py3file {plugin_root}/py/utils.py") +chat_py_imported = True -prompt = vim.eval("l:prompt") -config = make_config(vim.eval("l:config")) -config_options = config['options'] -config_ui = config['ui'] +def run_ai_chat(context): + prompt = context['prompt'] + config = make_config(context['config']) + config_options = config['options'] + config_ui = config['ui'] -def initialize_chat_window(): - lines = vim.eval('getline(1, "$")') - contains_user_prompt = '>>> user' in lines - if not contains_user_prompt: - # user role not found, put whole file content as an user prompt - vim.command("normal! gg") - populates_options = config_ui['populate_options'] == '1' - if populates_options: - vim.command("normal! O[chat-options]") - vim.command("normal! o") - for key, value in config_options.items(): - if key == 'initial_prompt': - value = "\\n".join(value) - vim.command("normal! i" + key + "=" + value + "\n") - vim.command("normal! " + ("o" if populates_options else "O")) - vim.command("normal! i>>> user\n") - - vim.command("normal! G") - vim_break_undo_sequence() - vim.command("redraw") - - file_content = vim.eval('trim(join(getline(1, "$"), "\n"))') - role_lines = re.findall(r'(^>>> user|^>>> system|^<<< assistant).*', file_content, flags=re.MULTILINE) - if not role_lines[-1].startswith(">>> user"): - # last role is not user, most likely completion was cancelled before - vim.command("normal! o") - vim.command("normal! i\n>>> user\n\n") + def initialize_chat_window(): + lines = vim.eval('getline(1, "$")') + contains_user_prompt = '>>> user' in lines + if not contains_user_prompt: + # user role not found, put whole file content as an user prompt + vim.command("normal! gg") + populates_options = config_ui['populate_options'] == '1' + if populates_options: + vim.command("normal! O[chat-options]") + vim.command("normal! o") + for key, value in config_options.items(): + if key == 'initial_prompt': + value = "\\n".join(value) + vim.command("normal! i" + key + "=" + value + "\n") + vim.command("normal! " + ("o" if populates_options else "O")) + vim.command("normal! i>>> user\n") - if prompt: - vim.command("normal! i" + prompt) + vim.command("normal! G") vim_break_undo_sequence() vim.command("redraw") -initialize_chat_window() + file_content = vim.eval('trim(join(getline(1, "$"), "\n"))') + role_lines = re.findall(r'(^>>> user|^>>> system|^<<< assistant).*', file_content, flags=re.MULTILINE) + if not role_lines[-1].startswith(">>> user"): + # last role is not user, most likely completion was cancelled before + vim.command("normal! o") + vim.command("normal! i\n>>> user\n\n") -chat_options = parse_chat_header_options() -options = {**config_options, **chat_options} + if prompt: + vim.command("normal! i" + prompt) + vim_break_undo_sequence() + vim.command("redraw") -initial_prompt = '\n'.join(options.get('initial_prompt', [])) -initial_messages = parse_chat_messages(initial_prompt) + initialize_chat_window() -chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))') -printDebug("[chat] text:\n" + chat_content) -chat_messages = parse_chat_messages(chat_content) -is_selection = vim.eval("l:is_selection") + chat_options = parse_chat_header_options() + options = {**config_options, **chat_options} -messages = initial_messages + chat_messages + initial_prompt = '\n'.join(options.get('initial_prompt', [])) + initial_messages = parse_chat_messages(initial_prompt) -try: - if messages[-1]["content"].strip(): - vim.command("normal! Go\n<<< assistant\n\n") - vim.command("redraw") + chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))') + print_debug("[chat] text:\n" + chat_content) + chat_messages = parse_chat_messages(chat_content) - print('Answering...') - vim.command("redraw") + messages = initial_messages + chat_messages - text_chunks = make_chat_text_chunks(messages, options) - render_text_chunks(text_chunks, is_selection) + try: + if messages[-1]["content"].strip(): + vim.command("normal! Go\n<<< assistant\n\n") + vim.command("redraw") - vim.command("normal! a\n\n>>> user\n\n") - vim.command("redraw") - clear_echo_message() -except BaseException as error: - handle_completion_error(error) - printDebug("[chat] error: {}", traceback.format_exc()) + print('Answering...') + vim.command("redraw") + + text_chunks = make_chat_text_chunks(messages, options) + render_text_chunks(text_chunks) + + vim.command("normal! a\n\n>>> user\n\n") + vim.command("redraw") + clear_echo_message() + except BaseException as error: + handle_completion_error(error) + print_debug("[chat] error: {}", traceback.format_exc()) diff --git a/py/complete.py b/py/complete.py index a7b9569..8078dea 100644 --- a/py/complete.py +++ b/py/complete.py @@ -1,52 +1,50 @@ import vim -# import utils -plugin_root = vim.eval("s:plugin_root") -vim.command(f"py3file {plugin_root}/py/utils.py") - -prompt = vim.eval("l:prompt") -config = make_config(vim.eval("l:config")) -config_options = config['options'] -config_ui = config['ui'] - -engine = config['engine'] -is_selection = vim.eval("l:is_selection") - -def complete_engine(prompt): - openai_options = make_openai_options(config_options) - http_options = make_http_options(config_options) - printDebug("[engine-complete] text:\n" + prompt) - - request = { - 'prompt': prompt, - **openai_options - } - printDebug("[engine-complete] request: {}", request) - url = config_options['endpoint_url'] - response = openai_request(url, request, http_options) - def map_chunk(resp): - printDebug("[engine-complete] response: {}", resp) - return resp['choices'][0].get('text', '') - text_chunks = map(map_chunk, response) - return text_chunks - -def chat_engine(prompt): - initial_prompt = config_options.get('initial_prompt', []) - initial_prompt = '\n'.join(initial_prompt) - chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip() - messages = parse_chat_messages(chat_content) - printDebug("[engine-chat] text:\n" + chat_content) - return make_chat_text_chunks(messages, config_options) - -engines = {"chat": chat_engine, "complete": complete_engine} - -try: - if prompt: - print('Completing...') - vim.command("redraw") - text_chunks = engines[engine](prompt) - render_text_chunks(text_chunks, is_selection) - clear_echo_message() -except BaseException as error: - handle_completion_error(error) - printDebug("[complete] error: {}", traceback.format_exc()) +complete_py_imported = True + +def run_ai_completition(context): + prompt = context['prompt'] + config = make_config(context['config']) + config_options = config['options'] + config_ui = config['ui'] + + engine = config['engine'] + + def complete_engine(prompt): + openai_options = make_openai_options(config_options) + http_options = make_http_options(config_options) + print_debug("[engine-complete] text:\n" + prompt) + + request = { + 'prompt': prompt, + **openai_options + } + print_debug("[engine-complete] request: {}", request) + url = config_options['endpoint_url'] + response = openai_request(url, request, http_options) + def map_chunk(resp): + print_debug("[engine-complete] response: {}", resp) + return resp['choices'][0].get('text', '') + text_chunks = map(map_chunk, response) + return text_chunks + + def chat_engine(prompt): + initial_prompt = config_options.get('initial_prompt', []) + initial_prompt = '\n'.join(initial_prompt) + chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip() + messages = parse_chat_messages(chat_content) + print_debug("[engine-chat] text:\n" + chat_content) + return make_chat_text_chunks(messages, config_options) + + engines = {"chat": chat_engine, "complete": complete_engine} + + try: + if prompt: + print('Completing...') + vim.command("redraw") + text_chunks = engines[engine](prompt) + render_text_chunks(text_chunks) + clear_echo_message() + except BaseException as error: + handle_completion_error(error) + print_debug("[complete] error: {}", traceback.format_exc()) diff --git a/py/config.py b/py/context.py index 7739b27..254cd3b 100644 --- a/py/config.py +++ b/py/context.py @@ -3,8 +3,10 @@ import re import os import configparser -def unwrap(input_var): - return vim.eval(input_var) +if "PYTEST_VERSION" in os.environ: + from utils import * + +context_py_imported = True def merge_deep_recursive(target, source = {}): source = source.copy() @@ -22,14 +24,6 @@ def merge_deep(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): @@ -113,7 +107,7 @@ def make_prompt(role_prompt, user_prompt, user_selection, selection_boundary): prompt = f"{role_prompt}{delimiter}{prompt}" return prompt -def make_config_and_prompt(params): +def make_ai_context(params): config_default = params['config_default'] config_extension = params['config_extension'] user_instruction = params['user_instruction'] diff --git a/py/roles.py b/py/roles.py index b7d19e1..16aa4e9 100644 --- a/py/roles.py +++ b/py/roles.py @@ -1,23 +1,17 @@ import vim -# import utils -plugin_root = vim.eval("s:plugin_root") -vim.command(f"py3file {plugin_root}/py/utils.py") +roles_py_imported = True -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}") +def load_ai_role_names(): + 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 = configparser.ConfigParser() + roles.read(roles_config_path) -enhance_roles_with_custom_function(roles) + enhance_roles_with_custom_function(roles) -role_names = [name for name in roles.sections() if not '.' in name] + role_names = [name for name in roles.sections() if not '.' in name] -role_list = [f'"{name}"' for name in role_names] -role_list = ", ".join(role_list) - -role_list = f"[{role_list}]" - -vim.command(f'let l:role_list = {role_list}') + return role_names diff --git a/py/utils.py b/py/utils.py index ae32bc2..ce95ed6 100644 --- a/py/utils.py +++ b/py/utils.py @@ -13,12 +13,17 @@ from urllib.error import HTTPError import traceback import configparser -is_debugging = vim.eval("g:vim_ai_debug") == "1" -debug_log_file = vim.eval("g:vim_ai_debug_log_file") +utils_py_imported = True + +def is_ai_debugging(): + return vim.eval("g:vim_ai_debug") == "1" class KnownError(Exception): pass +def unwrap(input_var): + return vim.eval(input_var) + def load_api_key(config_token_file_path): # token precedence: config file path, global file path, env variable global_token_file_path = vim.eval("g:vim_ai_token_file_path") @@ -72,32 +77,19 @@ def make_http_options(options): 'token_file_path': options['token_file_path'], } -# During text manipulation in Vim's visual mode, we utilize "normal! c" command. This command deletes the highlighted text, -# immediately followed by entering insert mode where it generates desirable text. - -# Normally, Vim contemplates the position of the first character in selection to decide whether to place the entered text -# before or after the cursor. For instance, if the given line is "abcd", and "abc" is selected for deletion and "1234" is -# written in its place, the result is as expected "1234d" rather than "d1234". However, if "bc" is chosen for deletion, the -# achieved output is "a1234d", whereas "1234ad" is not. - -# Despite this, post Vim script's execution of "normal! c", it takes an exit immediately returning to the normal mode. This -# might trigger a potential misalignment issue especially when the most extreme left character is the line’s second character. - -# To avoid such pitfalls, the method "need_insert_before_cursor" checks not only the selection status, but also the character -# at the first position of the highlighting. If the selection is off or the first position is not the second character in the line, -# it determines no need for prefixing the cursor. -def need_insert_before_cursor(is_selection): - if is_selection == False: - return False +# when running AIEdit on selection and cursor ends on the first column, it needs to +# be decided whether to append (a) or insert (i) to prevent missalignment. +# Example: helloxxx<Esc>hhhvb:AIE translate<CR> - expected Holaxxx, not xHolaxx +def need_insert_before_cursor(): pos = vim.eval("getpos(\"'<\")[1:2]") if not isinstance(pos, list) or len(pos) != 2: raise ValueError("Unexpected getpos value, it should be a list with two elements") return pos[1] == "1" # determines if visual selection starts on the first window column -def render_text_chunks(chunks, is_selection): +def render_text_chunks(chunks): generating_text = False full_text = '' - insert_before_cursor = need_insert_before_cursor(is_selection) + insert_before_cursor = need_insert_before_cursor() for text in chunks: if not generating_text: text = text.lstrip() # trim newlines from the beginning @@ -200,10 +192,10 @@ def vim_break_undo_sequence(): # breaks undo sequence (https://vi.stackexchange.com/a/29087) vim.command("let &ul=&ul") -def printDebug(text, *args): - if not is_debugging: +def print_debug(text, *args): + if not is_ai_debugging(): return - with open(debug_log_file, "a") as file: + with open(vim.eval("g:vim_ai_debug_log_file"), "a") as file: message = text.format(*args) if len(args) else text file.write(f"[{datetime.datetime.now()}] " + message + "\n") @@ -301,7 +293,7 @@ def make_chat_text_chunks(messages, config_options): 'messages': messages, **openai_options } - printDebug("[engine-chat] request: {}", request) + print_debug("[engine-chat] request: {}", request) url = config_options['endpoint_url'] response = openai_request(url, request, http_options) @@ -315,11 +307,11 @@ def make_chat_text_chunks(messages, config_options): return choices def map_chunk_no_stream(resp): - printDebug("[engine-chat] response: {}", resp) + print_debug("[engine-chat] response: {}", resp) return _choices(resp)[0].get('message', {}).get('content', '') def map_chunk_stream(resp): - printDebug("[engine-chat] response: {}", resp) + print_debug("[engine-chat] response: {}", resp) return _choices(resp)[0].get('delta', {}).get('content', '') map_chunk = map_chunk_stream if openai_options['stream'] else map_chunk_no_stream diff --git a/tests/config_test.py b/tests/context_test.py index a13d89f..16c6e10 100644 --- a/tests/config_test.py +++ b/tests/context_test.py @@ -1,5 +1,5 @@ import vim -from config import make_config_and_prompt, make_prompt +from context import make_ai_context, make_prompt default_config = { "options": { @@ -25,21 +25,21 @@ default_config = { } def test_default_config(): - actual_output = make_config_and_prompt({ + actual_context = make_ai_context({ 'config_default': default_config, 'config_extension': {}, 'user_instruction': 'translate to Slovak', 'user_selection': 'Hello world!', 'command_type': 'chat', }) - expected_output = { + expected_context = { 'config': default_config, 'prompt': 'translate to Slovak:\nHello world!', } - assert expected_output == actual_output + assert expected_context == actual_context def test_param_config(): - actual_config = make_config_and_prompt({ + actual_config = make_ai_context({ 'config_default': default_config, 'config_extension': { 'options': { @@ -54,58 +54,58 @@ def test_param_config(): assert 'gpt-4o' == actual_config['options']['model'] def test_role_config(): - config = make_config_and_prompt({ + context = make_ai_context({ 'config_default': default_config, 'config_extension': {}, 'user_instruction': '/test-role-simple user instruction', 'user_selection': 'selected text', 'command_type': 'chat', }) - actual_config = config['config'] - actual_prompt = config['prompt'] + 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(): - config = make_config_and_prompt({ + context = make_ai_context({ 'config_default': default_config, 'config_extension': {}, 'user_instruction': '/test-role hello', 'user_selection': '', 'command_type': 'chat', }) - actual_config = config['config'] - actual_prompt = config['prompt'] + 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 - config = make_config_and_prompt({ + context = make_ai_context({ 'config_default': default_config, 'config_extension': {}, 'user_instruction': '/test-role hello', 'user_selection': '', 'command_type': 'complete', }) - actual_config = config['config'] - actual_prompt = config['prompt'] + actual_config = context['config'] + actual_prompt = context['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 'hello' == actual_prompt def test_multiple_role_configs(): - config = make_config_and_prompt({ + context = make_ai_context({ 'config_default': default_config, 'config_extension': {}, 'user_instruction': '/test-role /test-role-simple hello', 'user_selection': '', 'command_type': 'chat', }) - actual_config = config['config'] - actual_prompt = config['prompt'] + 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 |