feat: add api function to cancel request

Signed-off-by: blob42 <[email protected]>

Author: blob42

lua/codegpt.lua

@@ -1,5 +1,6 @@
-local config = require("codegpt.config")
-local models = require("codegpt.models")
+local Config = require("codegpt.config")
+local Models = require("codegpt.models")
+local Api = require("codegpt.api")
 local M = {}
 
 local Commands = require("codegpt.commands")
@@ -17,16 +18,16 @@ end
 
 function M.run_cmd(opts)
 	if opts.name and opts.name:match("^V") then
-		config.popup_override = "vertical"
+		Config.popup_override = "vertical"
 	else
-		config.popup_override = nil
+		Config.popup_override = nil
 	end
 
 	-- bang makes popup persistent until closed
 	if opts.bang then
-		config.persistent_override = true
+		Config.persistent_override = true
 	else
-		config.persistent_override = false
+		Config.persistent_override = false
 	end
 
 	local text_selection, bounds = Utils.get_selected_lines(opts)
@@ -46,7 +47,7 @@ function M.run_cmd(opts)
 			command_args = ""
 		elseif text_selection == "" then
 			command = "chat"
-		elseif config.opts.commands[command] == nil then
+		elseif Config.opts.commands[command] == nil then
 			command = "code_edit"
 		end
 	elseif text_selection ~= "" and command_args == "" then
@@ -63,7 +64,8 @@ function M.run_cmd(opts)
 	Commands.run_cmd(command, command_args, text_selection, bounds)
 end
 
-M.setup = config.setup
-M.select_model = models.select_model
+M.setup = Config.setup
+M.select_model = Models.select_model
+M.cancel_request = Api.cancel_job
 
 return M

lua/codegpt/api.lua

@@ -1,8 +1,9 @@
 local Config = require("codegpt.config")
 
-local Api = {}
+local M = {}
 
-CODEGPT_CALLBACK_COUNTER = 0
+local CODEGPT_CALLBACK_COUNTER = 0
+M.current_job = nil
 
 local status_index = 0
 local timer = vim.uv.new_timer()
@@ -18,7 +19,7 @@ local function start_spinner_timer()
 	)
 end
 
-function Api.get_status(...)
+function M.get_status(...)
 	local spinners = Config.opts.ui.spinners or { "", "", "", "", "", "" }
 	local spinner_speed = Config.opts.ui.spinner_speed or 80
 	local ms = vim.uv.hrtime() / 1000000
@@ -40,15 +41,15 @@ function Api.get_status(...)
 	end
 end
 
-function Api.run_started_hook()
+function M.run_started_hook()
 	if Config.opts.hooks.request_started ~= nil then
 		Config.opts.hooks.request_started()
 	end
 
 	CODEGPT_CALLBACK_COUNTER = CODEGPT_CALLBACK_COUNTER + 1
 end
 
-function Api.run_finished_hook()
+function M.run_finished_hook()
 	if CODEGPT_CALLBACK_COUNTER > 0 then
 		CODEGPT_CALLBACK_COUNTER = CODEGPT_CALLBACK_COUNTER - 1
 	end
@@ -59,4 +60,11 @@ function Api.run_finished_hook()
 	end
 end
 
-return Api
+function M.cancel_job()
+	if M.current_job ~= nil then
+		M.current_job:shutdown()
+		M.run_finished_hook()
+	end
+end
+
+return M

lua/codegpt/providers/anthropic.lua

@@ -112,7 +112,7 @@ function M.make_call(payload, cb)
 	local url = "https://api.anthropic.com/v1/messages"
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		callback = function(response)

lua/codegpt/providers/azure.lua

@@ -198,7 +198,7 @@ function M.make_call(payload, cb)
 	local url = Config.opts.connection.chat_completions_url
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		callback = function(response)

lua/codegpt/providers/groq.lua

@@ -126,7 +126,7 @@ function M.make_call(payload, cb)
 	local url = "https://api.groq.com/openai/v1/chat/completions"
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		callback = function(response)

lua/codegpt/providers/ollama.lua

