From 81041eb341e4dd8ff1e1d2a141b90ce451a06808 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Fri, 29 Dec 2023 13:58:38 +0000
Subject: [PATCH 01/19] Tech: add tiktoken dependency

---
 .gitmodules           | 3 +++
 python3/deps/tiktoken | 1 +
 2 files changed, 4 insertions(+)
 create mode 100644 .gitmodules
 create mode 160000 python3/deps/tiktoken

diff --git a/.gitmodules b/.gitmodules
new file mode 100644
index 0000000..508abaf
--- /dev/null
+++ b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "python3/deps/tiktoken"]
+	path = python3/deps/tiktoken
+	url = https://github.com/openai/tiktoken.git
diff --git a/python3/deps/tiktoken b/python3/deps/tiktoken
new file mode 160000
index 0000000..9e79899
--- /dev/null
+++ b/python3/deps/tiktoken
@@ -0,0 +1 @@
+Subproject commit 9e79899bc248d5313c7dd73562b5e211d728723d

From e228d05d8713eb3858432b0bc16370a632a2fef5 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Mon, 1 Jan 2024 22:06:15 +0000
Subject: [PATCH 02/19] tech: add wip token count job

---
 autoload/neural/count.vim | 124 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 124 insertions(+)
 create mode 100644 autoload/neural/count.vim

diff --git a/autoload/neural/count.vim b/autoload/neural/count.vim
new file mode 100644
index 0000000..4aa1c86
--- /dev/null
+++ b/autoload/neural/count.vim
@@ -0,0 +1,124 @@
+" Author: Anexon <anexon@protonmail.com>
+" Description: Count the number of tokens in some input of text.
+
+let s:current_job = get(s:, 'current_count_job', 0)
+
+function! s:AddOutputLine(buffer, job_data, line) abort
+    call add(a:job_data.output_lines, a:line)
+endfunction
+
+function! s:AddErrorLine(buffer, job_data, line) abort
+    call add(a:job_data.error_lines, a:line)
+endfunction
+
+function! s:HandleOutputEnd(buffer, job_data, exit_code) abort
+    " Output an error message from the program if something goes wrong.
+    if a:exit_code != 0
+        " Complain when something goes wrong.
+        call neural#OutputErrorMessage(join(a:job_data.error_lines, "\n"))
+    else
+        if has('nvim')
+            execute 'lua require(''neural'').notify("' . a:job_data.output_lines[0] . '", "info")'
+        else
+            call neural#preview#Show(
+            \   a:job_data.output_lines[0],
+            \   {'stay_here': 1},
+            \)
+        endif
+    endif
+
+    let s:current_job = 0
+endfunction
+
+
+function! neural#count#Cleanup() abort
+    if s:current_job
+        call neural#job#Stop(s:current_job)
+        let s:current_job = 0
+    endif
+endfunction
+
+
+" TODO: Refactor
+" Get the path to the executable for a script language.
+function! s:GetScriptExecutable(source) abort
+    if a:source.script_language is# 'python'
+        let l:executable = ''
+
+        if has('win32')
+            " Try to automatically find Python on Windows, even if not in PATH.
+            let l:executable = expand('~/AppData/Local/Programs/Python/Python3*/python.exe')
+        endif
+
+        if empty(l:executable)
+            let l:executable = 'python3'
+        endif
+
+        return l:executable
+    endif
+
+    throw 'Unknown script language: ' . a:source.script_language
+endfunction
+
+" TODO: Refactor
+function! neural#count#GetCommand(buffer) abort
+    let l:source = {
+    \   'name': 'openai',
+    \   'script_language': 'python',
+    \   'script': neural#GetPythonDir() . '/utils.py',
+    \}
+
+    let l:script_exe = s:GetScriptExecutable(l:source)
+    let l:command = neural#Escape(l:script_exe)
+    \   . ' ' . neural#Escape(l:source.script)
+    let l:command = neural#job#PrepareCommand(a:buffer, l:command)
+
+    return [l:source, l:command]
+endfunction
+
+
+function! neural#count#SelectedLines() abort
+    " Reload the Neural config if needed.
+    " call neural#config#Load()
+    " Stop Neural doing anything else if explaining code.
+    " call neural#Cleanup()
+    let l:range = neural#visual#GetRange()
+    let l:buffer = bufnr('')
+
+    let [l:source, l:command] = neural#count#GetCommand(l:buffer)
+
+    let l:job_data = {
+    \   'output_lines': [],
+    \   'error_lines': [],
+    \}
+    let l:job_id = neural#job#Start(l:command, {
+    \   'mode': 'nl',
+    \   'out_cb': {job_id, line -> s:AddOutputLine(l:buffer, l:job_data, line)},
+    \   '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)},
+    \})
+
+    if l:job_id > 0
+        " let l:lines = neural#redact#PasswordsAndSecrets(l:range.selection)
+        let l:lines = l:range.selection
+
+        let l:config = get(g:neural.source, l:source.name, {})
+
+        " If the config is not a Dictionary, throw it away.
+        if type(l:config) isnot v:t_dict
+            let l:config = {}
+        endif
+
+        let l:input = {
+        \   'model': l:config,
+        \   'text': join(l:lines, "\n"),
+        \}
+        call neural#job#SendRaw(l:job_id, json_encode(l:input) . "\n")
+    else
+        call neural#OutputErrorMessage('Failed to run ' . l:source.name)
+
+        return
+    endif
+
+    let s:current_job = l:job_id
+endfunction

From ffd4a6d6cb1f2b537bc836d37f3c55d00a91cd16 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Mon, 1 Jan 2024 22:07:19 +0000
Subject: [PATCH 03/19] feat: add wip :NeuralCount command

---
 autoload/neural.vim | 6 ++++++
 plugin/neural.vim   | 2 ++
 2 files changed, 8 insertions(+)

diff --git a/autoload/neural.vim b/autoload/neural.vim
index b7381da..6aef3d8 100644
--- a/autoload/neural.vim
+++ b/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.
@@ -23,6 +24,11 @@ function! neural#GetScriptDir() abort
     return s:neural_script_dir
 endfunction
 
+" Get the Neural scripts directory in a way that makes it hard to modify.
+function! neural#GetPythonDir() abort
+    return s:neural_python_dir
+endfunction
+
 " Output an error message. The message should be a string.
 " The output error lines will be split in a platform-independent way.
 function! neural#OutputErrorMessage(message) abort
