Ollama Vim Plugin

ollama.vim

" Ollama Plugin: Send selected text or the whole buffer to Ollama

" Escape the input
function! EscapeSingleQuotes(input)
    " Escape single quotes for safe use in shell commands
    return substitute(a:input, "'", "'\"'\"'", "g")
endfunction

" Ensure the Command-Result (CR) window exists
function! EnsureCRWindow()
    if !exists('g:ollama_cr_window')
        let g:ollama_cr_window = 'CR'
    endif
endfunction

" Open the Command-Result (CR) window
function! OpenCRWindow()
    call EnsureCRWindow()
    execute "rightbelow new" g:ollama_cr_window
    setlocal buftype=nofile
    setlocal bufhidden=wipe
    setlocal nobuflisted
    setlocal filetype=markdown
endfunction

" Write output to the Command-Result (CR) window
function! WriteToCR(output)
    call OpenCRWindow()
    execute 'normal! ggdG'
    " Split the output into lines and append it cleanly
    let lines = split(a:output, "\n")
    call append(0, lines)
    normal! gg
endfunction

" Display a temporary status message
function! SetStatusMessage(message)
    let &statusline = '%=' . a:message . '%='
    redraw
endfunction

" Restore the status line to its default
function! RestoreStatusLine()
    let &statusline = ''
    redraw
endfunction

function! SendToOllama(mode, question)
    call EnsureCRWindow()
    let selection = ''

    " Check if we are in visual mode and there's an active selection
    if a:mode ==# 'v'
        let start = getpos("'<")
        let end = getpos("'>")
        if start[1] != 0 && end[1] != 0 && start[1] <= end[1]
            let selection = join(getline(start[1], end[1]), "\n")
        endif
    endif

    " If no selection or not in visual mode, use the whole buffer
    if selection == ''
        let selection = join(getline(1, '$'), "\n")
    endif

    " Escape the captured selection
    let escaped_selection = EscapeSingleQuotes(selection)

    " Build the command with cleaning filters
    if a:question !=# ''
        let prompt = input('Enter your question: ')
        let escaped_prompt = EscapeSingleQuotes(prompt)
        let command = printf(
        \   "echo '%s | Question: %s' | ollama run llama3.1 2> /dev/null | tr '\\0' '\\n' | sed -u 's/\r//g'", 
        \   escaped_selection, escaped_prompt)
    else
        let command = printf(
        \   "echo '%s' | ollama run llama3.1 2> /dev/null | tr '\\0' '\\n' | sed -u 's/\r//g'", 
        \   escaped_selection)
    endif

    " Display "Working..." status
    call SetStatusMessage('Working...')

    " Run the command and capture output
    let output = system(command)

    " Restore the status line
    call RestoreStatusLine()

    " Write the result to the CR window
    call WriteToCR(output)
endfunction

" Map <Leader>l for sending highlighted text or whole buffer
xnoremap <silent> <Leader>l :<C-u>call SendToOllama('v', '')<CR>
nnoremap <silent> <Leader>l :call SendToOllama('line', '')<CR>

" Map <Leader>L for sending text with a question or whole buffer
xnoremap <silent> <Leader>L :<C-u>call SendToOllama('v', 'question')<CR>
nnoremap <silent> <Leader>L :call SendToOllama('line', 'question')<CR>