diff --git a/lua/gp/init.lua b/lua/gp/init.lua
index 385839a..1c89fdd 100644
--- a/lua/gp/init.lua
+++ b/lua/gp/init.lua
@@ -1310,7 +1310,8 @@ end
 ---@param payload table # payload for api
 ---@param handler function # response handler
 ---@param on_exit function | nil # optional on_exit handler
-M.query = function(buf, provider, payload, handler, on_exit)
+---@param on_complete_callback function | nil # optional on_complete_callback handler
+M.query = function(buf, provider, payload, handler, on_exit, on_complete_callback)
 	-- make sure handler is a function
 	if type(handler) ~= "function" then
 		M.error(
@@ -1430,6 +1431,13 @@ M.query = function(buf, provider, payload, handler, on_exit)
 						end)
 					end
 				end
+
+				-- optional on_complete_callback handler
+				if type(on_complete_callback) == "function" then
+					vim.schedule(function()
+						on_complete_callback(qt.response)
+					end)
+				end
 			end
 		end
 	end
@@ -2958,7 +2966,8 @@ end,
 ---@param template string  # template with model instructions
 ---@param prompt string | nil  # nil for non interactive commads
 ---@param whisper string | nil  # predefined input (e.g. obtained from Whisper)
-M.Prompt = function(params, target, agent, template, prompt, whisper)
+---@param callback function | nil  # callback after completing the prompt
+M.Prompt = function(params, target, agent, template, prompt, whisper, callback)
 	if not agent or not type(agent) == "table" or not agent.provider then
 		M.warning(
 			"The `gp.Prompt` method signature has changed.\n"
@@ -3241,7 +3250,8 @@ M.Prompt = function(params, target, agent, template, prompt, whisper)
 			vim.schedule_wrap(function(qid)
 				on_exit(qid)
 				vim.cmd("doautocmd User GpDone")
-			end)
+			end),
+			on_complete_callback
 		)
 	end