llm.lua (4544B)
1 local ESC_FEEDKEY = vim.api.nvim_replace_termcodes("<ESC>", true, false, true) 2 local ollama = { 3 4 --- @param model string 5 --- @param context string 6 --- @param prompt string 7 --- @param buf_nr number 8 --- @param max_width number 9 --- @return number 10 run = function(model, context, prompt, buf_nr, max_width) 11 local cmd = ("ollama run $model $prompt") 12 :gsub("$model", model) 13 :gsub("$prompt", vim.fn.shellescape(context .. "\n" .. prompt)) 14 local header = vim.split(prompt, "\n") 15 table.insert(header, "----------------------------------------") 16 vim.api.nvim_buf_set_lines(buf_nr, 0, -1, false, header) 17 local line = vim.tbl_count(header) 18 local line_char_count = 0 19 local words = {} 20 return vim.fn.jobstart(cmd, { 21 on_stdout = function(_, data, _) 22 for i, token in ipairs(data) do 23 if i > 1 or (string.match(token, "^%s") and line_char_count > max_width) then -- if returned data array has more than one element, a line break occured. 24 line = line + 1 25 words = {} 26 line_char_count = 0 27 end 28 line_char_count = line_char_count + #token 29 table.insert(words, token) 30 vim.api.nvim_buf_set_lines(buf_nr, line, line + 1, false, { table.concat(words, "") }) 31 end 32 end, 33 }) 34 end, 35 36 --- @return string 37 get_visual_lines = function(bufnr) 38 vim.api.nvim_feedkeys(ESC_FEEDKEY, "n", true) 39 vim.api.nvim_feedkeys("gv", "x", false) 40 vim.api.nvim_feedkeys(ESC_FEEDKEY, "n", true) 41 42 local start_row, start_col = unpack(vim.api.nvim_buf_get_mark(bufnr, "<")) 43 local end_row, end_col = unpack(vim.api.nvim_buf_get_mark(bufnr, ">")) 44 local lines = vim.api.nvim_buf_get_lines(bufnr, start_row - 1, end_row, false) 45 46 if start_row == 0 then 47 lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false) 48 start_row = 1 49 start_col = 0 50 end_row = #lines 51 end_col = #lines[#lines] 52 end 53 54 start_col = start_col + 1 55 end_col = math.min(end_col, #lines[#lines] - 1) + 1 56 57 lines[#lines] = lines[#lines]:sub(1, end_col) 58 lines[1] = lines[1]:sub(start_col) 59 60 return table.concat(lines, "\n") 61 end 62 } 63 local plen = require("plenary.window.float") 64 -- define some styles and spawn a scratch buffer 65 local win_options = { 66 winblend = 10, 67 border = "rounded", 68 bufnr = vim.api.nvim_create_buf(false, true) 69 } 70 71 vim.keymap.set("n", "<leader>cc", function() 72 local prompt = vim.fn.input("Prompt: ") 73 if prompt == "" then 74 return 75 end 76 plen.clear(win_options.bufnr) 77 local float = plen.percentage_range_window(0.8, 0.8, win_options) 78 local win_width = vim.api.nvim_win_get_width(float.win_id) - 5 79 ollama.run("mistral:instruct", "", prompt, float.bufnr, win_width) 80 end, { silent = true, desc = "Mistral" }) 81 82 vim.keymap.set("v", "<leader>cp", function() 83 local prompt = ollama.get_visual_lines(0) 84 plen.clear(win_options.bufnr) 85 local float = plen.percentage_range_window(0.8, 0.8, win_options) 86 local win_width = vim.api.nvim_win_get_width(float.win_id) - 5 87 ollama.run("mistral:instruct", "", prompt, float.bufnr, win_width) 88 end, { silent = true, desc = "Mistral" }) 89 90 ------------------------------------------------------------------ 91 -- coding 92 vim.keymap.set("v", "<leader>cc", function() 93 local context = vim.fn.input("Prompt: ") 94 if context == "" then 95 return 96 end 97 local prompt = ollama.get_visual_lines(0) 98 plen.clear(win_options.bufnr) 99 local float = plen.percentage_range_window(0.8, 0.8, win_options) 100 local win_width = vim.api.nvim_win_get_width(float.win_id) - 5 101 ollama.run("codellama:7b", context, prompt, float.bufnr, win_width) 102 end, { silent = true, desc = "Codellama" }) 103 104 ------------------------------------------------------------------ 105 -- general purpose math & logic questions 106 vim.keymap.set("n", "<leader>cm", function() 107 local prompt = vim.fn.input("Prompt: ") 108 if prompt == "" then 109 return 110 end 111 plen.clear(win_options.bufnr) 112 local float = plen.percentage_range_window(0.8, 0.8, win_options) 113 local win_width = vim.api.nvim_win_get_width(float.win_id) - 5 114 ollama.run("wizard-math", "", prompt, float.bufnr, win_width) 115 end, { silent = true, desc = "Math" }) 116 117 vim.keymap.set("v", "<leader>cm", function() 118 local prompt = ollama.get_visual_lines(0) 119 plen.clear(win_options.bufnr) 120 local float = plen.percentage_range_window(0.8, 0.8, win_options) 121 local win_width = vim.api.nvim_win_get_width(float.win_id) - 5 122 ollama.run("wizard-math", "", prompt, float.bufnr, win_width) 123 end, { silent = true, desc = "Math" })