Feat: Add token counting

.gitignore

@@ -7,3 +7,5 @@ __pycache__
 tags
 # pyenv
 .python-version
+# python venv
+venv

README.md

@@ -11,7 +11,7 @@ A ChatGPT Vim plugin, an OpenAI Neovim plugin, and so much more! Neural integrat
 * Focused on privacy and avoiding leaking data to third parties
 * Easily ask AI to explain code or paragraphs `:NeuralExplain`
 * Compatible with Vim 8.0+ & Neovim 0.8+
-* Supported on Linux, Mac OSX, and Windows
+* Supports Linux, Mac OSX, and Windows
 * Only dependency is Python 3.7+
 
 Experience lightning-fast code generation and completion with asynchronous
@@ -28,6 +28,7 @@ them for a better experience.
 
 - [nui.nvim](https://github.com/MunifTanjim/nui.nvim) - for Neovim UI support
 - [significant.nvim](https://github.com/ElPiloto/significant.nvim) - for Neovim animated signs
+- [nvim-notify](https://github.com/rcarriga/nvim-notify) - for Neovim animated notifications
 - [ALE](https://github.com/dense-analysis/ale) - For correcting problems with
   generated code
 

autoload/neural.vim

@@ -3,6 +3,7 @@
 
 " The location of Neural source scripts
 let s:neural_script_dir = expand('<sfile>:p:h:h') . '/neural_sources'
+let s:neural_python_dir = expand('<sfile>:p:h:h') . '/python3'
 " Keep track of the current job.
 let s:current_job = get(s:, 'current_job', 0)
 " Keep track of the line the last request happened on.
@@ -50,47 +51,6 @@ function! neural#OutputErrorMessage(message) abort
     endif
 endfunction
 
-function! s:AddLineToBuffer(buffer, job_data, line) abort
-    " Add lines either if we can add them to the buffer which is no longer the
-    " current one, or otherwise only if we're still in the same buffer.
-    if bufnr('') isnot a:buffer && !exists('*appendbufline')
-        return
-    endif
-
-    let l:moving_line = a:job_data.moving_line
-    let l:started = a:job_data.content_started
-
-    " Skip introductory empty lines.
-    if !l:started && len(a:line) == 0
-        return
-    endif
-
-    " Check if we need to re-position the cursor to stop it appearing to move
-    " down as lines are added.
-    let l:pos = getpos('.')
-    let l:last_line = len(getbufline(a:buffer, 1, '$'))
-    let l:move_up = 0
-
-    if l:pos[1] == l:last_line
-        let l:move_up = 1
-    endif
-
-    " appendbufline isn't available in old Vim versions.
-    if bufnr('') is a:buffer
-        call append(l:moving_line, a:line)
-    else
-        call appendbufline(a:buffer, l:moving_line, a:line)
-    endif
-
-    " Move the cursor back up again to make content appear below.
-    if l:move_up
-        call setpos('.', l:pos)
-    endif
-
-    let a:job_data.moving_line = l:moving_line + 1
-    let a:job_data.content_started = 1
-endfunction
-
 function! s:AddErrorLine(buffer, job_data, line) abort
     call add(a:job_data.error_lines, a:line)
 endfunction
@@ -312,10 +272,11 @@ function! neural#Run(prompt, options) abort
     \   'moving_line': l:moving_line,
     \   'error_lines': [],
     \   'content_started': 0,
+    \   'content_ended': 0,
     \}
     let l:job_id = neural#job#Start(l:command, {
-    \   'mode': 'nl',
-    \   'out_cb': {job_id, line -> s:AddLineToBuffer(l:buffer, l:job_data, line)},
+    \   'mode': 'raw',
+    \   'out_cb': {job_id, text -> neural#handler#AddTextToBuffer(l:buffer, l:job_data, text)},
     \   'err_cb': {job_id, line -> s:AddErrorLine(l:buffer, l:job_data, line)},
     \   'exit_cb': {job_id, exit_code -> s:HandleOutputEnd(l:buffer, l:job_data, exit_code)},
     \})

autoload/neural/config.vim

@@ -14,7 +14,7 @@ let s:defaults = {
 \   'ui': {
 \       'prompt_enabled': v:true,
 \       'prompt_icon': '🗲',
-\       'animated_sign_enabled': v:true,
+\       'animated_sign_enabled': v:false,
 \       'echo_enabled': v:true,
 \   },
 \   'buffer': {
@@ -36,7 +36,7 @@ let s:defaults = {
 \           'api_key': '',
 \           'frequency_penalty': 0.1,
 \           'max_tokens': 2048,
-\           'model': 'gpt-3.5-turbo',
+\           'model': 'gpt-4',
 \           'presence_penalty': 0.1,
 \           'temperature': 0.2,
 \           'top_p': 1,

autoload/neural/count.vim

@@ -0,0 +1,58 @@
+" Author: Anexon <[email protected]>
+" Description: Count the number of tokens in a given input of text.
+
+let s:current_job = get(s:, 'current_count_job', 0)
+
+function! s:HandleOutputEnd(job_data, exit_code) abort
+    if a:exit_code == 0
+        let l:output = 'Tokens: ' . a:job_data.output_lines[0]
+
+        if has('nvim')
+            execute 'lua require(''neural.notify'').info("' . l:output . '")'
+        else
+            call neural#preview#Show(l:output, {'stay_here': 1})
+        endif
+    " Handle error
+    else
+        if has('nvim')
+            execute 'lua require(''neural.notify'').error("' . join(a:job_data.output_lines, "\n") . '")'
+        else
+            call neural#OutputErrorMessage(join(a:job_data.error_lines, "\n"))
+        endif
+    endif
+
+    " Cleanup
+    call neural#job#Stop(s:current_job)
+    let s:current_job = 0
+endfunction
+
+function! neural#count#SelectedLines() abort
+    " TODO: Reload the Neural config if needed and pass.
+    " call neural#config#Load()
+    " TODO: Should be able to get this elsewhere from a factory-like method.
+    let l:job_data = {
+    \   'output_lines': [],
+    \   'error_lines': [],
+    \}
+    let l:job_id = neural#job#Start(neural#utils#GetPythonCommand('utils.py'), {
+    \   'mode': 'nl',
+    \   'out_cb': {job_id, line -> add(l:job_data.output_lines, line)},
+    \   'err_cb': {job_id, line -> add(l:job_data.error_lines, line)},
+    \   'exit_cb': {job_id, exit_code -> s:HandleOutputEnd(l:job_data, exit_code)},
+    \})
+
+    if l:job_id > 0
+        let l:lines = neural#visual#GetRange().selection
+
+        let l:input = {
+        \   'text': join(l:lines, "\n"),
+        \}
+        call neural#job#SendRaw(l:job_id, json_encode(l:input) . "\n")
+    else
+        call neural#OutputErrorMessage('Failed to count tokens')
+
+        return
+    endif
+
+    let s:current_job = l:job_id
+endfunction

autoload/neural/handler.vim

@@ -0,0 +1,136 @@
+scriptencoding utf8
+
+" Author: Anexon <[email protected]>
+" Description: APIs for working with Asynchronous jobs, with an API normalised
+" between Vim 8 and NeoVim.
+
+
+if has('nvim') && !exists('s:ns_id')
+    let s:ns_id = nvim_create_namespace('neural')
+endif
+
+function! neural#handler#AddTextToBuffer(buffer, job_data, stream_data) abort
+    if (bufnr('') isnot a:buffer && !exists('*appendbufline')) || len(a:stream_data) == 0
+        return
+    endif
+
+    let l:leader = ' 🔸🔶'
+    let l:hl_group = 'ALEInfo'
+    let l:text = a:stream_data
+
+    " echoerr a:stream_data
+    "
+
+    " We need to handle creating new lines in the buffer separately to appending
+    " content to an existing line due to Vim/Neovim API design.
+    " if text is? ''
+    " endif
+
+
+    " Check if we need to re-position the cursor to stop it appearing to move
+    " down as lines are added.
+    let l:pos = getpos('.')
+    let l:last_line = len(getbufline(a:buffer, 1, '$'))
+    let l:move_up = 0
+    let l:new_lines = split(a:stream_data, "\n")
+
+    if l:pos[1] == l:last_line
+        let l:move_up = 1
+    endif
+
+    if empty(l:new_lines)
+        return
+    endif
+
+    " Cleanup leader
+    let l:line_content = getbufline(a:buffer, a:job_data.moving_line)
+    let l:new_lines[0] = get(l:line_content, 0, '') . l:new_lines[0]
+
+    if has('nvim')
+        call nvim_buf_set_lines(a:buffer, a:job_data.moving_line-1, a:job_data.moving_line, 0, l:new_lines)
+    else
+        echom string(l:new_lines)
+        call setbufline(a:buffer, a:job_data.moving_line, l:new_lines)
+    endif
+
+    " Move the cursor back up again to make content appear below.
+    if l:move_up
+        call setpos('.', l:pos)
+    endif
+
+    let a:job_data.moving_line += len(l:new_lines)-1
+
+    if has('nvim')
+        call nvim_buf_set_virtual_text(
+        \   a:buffer,
+        \   s:ns_id,  a:job_data.moving_line - 1,
+        \   [[l:leader, l:hl_group]],
+        \   {}
+        \)
+    endif
+    " elseif text is? '<<[EOF]>>'
+    " elseif match(text, '\%x04') != -1
+    "   " Strip out leader character/s at the end of the stream.
+    "     let l:line_content = getbufline(a:buffer, a:job_data.moving_line)
+    "
+    "     if len(l:line_content) != 0
+    "         let l:new_line_content = l:line_content[0][0:-len(l:leader)-1]
+    "     endif
+    "
+    "     call setbufline(a:buffer, a:job_data.moving_line, l:new_line_content)
+    "     let a:job_data.content_ended = 1
+    " else
+    "     let a:job_data.content_started = 1
+    "     " Prepend any current line content with the incoming stream text.
+    "     let l:line_content = getbufline(a:buffer, a:job_data.moving_line)
+    "
+    "     if len(l:line_content) == 0
+    "         let l:new_line_content = text . l:leader
+    "     else
+    "         let l:new_line_content = l:line_content[0][0:-len(l:leader)-1] . text . l:leader
+    "     endif
+    "
+    "     call setbufline(a:buffer, a:job_data.moving_line, l:new_line_content)
+    " endif
+endfunction
+
+function! neural#handler#AddLineToBuffer(buffer, job_data, line) abort
+    " Add lines either if we can add them to the buffer which is no longer the
+    " current one, or otherwise only if we're still in the same buffer.
+    if bufnr('') isnot a:buffer && !exists('*appendbufline')
+        return
+    endif
+
+    let l:moving_line = a:job_data.moving_line
+    let l:started = a:job_data.content_started
+
+    " Skip introductory empty lines.
+    if !l:started && len(a:line) == 0
+        return
+    endif
+
+    " Check if we need to re-position the cursor to stop it appearing to move
+    " down as lines are added.
+    let l:pos = getpos('.')
+    let l:last_line = len(getbufline(a:buffer, 1, '$'))
+    let l:move_up = 0
+
+    if l:pos[1] == l:last_line
+        let l:move_up = 1
+    endif
+
+    " appendbufline isn't available in old Vim versions.
+    if bufnr('') is a:buffer
+        call append(l:moving_line, a:line)
+    else
+        call appendbufline(a:buffer, l:moving_line, a:line)
+    endif
+
+    " Move the cursor back up again to make content appear below.
+    if l:move_up
+        call setpos('.', l:pos)
+    endif
+
+    let a:job_data.moving_line = l:moving_line + 1
+    let a:job_data.content_started = 1
+endfunction

autoload/neural/job.vim

@@ -28,9 +28,23 @@ endfunction
 
 function! s:JoinNeovimOutput(job, last_line, data, mode, callback) abort
     if a:mode is# 'raw'
+        " Neovim stream event handlers receive data as it becomes available
+        " from the OS, thus the first and last items in the data list may be
+        " partial lines.
+        " Each stream item is passed to the callback individually which can be
+        " a chunk of text or a newline character.
+        " echoerr a:data
         call a:callback(a:job, join(a:data, "\n"))
 
-        return ''
+        " if len(a:data) > 1
+        "     for text in a:data
+        "         call a:callback(a:job, [text])
+        "     endfor
+        " else
+        "     call a:callback(a:job, a:data)
+        " endif
+
+        return
     endif
 
     let l:lines = a:data[:-2]