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

This commit is contained in:
Georgi Gerganov 2024-10-10 08:38:57 +03:00
parent 7e0b5062af
commit 41053f92d3
No known key found for this signature in database
GPG key ID: BF970631944C16B7

View file

@ -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 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 <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])
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 <buffer> <Tab>
silent! iunmap <buffer> <Tab>
silent! iunmap <buffer> <Esc>
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>
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