diff --git a/plugin/neural.vim b/plugin/neural.vim
index d682d68..413fe5e 100644
--- a/plugin/neural.vim
+++ b/plugin/neural.vim
@@ -37,6 +37,8 @@ command! -nargs=0 NeuralStop :call neural#Stop()
 command! -nargs=? NeuralBuffer :call neural#buffer#CreateBuffer(<q-args>)
 " Have Neural explain the visually selected lines.
 command! -range NeuralExplain :call neural#explain#SelectedLines()
+" Get the token count for the visually selected lines.
+command! -range NeuralCountTokens :call neural#count#SelectedLines()
 
 " <Plug> mappings for commands
 nnoremap <silent> <Plug>(neural_prompt) :call neural#OpenPrompt()<Return>

From a6defcba58bb9ee1a72f16da623c5532cfd7594e Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Mon, 1 Jan 2024 22:11:04 +0000
Subject: [PATCH 04/19] tech: add wip count tokens from tiktoken

---
 python3/utils.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 python3/utils.py

diff --git a/python3/utils.py b/python3/utils.py
new file mode 100644
index 0000000..3c70ab1
--- /dev/null
+++ b/python3/utils.py
@@ -0,0 +1,65 @@
+import os
+import sys
+import json
+
+sys.path.append('./deps/tiktoken')
+#
+# # import tiktoken
+# from .deps import tiktoken
+
+# script_dir = os.path.dirname(os.path.realpath(__file__))
+# parent_dir = os.path.dirname(script_dir)
+# deps_path = os.path.join(parent_dir, 'deps')
+# sys.path.append(deps_path)
+
+
+# Calculate the absolute path to the 'deps' directory
+# /home/user/.config/nvim/projects/neural/python3
+script_dir = os.path.dirname(os.path.realpath(__file__))
+
+# /home/user/.config/nvim/projects/neural/python2/deps
+tiktoken_dir = os.path.abspath(os.path.join(script_dir, 'deps/tiktoken'))
+regex_dir = os.path.abspath(os.path.join(script_dir, 'deps/mrab-regex'))
+regex_dir2 = os.path.abspath(os.path.join(script_dir, 'deps/mrab-regex/regex_3'))
+
+# Add 'deps' directory to sys.path if it's not already there
+if tiktoken_dir not in sys.path:
+    sys.path.insert(0, tiktoken_dir)
+
+if regex_dir not in sys.path:
+    sys.path.insert(0, regex_dir)
+if regex_dir2 not in sys.path:
+    sys.path.insert(0, regex_dir2)
+
+# Now you can import tiktoken as if it was a top-level module
+# import deps.tiktoken.tiktoken as tiktoken
+# import regex
+# exit(regex)
+import tiktoken
+
+# import deps.mrabregex.regex_3 as regex
+# import tiktoken
+
+# Rest of your code...
+
+def count_tokens(text: str, model: str ="gpt-3.5-turbo") -> int:
+    """
+    Return the number of tokens from an input text using the appropriate
+    tokeniser for the given model.
+    """
+    encoder = tiktoken.encoding_for_model(model)
+
+    return len(encoder.encode(text))
+
+if __name__ == "__main__":
+    # Read input from command line
+    input_data = json.loads(sys.stdin.readline())
+    # input_text = sys.stdin.readline()
+
+    model = input_data["model"]
+
+    # Count tokens
+    count = count_tokens(input_data["text"])
+    print(count)
+
+    # sys.exit(count)

From 34e90c6836576cd92995dbfa7646820dbcae8a74 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Mon, 1 Jan 2024 22:11:59 +0000
Subject: [PATCH 05/19] feat: add notify integration for token count

---
 lua/neural.lua | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/lua/neural.lua b/lua/neural.lua
index 0d3fc3d..4d0ced0 100644
--- a/lua/neural.lua
+++ b/lua/neural.lua
@@ -1,8 +1,11 @@
 -- External dependencies
-local UI = {}
-local AnimatedSign = {}
 local has_nui, _ = pcall(require, 'nui.input')
 local has_significant, _ = pcall(require, 'significant')
+local has_notify, _ = pcall(require, 'notify')
+
+local UI = {}
+local AnimatedSign = {}
+local Notify = {}
 
 if has_nui then
     UI = require('neural.ui')
@@ -12,6 +15,10 @@ if has_significant then
     AnimatedSign = require('significant')
 end
 
+if has_notify then
+  Notify = require('notify')
+end
+
 local Neural = {}
 
 function Neural.setup(settings)
@@ -54,4 +61,16 @@ function Neural.stop_animated_sign(line)
     end
 end
 
+function Neural.notify(message, level)
+  if has_notify then
+    local opts = {
+      title = 'Neural - Token Count',
+      icon = vim.g.neural.ui.prompt_icon
+    }
+    Notify(message, level, opts)
+  else
+    vim.fn['neural#preview#Show'](message, {stay_here = 1})
+  end
+end
+
 return Neural

From e8ea921422636f37829e8caef204661f039ac3c0 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Tue, 2 Jan 2024 16:06:48 +0000
Subject: [PATCH 06/19] tech: remove incorrectly bundled tiktoken

---
 .gitmodules           |  3 ---
 python3/deps/tiktoken |  1 -
 python3/utils.py      | 41 +----------------------------------------
 3 files changed, 1 insertion(+), 44 deletions(-)
 delete mode 100644 .gitmodules
 delete mode 160000 python3/deps/tiktoken

diff --git a/.gitmodules b/.gitmodules
deleted file mode 100644
index 508abaf..0000000
--- a/.gitmodules
+++ /dev/null
@@ -1,3 +0,0 @@
-[submodule "python3/deps/tiktoken"]
-	path = python3/deps/tiktoken
-	url = https://github.com/openai/tiktoken.git
diff --git a/python3/deps/tiktoken b/python3/deps/tiktoken
deleted file mode 160000
index 9e79899..0000000
--- a/python3/deps/tiktoken
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 9e79899bc248d5313c7dd73562b5e211d728723d
diff --git a/python3/utils.py b/python3/utils.py
index 3c70ab1..e9925a7 100644
--- a/python3/utils.py
+++ b/python3/utils.py
@@ -1,47 +1,8 @@
-import os
 import sys
 import json
 
