feat: Add insert_after_selection and display_insert_after_selection actions

README.md

@@ -202,10 +202,14 @@ the prompt string with context in the following ways:
   - Uses the `extract` pattern to extract the response.
 - `insert`: Insert the response at the current cursor line
   - Uses the `extract` pattern to extract the response.
+- `insert_after_selection`: Insert the response after the previous selection.
+  - Uses the `extract` pattern to extract the response.
 - `display_replace`: Stream and display the response in a floating window, then replace the current selection with the response.
   - Uses the `extract` pattern to extract the response.
 - `display_insert`: Stream and display the response in a floating window, then insert the response at the current cursor line.
   - Uses the `extract` pattern to extract the response.
+- `display_insert_after_selection`: Stream and display the response in a floating window, then insert the response after the previous selection.
+  - Uses the `extract` pattern to extract the response.
 
 Sometimes, you may need functionality that is not provided by
 the built-in actions. In this case, you can write your own Custom Actions with the following interface:

lua/ollama/actions/factory.lua

@@ -3,6 +3,7 @@ local factory = {}
 ---@class Ollama.ActionFactoryBuildOpts
 ---@field display boolean? whether to display the response (default: true)
 ---@field insert boolean? whether to insert the response at the cursor (default: false)
+---@field insert_after_selection boolean? whether to insert the response after the previous selection (default: false)
 ---@field replace boolean? whether to replace the selection with the response. Precedes `insert` (default: false)
 ---@field show_prompt boolean? whether to prepend the display buffer with the parsed prompt (default: false)
 ---@field window "float" | "split" | "vsplit" | nil type of window to display the response in (default: "float") (NOT YET IMPLEMENTED)
@@ -11,6 +12,7 @@ local factory = {}
 local default_opts = {
 	display = true,
 	insert = false,
+	insert_after_selection = false,
 	replace = false,
 	show_prompt = false,
 	window = "float",
@@ -32,19 +34,21 @@ function factory.create_action(opts)
 			local out_win
 			local timer
 			local pre_lines
+			local bufnr
+      bufnr = vim.fn.bufnr("%") or 0
 
 			-- stuff for insert
-			local bufnr
-			local cursorLine
+			local insertLine
 			if opts.insert then
-				bufnr = vim.fn.bufnr("%") or 0
-				cursorLine = vim.fn.line(".") or 1
+				insertLine = vim.fn.line(".") or 1
+			end
+			if opts.insert_after_selection then
+				insertLine = vim.api.nvim_buf_get_mark(0, ">")[1] or 1
 			end
 
 			-- stuff for replace
 			local sel_pos
 			if opts.replace then
-				bufnr = vim.fn.bufnr("%") or 0
 				sel_pos = require("ollama.util").get_selection_pos()
 				if sel_pos == nil then
 					vim.api.nvim_notify("No selection found", vim.log.levels.INFO, { title = "Ollama" })
@@ -127,7 +131,7 @@ function factory.create_action(opts)
 						vim.api.nvim_set_option_value("modifiable", false, { buf = out_buf })
 					end
 
-					if opts.insert or opts.replace then
+					if opts.insert or opts.insert_after_selection or opts.replace then
 						local text = table.concat(lines, "\n")
 						if prompt.extract then
 							text = text:match(prompt.extract)
@@ -143,8 +147,8 @@ function factory.create_action(opts)
 						if opts.replace then
 							local start_line, start_col, end_line, end_col = unpack(sel_pos)
 							vim.api.nvim_buf_set_text(bufnr, start_line, start_col, end_line, end_col, lines)
-						elseif opts.insert then
-							vim.api.nvim_buf_set_lines(bufnr, cursorLine, cursorLine, false, lines)
+						elseif opts.insert or opts.insert_after_selection then
+							vim.api.nvim_buf_set_lines(bufnr, insertLine, insertLine, false, lines)
 						end
 
 						-- close floating window when done insert/replacing

lua/ollama/actions/init.lua

@@ -7,6 +7,8 @@ actions.display = factory.create_action({ display = true, show_prompt = true })
 
 actions.insert = factory.create_action({ display = false, insert = true })
 
+actions.insert_after_selection = factory.create_action({ display = false, insert_after_selection = true })
+
 actions.replace = factory.create_action({ display = false, replace = true })
 
 -- basically a merge of display -> replace actions
@@ -18,6 +20,8 @@ actions.display_replace = factory.create_action({
 
 actions.display_insert = factory.create_action({ insert = true, show_prompt = true })
 
+actions.display_insert_after_selection = factory.create_action({ insert_after_selection = true, show_prompt = true })
+
 -- TODO: remove this as its not used anymore
 -- if you use this in your config, please switch to "display" instead
 actions.display_prompt = actions.display

lua/ollama/init.lua

@@ -41,7 +41,7 @@ local M = {}
 ---@field top_p float? Works together with top-k. A higher value (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text. (Default: 0.9)
 
 ---Built-in actions
----@alias Ollama.PromptActionBuiltinEnum "display" | "replace" | "insert" | "display_replace" | "display_insert" | "display_prompt"
+---@alias Ollama.PromptActionBuiltinEnum "display" | "replace" | "insert" | "insert_after_selection" | "display_replace" | "display_insert" | "display_insert_after_selection" | "display_prompt"
 
 -- Handles the output of a prompt. Custom Actions can be defined in lieu of a builtin.
 ---@alias Ollama.PromptAction table | Ollama.PromptActionFields

lua/ollama/prompts.lua

@@ -31,6 +31,13 @@ local prompts = {
 		action = "replace",
 	},
 
+	Comment_Code = {
+		prompt = "Rewrite the following $ftype code."
+			.. response_format
+			.. "\n\n```$ftype\n$sel```",
+		action = "insert_after_selection",
+	},
+
 	Generate_Code = {
 		prompt = "Generate $ftype code that does the following: $input\n\n" .. response_format,
 		action = "insert",