llama.vim : simplify init and cancel + auto-fim

Author: ggerganov

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 <buffer> <silent> <C-F> <Esc>:call llama#fim()<CR>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 <buffer> <silent> <C-F> <C-O>:call llama#fim(v:false)<CR>
+    augroup END
+
+    silent! call llama#fim_cancel()
+endfunction
+
+" setup accept/cancel events
+function! llama#on_hint(id_timer)
+    inoremap <buffer> <Tab> <C-O>:call llama#fim_accept()<CR>
+    inoremap <buffer> <Esc> <C-O>:call llama#fim_cancel()<CR><Esc>
+
+    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 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        = 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 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 <buffer> <Tab> <C-O>:call llama#accept_vt_fim()<CR>
-    else
-        inoremap <buffer> <Tab> <C-O>:call llama#cancel_vt_fim()<CR>
-    endif
-
-    for l:key in range(32, 127) + [8, 27]
-        if l:key != 0x7C
-            if l:key == 8
-                execute 'inoremap <buffer> <Bs>    <C-O>:call llama#cancel_vt_fim()<CR><Bs>'
-            elseif l:key == 27
-                execute 'inoremap <buffer> <Esc>   <C-O>:call llama#cancel_vt_fim()<CR><Esc>'
-            elseif l:key == 32
-                execute 'inoremap <buffer> <Space> <C-O>:call llama#cancel_vt_fim()<CR><Space>'
-            elseif l:key == 127
-                execute 'inoremap <buffer> <Del>   <C-O>:call llama#cancel_vt_fim()<CR><Del>'
-            else
-                execute 'inoremap <buffer> ' . nr2char(l:key) . ' <C-O>:call llama#cancel_vt_fim()<CR>' . nr2char(l:key)
-            endif
-        endif
-    endfor
-
-    inoremap <buffer> <Up>    <C-O>:call llama#cancel_vt_fim()<CR><Up>
-    inoremap <buffer> <Down>  <C-O>:call llama#cancel_vt_fim()<CR><Down>
-    inoremap <buffer> <Left>  <C-O>:call llama#cancel_vt_fim()<CR><Left>
-    inoremap <buffer> <Right> <C-O>:call llama#cancel_vt_fim()<CR><Right>
+    " need to async this call because the <C-O> 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])
-    endif
+    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(l:pos_y + len(s:content) - 1, l:pos_x + s:pos_dx)
+        " 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
 
-    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 <buffer> <Tab>
-
-        for l:key in range(32, 127) + [8, 27]
-            if l:key != 0x7C
-                if l:key == 8
-                    execute 'iunmap <buffer> <Bs>'
-                elseif l:key == 27
-                    execute 'iunmap <buffer> <Esc>'
-                elseif l:key == 32
-                    execute 'iunmap <buffer> <Space>'
-                elseif l:key == 127
-                    execute 'iunmap <buffer> <Del>'
-                else
-                    execute 'iunmap <buffer> ' . nr2char(l:key)
-                endif
-            endif
-        endfor
-
-        iunmap <buffer> <Up>
-        iunmap <buffer> <Down>
-        iunmap <buffer> <Left>
-        iunmap <buffer> <Right>
+    silent! iunmap <buffer> <Tab>
+    silent! iunmap <buffer> <Esc>
 
-        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