-sys.path.append('./deps/tiktoken')
-#
-# # import tiktoken
-# from .deps import tiktoken
-
-# script_dir = os.path.dirname(os.path.realpath(__file__))
-# parent_dir = os.path.dirname(script_dir)
-# deps_path = os.path.join(parent_dir, 'deps')
-# sys.path.append(deps_path)
-
-
-# Calculate the absolute path to the 'deps' directory
-# /home/user/.config/nvim/projects/neural/python3
-script_dir = os.path.dirname(os.path.realpath(__file__))
-
-# /home/user/.config/nvim/projects/neural/python2/deps
-tiktoken_dir = os.path.abspath(os.path.join(script_dir, 'deps/tiktoken'))
-regex_dir = os.path.abspath(os.path.join(script_dir, 'deps/mrab-regex'))
-regex_dir2 = os.path.abspath(os.path.join(script_dir, 'deps/mrab-regex/regex_3'))
-
-# Add 'deps' directory to sys.path if it's not already there
-if tiktoken_dir not in sys.path:
-    sys.path.insert(0, tiktoken_dir)
-
-if regex_dir not in sys.path:
-    sys.path.insert(0, regex_dir)
-if regex_dir2 not in sys.path:
-    sys.path.insert(0, regex_dir2)
-
-# Now you can import tiktoken as if it was a top-level module
-# import deps.tiktoken.tiktoken as tiktoken
-# import regex
-# exit(regex)
 import tiktoken
 
