From 41053f92d305468844cb7ad539d0ff752c1e9d6a Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 10 Oct 2024 08:38:57 +0300 Subject: [PATCH] llama.vim : simplify init and cancel + auto-fim --- examples/llama.vim | 210 +++++++++++++++++++++++---------------------- 1 file changed, 109 insertions(+), 101 deletions(-) diff --git a/examples/llama.vim b/examples/llama.vim index b8cfa5906..de889678d 100644 --- a/examples/llama.vim +++ b/examples/llama.vim @@ -2,12 +2,9 @@ " " - Ctrl+F - trigger FIM completion " -" copy paste this in your .vimrc: +" run this once to initialise the plugin: " -"augroup llama_cpp -" autocmd! -" autocmd InsertEnter * inoremap :call llama#fim()a -"augroup END +" :call llama#init() " " color of the suggested text @@ -21,24 +18,76 @@ let s:default_config = { \ 'n_predict': 64, \ 'n_probs': 3, \ 'temperature': 0.1, + \ 'auto_fim': v:true, \ 'stop': ["\n"] \ } let g:llama_config = get(g:, 'llama_config', s:default_config) -function! llama#fim() abort +function! llama#init() + let s:pos_x = 0 + let s:pos_y = 0 + let s:pos_x0 = 0 " pos_x corrected for end-of-line edge case + + let s:line_cur = '' + + let s:pos_dx = 0 + let s:content = [] + let s:can_accept = v:false + + let s:timer_fim = -1 + let s:t_fim_last = reltime() + + augroup llama + autocmd! + autocmd InsertEnter * inoremap :call llama#fim(v:false) + augroup END + + silent! call llama#fim_cancel() +endfunction + +" setup accept/cancel events +function! llama#on_hint(id_timer) + inoremap :call llama#fim_accept() + inoremap :call llama#fim_cancel() + + augroup llama_insert + autocmd! + autocmd CursorMovedI * call llama#fim_cancel() + augroup END +endfunction + +function! llama#fim_auto() + if reltimefloat(reltime(s:t_fim_last)) < 0.50 + if s:timer_fim != -1 + call timer_stop(s:timer_fim) + let s:timer_fim = -1 + endif + endif + + let s:t_fim_last = reltime() + let s:timer_fim = timer_start(500, {-> llama#fim(v:true)}) +endfunction + +function! llama#fim(is_auto) abort let l:t_start = reltime() - let l:pos_x = col('.') - let l:pos_y = line('.') + let s:content = [] + let s:can_accept = v:false + + let s:pos_x = col('.') + let s:pos_y = line('.') let l:max_y = line('$') - let l:lines_prefix = getline(max([1, l:pos_y - g:llama_config.n_prefix]), l:pos_y - 1) - let l:lines_suffix = getline(l:pos_y + 1, min([l:max_y, l:pos_y + g:llama_config.n_suffix])) + let l:lines_prefix = getline(max([1, s:pos_y - g:llama_config.n_prefix]), s:pos_y - 1) + let l:lines_suffix = getline(s:pos_y + 1, min([l:max_y, s:pos_y + g:llama_config.n_suffix])) - let l:line_cur = getline('.') - let l:line_cur_prefix = strpart(l:line_cur, 0, l:pos_x) - let l:line_cur_suffix = strpart(l:line_cur, l:pos_x) + let s:line_cur = getline('.') + + let s:pos_x0 = s:pos_x == len(s:line_cur) ? s:pos_x : s:pos_x - 1 + + let l:line_cur_prefix = strpart(s:line_cur, 0, s:pos_x0) + let l:line_cur_suffix = strpart(s:line_cur, s:pos_x0) let l:prefix = "" \ . join(l:lines_prefix, "\n") @@ -73,7 +122,7 @@ function! llama#fim() abort \ g:llama_config.endpoint, shellescape(l:request) \ ) - let l:can_accept = v:true + let s:can_accept = v:true let l:has_info = v:false let l:n_prompt = 0 @@ -84,21 +133,24 @@ function! llama#fim() abort let l:t_gen_ms = 1.0 let l:s_gen = 0 - let s:content = [] - + " TODO: async this let l:raw = system(l:curl_command) - if l:can_accept && v:shell_error - call add(s:content, "<| curl error: is the server on? |>") - let l:can_accept = v:false + if s:can_accept && v:shell_error + if !a:is_auto + call add(s:content, "<| curl error: is the server on? |>") + endif + let s:can_accept = v:false endif - if l:can_accept && l:raw == "" - call add(s:content, "<| empty response: is the server on? |>") - let l:can_accept = v:false + if s:can_accept && l:raw == "" + if !a:is_auto + call add(s:content, "<| empty response: is the server on? |>") + endif + let s:can_accept = v:false endif " get the generated suggestion - if l:can_accept + if s:can_accept let l:response = json_decode(l:raw) for l:part in split(get(l:response, 'content', ''), "\n", 1) @@ -126,14 +178,20 @@ function! llama#fim() abort endif if len(s:content) == 0 - call add(s:content, "<| nothing to suggest |>") - let l:can_accept = v:false + if !a:is_auto + call add(s:content, "<| nothing to suggest |>") + endif + let s:can_accept = v:false + endif + + if len(s:content) == 0 + return endif let s:pos_dx = len(s:content[-1]) let s:content[-1] .= l:line_cur_suffix - call llama#cancel_vt_fim() + call llama#fim_cancel() " display virtual text with the suggestion let l:bufnr = bufnr('%') @@ -153,74 +211,42 @@ function! llama#fim() abort \ 1000.0 * reltimefloat(reltime(l:t_start)) \ ) - call nvim_buf_set_extmark(l:bufnr, l:id_vt_info, l:pos_y - 1, l:pos_x - 1, { + call nvim_buf_set_extmark(l:bufnr, l:id_vt_info, s:pos_y - 1, s:pos_x - 1, { \ 'virt_text': [[l:info, 'llama_hl_info']], \ 'virt_text_pos': 'eol', \ }) endif - call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, l:pos_y - 1, l:pos_x - 1, { + call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, s:pos_y - 1, s:pos_x - 1, { \ 'virt_text': [[s:content[0], 'llama_hl_hint']], - \ 'virt_text_win_col': l:pos_x == 1 ? 0 : virtcol('.') + \ 'virt_text_win_col': s:pos_x == len(s:line_cur) ? virtcol('.') : virtcol('.') - 1 \ }) - call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, l:pos_y - 1, 0, { + call nvim_buf_set_extmark(l:bufnr, l:id_vt_fim, s:pos_y - 1, 0, { \ 'virt_lines': map(s:content[1:], {idx, val -> [[val, 'llama_hl_hint']]}), \ 'virt_text_win_col': virtcol('.') \ }) - " accept suggestion with Tab and reject it with any other key - let s:mapping_on = v:true - - if l:can_accept - inoremap :call llama#accept_vt_fim() - else - inoremap :call llama#cancel_vt_fim() - endif - - for l:key in range(32, 127) + [8, 27] - if l:key != 0x7C - if l:key == 8 - execute 'inoremap :call llama#cancel_vt_fim()' - elseif l:key == 27 - execute 'inoremap :call llama#cancel_vt_fim()' - elseif l:key == 32 - execute 'inoremap :call llama#cancel_vt_fim()' - elseif l:key == 127 - execute 'inoremap :call llama#cancel_vt_fim()' - else - execute 'inoremap ' . nr2char(l:key) . ' :call llama#cancel_vt_fim()' . nr2char(l:key) - endif - endif - endfor - - inoremap :call llama#cancel_vt_fim() - inoremap :call llama#cancel_vt_fim() - inoremap :call llama#cancel_vt_fim() - inoremap :call llama#cancel_vt_fim() + " need to async this call because the in insert mode causes the cursor to move when at the end of the line + call timer_start(0, 'llama#on_hint') endfunction -function! llama#accept_vt_fim() - let l:pos_x = col('.') - let l:pos_y = line('.') - - let l:line_cur = getline('.') - - let l:pos0 = l:pos_x == len(l:line_cur) ? l:pos_x - 1 : l:pos_x - 2 - +function! llama#fim_accept() " insert the suggestion at the cursor location - call setline(l:pos_y, l:line_cur[:l:pos0] . s:content[0]) - if len(s:content) > 1 - call append(l:pos_y, s:content[1:-1]) + if s:can_accept && len(s:content) > 0 + call setline(s:pos_y, s:line_cur[:(s:pos_x0 - 1)] . s:content[0]) + if len(s:content) > 1 + call append(s:pos_y, s:content[1:-1]) + endif + + " move the cursor to the end of the accepted text + call cursor(s:pos_y + len(s:content) - 1, s:pos_x + s:pos_dx) endif - " move the cursor to the end of the accepted text - call cursor(l:pos_y + len(s:content) - 1, l:pos_x + s:pos_dx) - - call llama#cancel_vt_fim() + call llama#fim_cancel() endfunction -function! llama#cancel_vt_fim() +function! llama#fim_cancel() " clear the virtual text let l:bufnr = bufnr('%') @@ -230,31 +256,13 @@ function! llama#cancel_vt_fim() call nvim_buf_clear_namespace(l:bufnr, l:id_vt_fim, 0, -1) call nvim_buf_clear_namespace(l:bufnr, l:id_vt_info, 0, -1) - " remove the key mappings - if exists('s:mapping_on') && s:mapping_on - iunmap + silent! iunmap + silent! iunmap - for l:key in range(32, 127) + [8, 27] - if l:key != 0x7C - if l:key == 8 - execute 'iunmap ' - elseif l:key == 27 - execute 'iunmap ' - elseif l:key == 32 - execute 'iunmap ' - elseif l:key == 127 - execute 'iunmap ' - else - execute 'iunmap ' . nr2char(l:key) - endif - endif - endfor - - iunmap - iunmap - iunmap - iunmap - - let s:mapping_on = v:false - endif + augroup llama_insert + autocmd! + if g:llama_config.auto_fim + autocmd CursorMovedI * call llama#fim_auto() + endif + augroup END endfunction