@@ -3,7 +3,7 @@ local Render = require("codegpt.template_render")
 local Utils = require("codegpt.utils")
 local Api = require("codegpt.api")
 local Config = require("codegpt.config")
-local Tokens = require("codegpt.tokens")
+local tokens = require("codegpt.tokens")
 local errors = require("codegpt.errors")
 
 local M = {}
@@ -30,7 +30,7 @@ local function generate_messages(command, cmd_opts, command_args, text_selection
 end
 
 local function get_max_tokens(max_tokens, prompt)
-	local total_length = Tokens.get_tokens(prompt)
+	local total_length = tokens.get_tokens(prompt)
 
 	if total_length >= max_tokens then
 		error("Total length of messages exceeds max_tokens: " .. total_length .. " > " .. max_tokens)
@@ -160,7 +160,7 @@ function M.make_call(payload, cb)
 	local url = Config.opts.connection.ollama_base_url:gsub("/$", "") .. "/api/chat"
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		callback = function(response)
@@ -173,13 +173,13 @@ function M.make_call(payload, cb)
 end
 
 ---@param payload table payload sent to api
----@param stream_cb fun(data: table) callback to handle the resonse json stream
+---@param stream_cb fun(data: table, job: table) callback to handle the resonse json stream
 function M.make_stream_call(payload, stream_cb)
 	local payload_str = vim.fn.json_encode(payload)
 	local url = Config.opts.connection.ollama_base_url:gsub("/$", "") .. "/api/chat"
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		stream = function(error, data, job)

lua/codegpt/providers/openai.lua

@@ -3,7 +3,7 @@ local Render = require("codegpt.template_render")
 local Utils = require("codegpt.utils")
 local Api = require("codegpt.api")
 local Config = require("codegpt.config")
-local Tokens = require("codegpt.tokens")
+local tokens = require("codegpt.tokens")
 local errors = require("codegpt.errors")
 
 -- TODO: handle streaming mode
@@ -32,7 +32,7 @@ local function generate_messages(command, cmd_opts, command_args, text_selection
 end
 
 local function get_max_tokens(max_tokens, messages)
-	local total_length = Tokens.get_tokens(messages)
+	local total_length = tokens.get_tokens(messages)
 
 	if total_length >= max_tokens then
 		error("Total length of messages exceeds max_tokens: " .. total_length .. " > " .. max_tokens)
@@ -151,7 +151,7 @@ function M.make_call(payload, cb)
 	local url = Config.opts.connection.chat_completions_url
 	local headers = M.make_headers()
 	Api.run_started_hook()
-	curl.post(url, {
+	Api.current_job = curl.post(url, {
 		body = payload_str,
 		headers = headers,
 		callback = function(response)

lua/codegpt/ui.lua

@@ -1,6 +1,6 @@
 local Popup = require("nui.popup")
 local Split = require("nui.split")
-local config = require("codegpt.config")
+local Config = require("codegpt.config")
 local event = require("nui.utils.autocmd").event
 
 local M = {}
@@ -14,7 +14,7 @@ local function create_horizontal()
 		split = Split({
 			relative = "editor",
 			position = "bottom",
-			size = config.opts.ui.horizontal_popup_size,
+			size = Config.opts.ui.horizontal_popup_size,
 		})
 	end
 
@@ -26,7 +26,7 @@ local function create_vertical()
 		split = Split({
 			relative = "editor",
 			position = "right",
-			size = config.opts.ui.vertical_popup_size,
+			size = Config.opts.ui.vertical_popup_size,
 		})
 	end
 
@@ -35,15 +35,15 @@ end
 
 local function create_floating()
 	if not popup then
-		local window_options = config.opts.ui.popup_window_options
+		local window_options = Config.opts.ui.popup_window_options
 		if window_options == nil then
 			window_options = {}
 		end
 
 		local popupOpts = {
 			enter = true,
 			focusable = true,
-			border = config.opts.ui.popup_border,
+			border = Config.opts.ui.popup_border,
 			position = "50%",
 			size = {
 				width = "80%",
@@ -58,13 +58,13 @@ local function create_floating()
 		popup = Popup(popupOpts)
 	end
 
-	popup:update_layout(config.opts.ui.popup_options)
+	popup:update_layout(Config.opts.ui.popup_options)
 
 	return popup
 end
 
 local function create_window()
-	local popup_type = config.popup_override or config.opts.ui.popup_type
+	local popup_type = Config.popup_override or Config.opts.ui.popup_type
 	local ui_elem = nil
 	if popup_type == "horizontal" then
 		ui_elem = create_horizontal()
@@ -78,24 +78,27 @@ local function create_window()
 end
 
 function M.popup(job, lines, filetype, bufnr, start_row, start_col, end_row, end_col)
+	if job ~= nil and job.is_shutdown then
+		return
+	end
 	local ui_elem = create_window()
 	-- mount/open the component
 	ui_elem:mount()
 
-	if not (config.persistent_override or config.opts.ui.persistent) then
+	if not (Config.persistent_override or Config.opts.ui.persistent) then
 		-- unmount component when cursor leaves buffer
 		ui_elem:on(event.BufLeave, function()
 			ui_elem:unmount()
 		end)
 	end
 
 	-- unmount component when key 'q'
-	ui_elem:map("n", config.opts.ui.actions.quit, function()
+	ui_elem:map("n", Config.opts.ui.actions.quit, function()
 		ui_elem:unmount()
 	end, { noremap = true, silent = true })
 	--
 	-- cancel job if actions.cancel is called
-	ui_elem:map("n", config.opts.ui.actions.cancel, function()
+	ui_elem:map("n", Config.opts.ui.actions.cancel, function()
 		job:shutdown()
 	end, { noremap = true, silent = true })
 
@@ -104,19 +107,19 @@ function M.popup(job, lines, filetype, bufnr, start_row, start_col, end_row, end
 	vim.api.nvim_buf_set_lines(ui_elem.bufnr, 0, 1, false, lines)
 
 	-- replace lines when ctrl-o pressed
-	ui_elem:map("n", config.opts.ui.actions.use_as_output, function()
+	ui_elem:map("n", Config.opts.ui.actions.use_as_output, function()
 		vim.api.nvim_buf_set_text(bufnr, start_row, start_col, end_row, end_col, lines)
 		ui_elem:unmount()
 	end)
 
 	-- selecting all the content when ctrl-i is pressed
 	-- so the user can proceed with another API request
-	ui_elem:map("n", config.opts.ui.actions.use_as_input, function()
+	ui_elem:map("n", Config.opts.ui.actions.use_as_input, function()
 		vim.api.nvim_feedkeys("ggVG:Chat ", "n", false)
 	end, { noremap = false })
 
 	-- mapping custom commands
-	for _, command in ipairs(config.opts.ui.actions.custom) do
+	for _, command in ipairs(Config.opts.ui.actions.custom) do
 		ui_elem:map(command[1], command[2], command[3], command[4])
 	end
 end
@@ -125,27 +128,30 @@ local streaming = false
 local stream_ui_elem = nil
 
 function M.popup_stream(job, stream, filetype, bufnr, start_row, start_col, end_row, end_col)
+	if job ~= nil and job.is_shutdown then
+		return
+	end
 	if not streaming then
 		streaming = true
 		stream_ui_elem = create_window()
 
 		-- mount/open the component
 		stream_ui_elem:mount()
 
-		if not (config.persistent_override or config.opts.ui.persistent) then
+		if not (Config.persistent_override or Config.opts.ui.persistent) then
 			-- unmount component when cursor leaves buffer
 			stream_ui_elem:on(event.BufLeave, function()
 				stream_ui_elem:unmount()
 			end)
 		end
 
 		-- unmount component when key 'q'
-		stream_ui_elem:map("n", config.opts.ui.actions.quit, function()
+		stream_ui_elem:map("n", Config.opts.ui.actions.quit, function()
 			stream_ui_elem:unmount()
 		end, { noremap = true, silent = true })
 
 		-- cancel job if actions.cancel is called
-		stream_ui_elem:map("n", config.opts.ui.actions.cancel, function()
+		stream_ui_elem:map("n", Config.opts.ui.actions.cancel, function()
 			job:shutdown()
 		end, { noremap = true, silent = true })