-# import deps.mrabregex.regex_3 as regex
-# import tiktoken
-
-# Rest of your code...
-
 def count_tokens(text: str, model: str ="gpt-3.5-turbo") -> int:
     """
     Return the number of tokens from an input text using the appropriate
@@ -54,8 +15,8 @@ def count_tokens(text: str, model: str ="gpt-3.5-turbo") -> int:
 if __name__ == "__main__":
     # Read input from command line
     input_data = json.loads(sys.stdin.readline())
-    # input_text = sys.stdin.readline()
 
+    # TODO: Read config
     model = input_data["model"]
 
     # Count tokens

From 346c45c1ce86be4496acb175687822c332a9de44 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 16:44:07 +0000
Subject: [PATCH 07/19] tech: temp remove model args for token count

---
 python3/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python3/utils.py b/python3/utils.py
index e9925a7..1a5e597 100644
--- a/python3/utils.py
+++ b/python3/utils.py
@@ -17,7 +17,7 @@ def count_tokens(text: str, model: str ="gpt-3.5-turbo") -> int:
     input_data = json.loads(sys.stdin.readline())
 
     # TODO: Read config
-    model = input_data["model"]
+    # model = input_data["model"]
 
     # Count tokens
     count = count_tokens(input_data["text"])

From 903921f40835727ba72d72bbfd3ebbe37f80ce42 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 16:44:44 +0000
Subject: [PATCH 08/19] tidy: ignore python virtual envs

---
 .gitignore | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.gitignore b/.gitignore
index 8df500d..13c8b2d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,3 +7,5 @@ __pycache__
 tags
 # pyenv
 .python-version
+# python venv
+venv

From 4b896c9edfda4288aea5bcea3640e6b9027d60ff Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 16:45:13 +0000
Subject: [PATCH 09/19] tech: add common utils for running python scripts

---
 autoload/neural/utils.vim | 84 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 autoload/neural/utils.vim

diff --git a/autoload/neural/utils.vim b/autoload/neural/utils.vim
new file mode 100644
index 0000000..138afc5
--- /dev/null
+++ b/autoload/neural/utils.vim
@@ -0,0 +1,84 @@
+" Author: Anexon <anexon@protonmail.com>
+" Co-author: w0rp <devw0rp@gmail.com>
+" Description: Utils and helpers with API normalised between Neovim/Vim 8 and
+" platform independent.
+
+let s:python_script_dir = expand('<sfile>:p:h:h:h') . '/python3'
+
+function! s:IsWindows() abort
+    return has('win32') || has('win64')
+endfunction
+
+" Return string of full neural python script path.
+function! s:GetPythonScript(script) abort
+    return s:python_script_dir . '/' . a:script
+endfunction
+
+" Return path of python executable.
+function! s:GetPython() abort
+    " Use the virtual environment if it exists.
+    if neural#utils#IsVenvAvailable()
+        return s:python_script_dir . '/venv/bin/python3'
+    else
+        let l:python = ''
+
+        " Try to automatically find Python on Windows, even if not in PATH.
+        if s:IsWindows()
+            let l:python = expand('~/AppData/Local/Programs/Python/Python3*/python.exe')
+        endif
+
+        " Fallback to the system Python path
+        if empty(l:python)
+            let l:python = 'python3'
+        endif
+
+        return l:python
+    endif
+endfunction
+
+" Check the virtual environment exist.
+function! neural#utils#IsVenvAvailable() abort
+    let l:venv_dir = s:python_script_dir . '/venv'
+    let l:venv_python = l:venv_dir . (s:IsWindows() ? '\Scripts\python.exe' : '/bin/python')
+
+    return isdirectory(l:venv_dir) && filereadable(l:venv_python) && executable(l:venv_python)
+endfunction
+
+" Returns python command call for a given neural python script.
+function! neural#utils#GetPythonCommand(script) abort
+    let l:script = neural#utils#StringEscape(s:GetPythonScript(a:script))
+    let l:python = neural#utils#StringEscape(s:GetPython())
+
+    return neural#utils#GetCommand(l:python . ' ' . l:script)
+endfunction
+
+" Return a command that should be executed in a subshell.
+"
+" This fixes issues related to PATH variables, %PATHEXT% in Windows, etc.
+" Neovim handles this automatically if the command is a String, but we do
+" this explicitly for consistency.
+function! neural#utils#GetCommand(command) abort
+    if s:IsWindows()
+        return 'cmd /s/c "' . a:command . '"'
+    endif
+
+    return ['/bin/sh', '-c', a:command]
+endfunction
+
+" Return platform independent escaped String.
+function! neural#utils#StringEscape(str) abort
+    if fnamemodify(&shell, ':t') is? 'cmd.exe'
+        " If the string contains spaces, it will be surrounded by quotes.
+        " Otherwise, special characters will be escaped with carets (^).
+        return substitute(
+        \   a:str =~# ' '
+        \       ?  '"' .  substitute(a:str, '"', '""', 'g') . '"'
+        \       : substitute(a:str, '\v([&|<>^])', '^\1', 'g'),
+        \   '%',
+        \   '%%',
+        \   'g',
+        \)
+    endif
+
+    return shellescape (a:str)
+endfunction

From 4f0b8e1b618041ffb083dfa7e152795604bcf297 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 16:46:44 +0000
Subject: [PATCH 10/19] feat: setup venv through nvim lazy package manager

---
 build.lua                | 45 ++++++++++++++++++++++++++++++++++++++++
 python3/__init__.py      |  0
 python3/requirements.txt |  1 +
 3 files changed, 46 insertions(+)
 create mode 100644 build.lua
 create mode 100644 python3/__init__.py
 create mode 100644 python3/requirements.txt

diff --git a/build.lua b/build.lua
new file mode 100644
index 0000000..056d4cc
--- /dev/null
+++ b/build.lua
@@ -0,0 +1,45 @@
+local function is_windows()
+    return vim.fn.has("win32") == 1 or vim.fn.has("win64") == 1
+end
+
+local function get_path_separator()
+    return is_windows() and "\\" or "/"
+end
+
+local function get_script_path()
+    local path_sep = get_path_separator()
+    local script_path = debug.getinfo(1).source:match("@?(.*)" .. path_sep)
+
+    return script_path or ""
+end
+
+local function run_command(command)
+    local status, _, code = os.execute(command)
+
+    require('notify').notify('Command: ' .. command .. ' || status: ' .. status, 'info')
+    require('notify').notify('status: ' .. status, 'info')
+
+    if status == false then
+        error("Command failed: " .. command .. ". Exit code: " .. tostring(code))
+    end
+end
+
+local function setup()
+    local python = is_windows() and "python" or "python3"
+    local sep = get_path_separator()
+    local script_path = get_script_path()
+    local venv_path = script_path .. sep .. "python3" .. sep .. "venv"
+
+    -- Create virtual environment if it does not exist.
+    if vim.fn.isdirectory(venv_path) == 0 then
+        run_command(python .. " -m venv " .. venv_path)
+    end
+
+    -- Install requirements via pip
+    local pip_cmd = venv_path .. (is_windows() and sep .. "Scripts" .. sep .. "pip" or sep .. "bin" .. sep .. "pip")
+    local requirements_path = script_path .. sep .. "python3" .. sep .. "requirements.txt"
+
+    run_command(pip_cmd .. " install -r " .. requirements_path)
+end
+
+setup()
diff --git a/python3/__init__.py b/python3/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/python3/requirements.txt b/python3/requirements.txt
new file mode 100644
index 0000000..77778b9
--- /dev/null
+++ b/python3/requirements.txt
@@ -0,0 +1 @@
+tiktoken

From 87a87709d07e07408e2903ce5e557d4bbf527bea Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 16:48:00 +0000
Subject: [PATCH 11/19] tidy: simplify token count logic to use utils

---
 autoload/neural.vim       |  5 --
 autoload/neural/count.vim | 99 ++++++---------------------------------
 2 files changed, 14 insertions(+), 90 deletions(-)

diff --git a/autoload/neural.vim b/autoload/neural.vim
index 6aef3d8..55b53f4 100644
--- a/autoload/neural.vim
+++ b/autoload/neural.vim
@@ -24,11 +24,6 @@ function! neural#GetScriptDir() abort
     return s:neural_script_dir
 endfunction
 
-" Get the Neural scripts directory in a way that makes it hard to modify.
-function! neural#GetPythonDir() abort
-    return s:neural_python_dir
-endfunction
-
 " Output an error message. The message should be a string.
 " The output error lines will be split in a platform-independent way.
 function! neural#OutputErrorMessage(message) abort
diff --git a/autoload/neural/count.vim b/autoload/neural/count.vim
index 4aa1c86..5256f56 100644
--- a/autoload/neural/count.vim
+++ b/autoload/neural/count.vim
@@ -1,22 +1,10 @@
 " Author: Anexon <anexon@protonmail.com>
-" Description: Count the number of tokens in some input of text.
+" Description: Count the number of tokens in a given input of text.
 
 let s:current_job = get(s:, 'current_count_job', 0)
 
-function! s:AddOutputLine(buffer, job_data, line) abort
-    call add(a:job_data.output_lines, a:line)
-endfunction
-
-function! s:AddErrorLine(buffer, job_data, line) abort
-    call add(a:job_data.error_lines, a:line)
-endfunction
-
-function! s:HandleOutputEnd(buffer, job_data, exit_code) abort
-    " Output an error message from the program if something goes wrong.
-    if a:exit_code != 0
-        " Complain when something goes wrong.
-        call neural#OutputErrorMessage(join(a:job_data.error_lines, "\n"))
-    else
+function! s:HandleOutputEnd(job_data, exit_code) abort
+    if a:exit_code == 0
         if has('nvim')
             execute 'lua require(''neural'').notify("' . a:job_data.output_lines[0] . '", "info")'
         else
@@ -25,97 +13,38 @@ function! s:HandleOutputEnd(buffer, job_data, exit_code) abort
             \   {'stay_here': 1},
             \)
         endif
+    else
+        call neural#OutputErrorMessage(join(a:job_data.error_lines, "\n"))
     endif
 
+    call neural#job#Stop(s:current_job)
     let s:current_job = 0
 endfunction
 
-
-function! neural#count#Cleanup() abort
-    if s:current_job
-        call neural#job#Stop(s:current_job)
-        let s:current_job = 0
-    endif
-endfunction
-
-
-" TODO: Refactor
-" Get the path to the executable for a script language.
-function! s:GetScriptExecutable(source) abort
-    if a:source.script_language is# 'python'
-        let l:executable = ''
-
-        if has('win32')
-            " Try to automatically find Python on Windows, even if not in PATH.
-            let l:executable = expand('~/AppData/Local/Programs/Python/Python3*/python.exe')
-        endif
-
-        if empty(l:executable)
-            let l:executable = 'python3'
-        endif
-
-        return l:executable
-    endif
-
-    throw 'Unknown script language: ' . a:source.script_language
-endfunction
-
-" TODO: Refactor
-function! neural#count#GetCommand(buffer) abort
-    let l:source = {
-    \   'name': 'openai',
-    \   'script_language': 'python',
-    \   'script': neural#GetPythonDir() . '/utils.py',
-    \}
-
-    let l:script_exe = s:GetScriptExecutable(l:source)
-    let l:command = neural#Escape(l:script_exe)
-    \   . ' ' . neural#Escape(l:source.script)
-    let l:command = neural#job#PrepareCommand(a:buffer, l:command)
-
-    return [l:source, l:command]
-endfunction
-
-
 function! neural#count#SelectedLines() abort
-    " Reload the Neural config if needed.
+    " TODO: Reload the Neural config if needed and pass.
     " call neural#config#Load()
-    " Stop Neural doing anything else if explaining code.
-    " call neural#Cleanup()
-    let l:range = neural#visual#GetRange()
-    let l:buffer = bufnr('')
-
-    let [l:source, l:command] = neural#count#GetCommand(l:buffer)
-
+    " 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(l:command, {
+    let l:job_id = neural#job#Start(neural#utils#GetPythonCommand('utils.py'), {
     \   'mode': 'nl',
-    \   'out_cb': {job_id, line -> s:AddOutputLine(l:buffer, l:job_data, line)},
-    \   '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)},
+    \   '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#redact#PasswordsAndSecrets(l:range.selection)
-        let l:lines = l:range.selection
-
-        let l:config = get(g:neural.source, l:source.name, {})
-
-        " If the config is not a Dictionary, throw it away.
-        if type(l:config) isnot v:t_dict
-            let l:config = {}
-        endif
+        let l:lines = neural#visual#GetRange().selection
 
         let l:input = {
-        \   'model': l:config,
         \   'text': join(l:lines, "\n"),
         \}
         call neural#job#SendRaw(l:job_id, json_encode(l:input) . "\n")
     else
-        call neural#OutputErrorMessage('Failed to run ' . l:source.name)
+        call neural#OutputErrorMessage('Failed to cound tokens')
 
         return
     endif

From 48299ed39f6d3a63574746b8d80b7c2e74caa701 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 17:06:01 +0000
Subject: [PATCH 12/19] tidy: util script authorship

---
 autoload/neural/utils.vim | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/autoload/neural/utils.vim b/autoload/neural/utils.vim
index 138afc5..4ed8504 100644
--- a/autoload/neural/utils.vim
+++ b/autoload/neural/utils.vim
@@ -1,5 +1,4 @@
-" Author: Anexon <anexon@protonmail.com>
-" Co-author: w0rp <devw0rp@gmail.com>
+" Author: Anexon <anexon@protonmail.com>, w0rp <devw0rp@gmail.com>
 " Description: Utils and helpers with API normalised between Neovim/Vim 8 and
 " platform independent.
 

From 85cea4533dd61d8ef077618042ded8d5cc6a2b78 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 23:24:21 +0000
Subject: [PATCH 13/19] tidy: cleanup lua neural module init

---
 lua/neural.lua | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/lua/neural.lua b/lua/neural.lua
index 4d0ced0..7ccb645 100644
--- a/lua/neural.lua
+++ b/lua/neural.lua
@@ -1,24 +1,11 @@
 -- External dependencies
 local has_nui, _ = pcall(require, 'nui.input')
-local has_significant, _ = pcall(require, 'significant')
-local has_notify, _ = pcall(require, 'notify')
-
-local UI = {}
-local AnimatedSign = {}
-local Notify = {}
+local has_significant, AnimatedSign = pcall(require, 'significant')
 
 if has_nui then
     UI = require('neural.ui')
 end
 
-if has_significant then
-    AnimatedSign = require('significant')
-end
-
-if has_notify then
-  Notify = require('notify')
-end
-
 local Neural = {}
 
 function Neural.setup(settings)

From 5bced5cc29b658be4abd5bf4bb264d0b66165f51 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 23:25:10 +0000
Subject: [PATCH 14/19] fix: token count error message spelling

---
 autoload/neural/count.vim | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autoload/neural/count.vim b/autoload/neural/count.vim
index 5256f56..382119f 100644
--- a/autoload/neural/count.vim
+++ b/autoload/neural/count.vim
@@ -44,7 +44,7 @@ function! neural#count#SelectedLines() abort
         \}
         call neural#job#SendRaw(l:job_id, json_encode(l:input) . "\n")
     else
-        call neural#OutputErrorMessage('Failed to cound tokens')
+        call neural#OutputErrorMessage('Failed to count tokens')
 
         return
     endif

From ab093ad2dc1533f0585d89e601af8663f08e3420 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 23:25:36 +0000
Subject: [PATCH 15/19] feat: add improved notifications for token count

---
 autoload/neural/count.vim | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/autoload/neural/count.vim b/autoload/neural/count.vim
index 382119f..24093d2 100644
--- a/autoload/neural/count.vim
+++ b/autoload/neural/count.vim
@@ -5,18 +5,23 @@ 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 = a:job_data.output_lines[0]
+
         if has('nvim')
-            execute 'lua require(''neural'').notify("' . a:job_data.output_lines[0] . '", "info")'
+            execute 'lua require(''neural.notify'').info("' . l:output . '")'
         else
-            call neural#preview#Show(
-            \   a:job_data.output_lines[0],
-            \   {'stay_here': 1},
-            \)
+            call neural#preview#Show(l:output, {'stay_here': 1})
         endif
+    " Handle error
     else
-        call neural#OutputErrorMessage(join(a:job_data.error_lines, "\n"))
+        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

From c57e738b3c01a4a6003d639b2c69cd2a4bbdb849 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Wed, 3 Jan 2024 23:27:15 +0000
Subject: [PATCH 16/19] docs: update README about nvim-notify integration

---
 README.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index bf5178c..e351615 100644
--- a/README.md
+++ b/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
 

From 1602324cee7f90ad9950df267d20ee8625c79b64 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Thu, 4 Jan 2024 00:01:17 +0000
Subject: [PATCH 17/19] fix: Add context to the token count message

---
 autoload/neural/count.vim | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autoload/neural/count.vim b/autoload/neural/count.vim
index 24093d2..39005a2 100644
--- a/autoload/neural/count.vim
+++ b/autoload/neural/count.vim
@@ -5,7 +5,7 @@ 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 = a:job_data.output_lines[0]
+        let l:output = 'Tokens: ' . a:job_data.output_lines[0]
 
         if has('nvim')
             execute 'lua require(''neural.notify'').info("' . l:output . '")'

From 4e61a909fc69f5077edeb15d308e5c53d4f52252 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Sun, 2 Jun 2024 14:34:32 +0100
Subject: [PATCH 18/19] Dev: Async write to buffer WIP

---
 autoload/neural.vim         |  46 +------------
 autoload/neural/config.vim  |   4 +-
 autoload/neural/handler.vim | 126 ++++++++++++++++++++++++++++++++++++
 autoload/neural/job.vim     |  16 ++++-
 build/init.lua              |  60 +++++++++++++++++
 lua/.lua-format             |  29 +++++++++
 lua/neural/notify.lua       |  47 ++++++++++++++
 neural_sources/openai.py    |  71 ++++++++++++++------
 8 files changed, 331 insertions(+), 68 deletions(-)
 create mode 100644 autoload/neural/handler.vim
 create mode 100644 build/init.lua
 create mode 100644 lua/.lua-format
 create mode 100644 lua/neural/notify.lua

diff --git a/autoload/neural.vim b/autoload/neural.vim
index 55b53f4..fb11011 100644
--- a/autoload/neural.vim
+++ b/autoload/neural.vim
@@ -51,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
@@ -313,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)},
     \})
diff --git a/autoload/neural/config.vim b/autoload/neural/config.vim
index 933ee4f..e7f4c75 100644
--- a/autoload/neural/config.vim
+++ b/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,
diff --git a/autoload/neural/handler.vim b/autoload/neural/handler.vim
new file mode 100644
index 0000000..b863dd3
--- /dev/null
+++ b/autoload/neural/handler.vim
@@ -0,0 +1,126 @@
+" Author: Anexon <anexon@protonmail.com>
+" Description: APIs for working with Asynchronous jobs, with an API normalised
+" between Vim 8 and NeoVim.
+
+function! neural#handler#AddTextToBuffer(buffer, job_data, stream_data) abort
+    if bufnr('') isnot a:buffer && !exists('*appendbufline')
+    \ || !a:job_data.content_started && len(a:stream_data) == 0
+    \ || a:job_data.content_ended
+        return
+    endif
+
+    let l:leader = ' 🔸🔶'
+
+    " 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.
+    for text in a:stream_data
+        " Don't write empty or null characters to the buffer.
+        " Replace null characters (^@) with nothing or a space, depending on your needs
+        let text = substitute(text, '\%x03', '<X>', 'g')
+        let text = substitute(text, '\%x00', '', 'g')
+        " let text = substitute(text, '\\n', '|||', 'g')
+
+        " echoerr text
+        " if text is? '' || match(text, '\%x10') != -1 || match(text, '\%x00') != -1
+        if text is? ''
+            continue
+        elseif text is? '\n' || text is? '<X>'|| match(text, '\%x03') != -1 || text is? '^C'
+            " 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(a:job_data.moving_line, '')
+            else
+                call appendbufline(a:buffer, a:job_data.moving_line, '')
+            endif
+
+            " Cleanup leader
+            let l:line_content = getbufline(a:buffer, a:job_data.moving_line)
+            call setbufline(a:buffer, a:job_data.moving_line, l:line_content[0][0:-len(l:leader)-1])
+
+            " call setpos('.', getpos('.')[1])
+
+            " 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 += 1
+        " 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
+    endfor
+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
diff --git a/autoload/neural/job.vim b/autoload/neural/job.vim
index 4235f6f..59f34be 100644
--- a/autoload/neural/job.vim
+++ b/autoload/neural/job.vim
@@ -28,9 +28,21 @@ endfunction
 
 function! s:JoinNeovimOutput(job, last_line, data, mode, callback) abort
     if a:mode is# 'raw'
-        call a:callback(a:job, join(a:data, "\n"))
+        " 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
+        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 ''
+        return
     endif
 
     let l:lines = a:data[:-2]
diff --git a/build/init.lua b/build/init.lua
new file mode 100644
index 0000000..61d0b1f
--- /dev/null
+++ b/build/init.lua
@@ -0,0 +1,60 @@
+local function is_windows()
+    return vim.fn.has("win32") == 1 or vim.fn.has("win64") == 1
+end
+
+local function get_path_separator()
+    return is_windows() and "\\" or "/"
+end
+
+local function get_script_path()
+    local path_sep = get_path_separator()
+    local script_path = debug.getinfo(1).source:match("@?(.*)" .. path_sep)
+
+    return script_path or ""
+end
+
+local function run_command(command)
+    local status, _, code = os.execute(command)
+
+    require('notify').notify('Command: ' .. command .. ' || status: ' .. status, 'info')
+    require('notify').notify('status: ' .. status, 'info')
+
+    if status == false then
+        error("Command failed: " .. command .. ". Exit code: " .. tostring(code))
+    end
+end
+
+local function setup()
+    local python = is_windows() and "python" or "python3"
+    local sep = get_path_separator()
+    local script_path = get_script_path()
+    local venv_path = script_path .. sep .. "python3" .. sep .. "venv"
+
+    -- Create virtual environment if it does not exist.
+    if vim.fn.isdirectory(venv_path) == 0 then
+        run_command(python .. " -m venv " .. venv_path)
+    end
+
+    -- Install requirements via pip
+    local pip_cmd = venv_path .. (is_windows() and sep .. "Scripts" .. sep .. "pip" or sep .. "bin" .. sep .. "pip")
+    local requirements_path = script_path .. sep .. "python3" .. sep .. "requirements.txt"
+
+    run_command(pip_cmd .. " install -r " .. requirements_path)
+end
+
+setup()
+
+local function setup()
+    local setup_script = debug.getinfo(1).source:match("@?(.*/)") .. "setup.sh"
+
+    if vim.fn.filereadable(setup_script) == 1 then
+        local status, _, code = os.execute("sh " .. build_script_path)
+        if status == false then
+            error("Failed to execute build script. Exit code: " .. tostring(code))
+        end
+    else
+        error("Build script not found: " .. build_script_path)
+    end
+end
+
+setup()
diff --git a/lua/.lua-format b/lua/.lua-format
new file mode 100644
index 0000000..b8ac3d8
--- /dev/null
+++ b/lua/.lua-format
@@ -0,0 +1,29 @@
+column_limit: 120
+indent_width: 4
+spaces_before_call: 1
+keep_simple_control_block_one_line: true
+keep_simple_function_one_line: true
+align_args: true
+break_after_functioncall_lp: false
+break_before_functioncall_rp: false
+spaces_inside_functioncall_parens: false
+spaces_inside_functiondef_parens: false
+align_parameter: true
+chop_down_parameter: false
+break_after_functiondef_lp: false
+break_before_functiondef_rp: false
+align_table_field: true
+break_after_table_lb: true
+break_before_table_rb: true
+chop_down_table: false
+chop_down_kv_table: true
+table_sep: ","
+column_table_limit: 0
+extra_sep_at_table_end: false
+spaces_inside_table_braces: false
+break_after_operator: true
+double_quote_to_single_quote: false
+single_quote_to_double_quote: false
+spaces_around_equals_in_field: true
+line_breaks_after_function_body: 1
+line_separator: input
diff --git a/lua/neural/notify.lua b/lua/neural/notify.lua
new file mode 100644
index 0000000..c9ecfc6
--- /dev/null
+++ b/lua/neural/notify.lua
@@ -0,0 +1,47 @@
+-- Author: Anexon <anexon@protonmail.com>
+-- Description: Show messages with optional nvim-notify integration.
+
+-- nvim-notify plugin
+local has_notify, Notify = pcall(require, 'notify')
+
+local M = {}
+
+local opts = {
+    title = 'Neural',
+    icon = vim.g.neural.ui.prompt_icon
+}
+
+-- Show a message with nvim-notify or fallback.
+--- @param message string
+--- @param level string Level following vim.log.levels spec.
+function M.show_message(message, level)
+    if has_notify then
+        Notify(message, level, opts)
+    else
+      vim.fn['neural#preview#Show'](message, {stay_here = 1})
+    end
+end
+
+-- Show info message.
+--- @param message string
+function M.info(message)
+    M.show_message(message, 'info')
+end
+
+-- Show warning message.
+--- @param message string
+function M.warn(message)
+    M.show_message(message, 'warn')
+end
+
+-- Show error message.
+--- @param message string
+function M.error(message)
+    if has_notify then
+        Notify(message, 'error', opts)
+    else
+      vim.fn['neural#OutputErrorMessage'](message)
+    end
+end
+
+return M
diff --git a/neural_sources/openai.py b/neural_sources/openai.py
index 5e31eb5..2402483 100644
--- a/neural_sources/openai.py
+++ b/neural_sources/openai.py
@@ -3,17 +3,21 @@
 """
 import json
 import platform
+import re
 import ssl
 import sys
 import urllib.error
 import urllib.request
 from typing import Any, Dict
 
-API_ENDPOINT = 'https://api.openai.com/v1/completions'
+API_ENDPOINT = "https://api.openai.com/v1/completions"
 
-OPENAI_DATA_HEADER = 'data: '
-OPENAI_DONE = '[DONE]'
+OPENAI_DATA_HEADER = "data: "
+OPENAI_DONE = "[DONE]"
+END_OF_STREAM = "<<[EOF]>>"
 
+ETX = chr(3) # End of Text - Signals end of text for some line buffer.
+EOT = chr(4) # End of Transmission - Signals end of text to write to buffer.
 
 class Config:
     """
@@ -37,15 +41,29 @@ def __init__(
         self.presence_penalty = presence_penalty
         self.frequency_penalty = frequency_penalty
 
+def format_prompt(prompt: str) -> str:
+    """
+    OpenAI models use `<|endoftext|>` as the document separator during
+    training. The completion models will attempt to complete the prompt before
+    returning a response for the prompt via the completion API with `\n\n`.
+
+    Appending `\n\n` to the prompt ensures the model responds with only a pure
+    completion response. Any proceeding newline characters are therefore
+    considered intentional to the response.
+    """
+    return prompt + '\n\n'
 
 def get_openai_completion(config: Config, prompt: str) -> None:
+    """
+    Get a completion API response from a given prompt.
+    """
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {config.api_key}"
+        "Authorization": f"Bearer {config.api_key}",
     }
     data = {
         "model": config.model,
-        "prompt": prompt,
+        "prompt": format_prompt(prompt),
         "temperature": config.temperature,
         "max_tokens": config.max_tokens,
         "top_p": 1,
@@ -69,8 +87,8 @@ def get_openai_completion(config: Config, prompt: str) -> None:
     # urllib.error.URLError: <urlopen error [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate (_ssl.c:997)>  # noqa
     context = (
         ssl._create_unverified_context()  # type: ignore
-        if platform.system() == "Darwin" else
-        None
+        if platform.system() == "Darwin"
+        else None
     )
 
     with urllib.request.urlopen(req, context=context) as response:
@@ -83,16 +101,27 @@ def get_openai_completion(config: Config, prompt: str) -> None:
             line = line_bytes.decode("utf-8", errors="replace")
 
             if line.startswith(OPENAI_DATA_HEADER):
-                line_data = line[len(OPENAI_DATA_HEADER):-1]
+                line_data = line[len(OPENAI_DATA_HEADER) : -1]
 
                 if line_data == OPENAI_DONE:
                     pass
                 else:
                     openai_obj = json.loads(line_data)
+                    openai_text = openai_obj["choices"][0]["text"]
+
+                    # Split the text at each newline, keeping the newline characters
+                    split_text = re.split(r"(\n)", openai_text)
 
-                    print(openai_obj["choices"][0]["text"], end="", flush=True)
+                    for segment in split_text:
+                        # time.sleep(0.05)
+                        if segment == "\n":
+                            # Signal End of Text for buffer line.
+                            print(end=ETX, flush=True)
+                        else:
+                            print(segment, flush=True)
 
-    print()
+    # Signal End of Transmission.
+    print(end=EOT, flush=True)
 
 
 def load_config(raw_config: Dict[str, Any]) -> Config:
@@ -100,37 +129,37 @@ def load_config(raw_config: Dict[str, Any]) -> Config:
     if not isinstance(raw_config, dict):  # type: ignore
         raise ValueError("openai config is not a dictionary")
 
-    api_key = raw_config.get('api_key')
+    api_key = raw_config.get("api_key")
 
     if not isinstance(api_key, str) or not api_key:  # type: ignore
         raise ValueError("openai.api_key is not defined")
 
-    model = raw_config.get('model')
+    model = raw_config.get("model")
 
     if not isinstance(model, str) or not model:
         raise ValueError("openai.model is not defined")
 
-    temperature = raw_config.get('temperature', 0.2)
+    temperature = raw_config.get("temperature", 0.2)
 
     if not isinstance(temperature, (int, float)):
         raise ValueError("openai.temperature is invalid")
 
-    top_p = raw_config.get('top_p', 1)
+    top_p = raw_config.get("top_p", 1)
 
     if not isinstance(top_p, (int, float)):
         raise ValueError("openai.top_p is invalid")
 
-    max_tokens = raw_config.get('max_tokens', 1024)
+    max_tokens = raw_config.get("max_tokens", 1024)
 
     if not isinstance(max_tokens, (int)):
         raise ValueError("openai.max_tokens is invalid")
 
-    presence_penalty = raw_config.get('presence_penalty', 0)
+    presence_penalty = raw_config.get("presence_penalty", 0)
 
     if not isinstance(presence_penalty, (int, float)):
         raise ValueError("openai.presence_penalty is invalid")
 
-    frequency_penalty = raw_config.get('frequency_penalty', 0)
+    frequency_penalty = raw_config.get("frequency_penalty", 0)
 
     if not isinstance(frequency_penalty, (int, float)):
         raise ValueError("openai.frequency_penalty is invalid")
@@ -147,7 +176,7 @@ def load_config(raw_config: Dict[str, Any]) -> Config:
 
 
 def get_error_message(error: urllib.error.HTTPError) -> str:
-    message = error.read().decode('utf-8', errors='ignore')
+    message = error.read().decode("utf-8", errors="ignore")
 
     try:
         # JSON data might look like this:
@@ -159,10 +188,10 @@ def get_error_message(error: urllib.error.HTTPError) -> str:
         #       "code": null
         #   }
         # }
-        message = json.loads(message)['error']['message']
+        message = json.loads(message)["error"]["message"]
 
         if "This model's maximum context length is" in message:
-            message = 'Too much text for a request!'
+            message = "Too much text for a request!"
     except Exception:
         # If we can't get a better message use the JSON payload at least.
         pass
@@ -183,7 +212,7 @@ def main() -> None:
     except urllib.error.HTTPError as error:
         if error.code == 400:
             message = get_error_message(error)
-            sys.exit('Neural error: OpenAI request failure: ' + message)
+            sys.exit("Neural error: OpenAI request failure: " + message)
         elif error.code == 429:
             sys.exit("Neural error: OpenAI request limit reached!")
         else:

From 8f11c79daa1a61b9ee13e81a9cc24fd8efa4f122 Mon Sep 17 00:00:00 2001
From: Anexon <anexon@protonmail.com>
Date: Thu, 6 Jun 2024 00:35:30 +0100
Subject: [PATCH 19/19] Dev: WIP change async write callback

---
 autoload/neural/handler.vim | 148 +++++++++++++++++++-----------------
 autoload/neural/job.vim     |  16 ++--
 neural_sources/openai.py    |  23 +++---
 3 files changed, 100 insertions(+), 87 deletions(-)

diff --git a/autoload/neural/handler.vim b/autoload/neural/handler.vim
index b863dd3..0dd0bfa 100644
--- a/autoload/neural/handler.vim
+++ b/autoload/neural/handler.vim
@@ -1,87 +1,97 @@
+scriptencoding utf8
+
 " Author: Anexon <anexon@protonmail.com>
 " 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')
-    \ || !a:job_data.content_started && len(a:stream_data) == 0
-    \ || a:job_data.content_ended
+    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.
-    for text in a:stream_data
-        " Don't write empty or null characters to the buffer.
-        " Replace null characters (^@) with nothing or a space, depending on your needs
-        let text = substitute(text, '\%x03', '<X>', 'g')
-        let text = substitute(text, '\%x00', '', 'g')
-        " let text = substitute(text, '\\n', '|||', 'g')
-
-        " echoerr text
-        " if text is? '' || match(text, '\%x10') != -1 || match(text, '\%x00') != -1
-        if text is? ''
-            continue
-        elseif text is? '\n' || text is? '<X>'|| match(text, '\%x03') != -1 || text is? '^C'
-            " 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(a:job_data.moving_line, '')
-            else
-                call appendbufline(a:buffer, a:job_data.moving_line, '')
-            endif
-
-            " Cleanup leader
-            let l:line_content = getbufline(a:buffer, a:job_data.moving_line)
-            call setbufline(a:buffer, a:job_data.moving_line, l:line_content[0][0:-len(l:leader)-1])
-
-            " call setpos('.', getpos('.')[1])
-
-            " 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 += 1
-        " 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
-    endfor
+    " 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
diff --git a/autoload/neural/job.vim b/autoload/neural/job.vim
index 59f34be..2a0450a 100644
--- a/autoload/neural/job.vim
+++ b/autoload/neural/job.vim
@@ -34,13 +34,15 @@ function! s:JoinNeovimOutput(job, last_line, data, mode, callback) abort
         " Each stream item is passed to the callback individually which can be
         " a chunk of text or a newline character.
         " echoerr a:data
-        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
+        call a:callback(a:job, join(a:data, "\n"))
+
+        " 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
diff --git a/neural_sources/openai.py b/neural_sources/openai.py
index 2402483..4f1c5bb 100644
--- a/neural_sources/openai.py
+++ b/neural_sources/openai.py
@@ -3,7 +3,6 @@
 """
 import json
 import platform
-import re
 import ssl
 import sys
 import urllib.error
@@ -109,19 +108,21 @@ def get_openai_completion(config: Config, prompt: str) -> None:
                     openai_obj = json.loads(line_data)
                     openai_text = openai_obj["choices"][0]["text"]
 
-                    # Split the text at each newline, keeping the newline characters
-                    split_text = re.split(r"(\n)", openai_text)
-
-                    for segment in split_text:
+                    print(openai_text, end="", flush=True)
+                    # # Split the text at each newline, keeping the newline characters
+                    # split_text = re.split(r"(\n)", openai_text)
+                    #
+                    # for segment in split_text:
+                    #     print(segment, flush=True)
                         # time.sleep(0.05)
-                        if segment == "\n":
-                            # Signal End of Text for buffer line.
-                            print(end=ETX, flush=True)
-                        else:
-                            print(segment, flush=True)
+                        # if segment == "\n":
+                        #     # Signal End of Text for buffer line.
+                        #     print(end=ETX, flush=True)
+                        # else:
+                        #     print(segment, flush=True)
 
     # Signal End of Transmission.
-    print(end=EOT, flush=True)
+    # print(end=EOT, flush=True)
 
 
 def load_config(raw_config: Dict[str, Any]) -> Config: