feat: pasting image within buffer (#331)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
This commit is contained in:
Aaron Pham
2024-08-28 14:43:14 -04:00
committed by GitHub
parent 46a621e9de
commit c635f73748
12 changed files with 166 additions and 295 deletions

74
lua/avante/clipboard.lua Normal file
View File

@@ -0,0 +1,74 @@
---NOTE: this module is inspired by https://github.com/HakonHarnes/img-clip.nvim/tree/main
---@see https://github.com/ekickx/clipboard-image.nvim/blob/main/lua/clipboard-image/paste.lua
local Path = require("plenary.path")
local Utils = require("avante.utils")
local Config = require("avante.config")
---@module "img-clip"
local ImgClip = nil
---@class AvanteClipboard
---@field get_base64_content fun(filepath: string): string | nil
---
---@class avante.Clipboard: AvanteClipboard
local M = {}
---@type Path
local paste_directory = nil
---@return Path
local function get_paste_directory()
if paste_directory then
return paste_directory
end
paste_directory = Path:new(Config.history.storage_path):joinpath("pasted_images")
return paste_directory
end
M.support_paste_image = Config.support_paste_image
M.setup = function()
get_paste_directory()
if not paste_directory:exists() then
paste_directory:mkdir({ parent = true })
end
if M.support_paste_image() and ImgClip == nil then
ImgClip = require("img-clip")
end
end
---@param line string
M.paste_image = function(line)
if not Config.support_paste_image() then
return false
end
return ImgClip.paste_image({
dir_path = paste_directory:absolute(),
prompt_for_file_name = false,
filetypes = {
AvanteInput = { url_encode_path = true, template = "\nimage: $FILE_PATH\n" },
},
}, line)
end
---@param filepath string
M.get_base64_content = function(filepath)
local os_mapping = Utils.get_os_name()
---@type vim.SystemCompleted
local output
if os_mapping == "darwin" or os_mapping == "linux" then
output = Utils.shell_run(("cat %s | base64 | tr -d '\n'"):format(filepath))
if output.code == 0 then
return output.stdout
else
error("Failed to convert image to base64")
end
else
Utils.warn("Windows is not supported yet", { title = "Avante" })
end
end
return M

View File

@@ -1,85 +0,0 @@
local Utils = require("avante.utils")
---@class AvanteClipboard
local M = {}
---@alias DarwinClipboardCommand "pngpaste" | "osascript"
M.clip_cmd = nil
---@return DarwinClipboardCommand
M.get_clip_cmd = function()
if M.clip_cmd then
return M.clip_cmd
end
if vim.fn.executable("pngpaste") == 1 then
M.clip_cmd = "pngpaste"
elseif vim.fn.executable("osascript") == 1 then
M.clip_cmd = "osascript"
end
return M.clip_cmd
end
M.has_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "pngpaste" then
output = Utils.shell_run("pngpaste -")
return output.code == 0
elseif cmd == "osascript" then
output = Utils.shell_run("osascript -e 'clipboard info'")
return output.code == 0 and output.stdout ~= nil and output.stdout:find("class PNGf") ~= nil
end
Utils.warn("Failed to validate clipboard content", { title = "Avante" })
return false
end
M.save_content = function(filepath)
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "pngpaste" then
output = Utils.shell_run(('pngpaste - > "%s"'):format(filepath))
return output.code == 0
elseif cmd == "osascript" then
output = Utils.shell_run(
string.format(
[[osascript -e 'set theFile to (open for access POSIX file "%s" with write permission)' ]]
.. [[-e 'try' -e 'write (the clipboard as «class PNGf») to theFile' -e 'end try' ]]
.. [[-e 'close access theFile' -e 'do shell script "cat %s > %s"']],
filepath,
filepath,
filepath
)
)
return output.code == 0
end
return false
end
M.get_base64_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "pngpaste" then
output = Utils.shell_run("pngpaste - | base64 | tr -d '\n'")
if output.code == 0 then
return output.stdout
end
elseif cmd == "osascript" then
output = Utils.shell_run(
[[osascript -e 'set theFile to (open for access POSIX file "/tmp/image.png" with write permission)' -e 'try' -e 'write (the clipboard as «class PNGf») to theFile' -e 'end try' -e 'close access theFile'; ]]
.. [[cat /tmp/image.png | base64 | tr -d '\n']]
)
if output.code == 0 then
return output.stdout
end
end
error("Failed to get clipboard content")
end
return M

View File

@@ -1,48 +0,0 @@
---NOTE: this module is inspired by https://github.com/HakonHarnes/img-clip.nvim/tree/main
---@see https://github.com/ekickx/clipboard-image.nvim/blob/main/lua/clipboard-image/paste.lua
local Path = require("plenary.path")
local Utils = require("avante.utils")
local Config = require("avante.config")
---@class AvanteClipboard
---@field clip_cmd string
---@field get_clip_cmd fun(): string
---@field has_content fun(): boolean
---@field get_base64_content fun(): string
---@field save_content fun(filename: string): boolean
---
---@class avante.Clipboard: AvanteClipboard
local M = {}
---@type Path
local paste_directory = nil
---@return Path
local function get_paste_directory()
if paste_directory then
return paste_directory
end
paste_directory = Path:new(Config.history.storage_path):joinpath("pasted_images")
return paste_directory
end
M.setup = function()
get_paste_directory()
if not paste_directory:exists() then
paste_directory:mkdir({ parent = true })
end
end
return setmetatable(M, {
__index = function(t, k)
local os_mapping = Utils.get_os_name()
---@type AvanteClipboard
local impl = require("avante.clipboard." .. os_mapping)
if impl[k] ~= nil then
return impl[k]
end
return t[k]
end,
})

View File

@@ -1,74 +0,0 @@
local Utils = require("avante.utils")
---@class AvanteClipboard
local M = {}
M.clip_cmd = nil
M.get_clip_cmd = function()
if M.clip_cmd then
return M.clip_cmd
end
-- Wayland
if os.getenv("WAYLAND_DISPLAY") ~= nil and vim.fn.executable("wl-paste") == 1 then
M.clip_cmd = "wl-paste"
-- X11
elseif os.getenv("DISPLAY") ~= nil and vim.fn.executable("xclip") == 1 then
M.clip_cmd = "xclip"
end
return M.clip_cmd
end
M.has_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
-- X11
if cmd == "xclip" then
output = Utils.shell_run("xclip -selection clipboard -t TARGETS -o")
return output.code == 0 and output.stdout:find("image/png") ~= nil
elseif cmd == "wl-paste" then
output = Utils.shell_run("wl-paste --list-types")
return output.code == 0 and output.stdout:find("image/png") ~= nil
end
Utils.warn("Failed to validate clipboard content", { title = "Avante" })
return false
end
M.save_content = function(filepath)
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "xclip" then
output = Utils.shell_run(('xclip -selection clipboard -o -t image/png > "%s"'):format(filepath))
return output.code == 0
elseif cmd == "wl-paste" then
output = Utils.shell_run(('wl-paste --type image/png > "%s"'):format(filepath))
return output.code == 0
end
return false
end
M.get_base64_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "xclip" then
output = Utils.shell_run("xclip -selection clipboard -o -t image/png | base64 | tr -d '\n'")
if output.code == 0 then
return output.stdout
end
elseif cmd == "osascript" then
output = Utils.shell_run("wl-paste --type image/png | base64 | tr -d '\n'")
if output.code == 0 then
return output.stdout
end
end
error("Failed to get clipboard content")
end
return M

View File

@@ -1,67 +0,0 @@
local Utils = require("avante.utils")
---@class AvanteClipboard
local M = {}
M.clip_cmd = nil
M.get_clip_cmd = function()
if M.clip_cmd then
return M.clip_cmd
end
if (vim.fn.has("win32") > 0 or vim.fn.has("wsl") > 0) and vim.fn.executable("powershell.exe") then
M.clip_cmd = "powershell.exe"
end
return M.clip_cmd
end
M.has_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "powershell.exe" then
output =
Utils.shell_run("Add-Type -AssemblyName System.Windows.Forms; [System.Windows.Forms.Clipboard]::GetImage()")
return output.code == 0 and output.stdout:find("Width") ~= nil
end
Utils.warn("Failed to validate clipboard content", { title = "Avante" })
return false
end
M.save_content = function(filepath)
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "powershell.exe" then
output = Utils.shell_run(
("Add-Type -AssemblyName System.Windows.Forms; [System.Windows.Forms.Clipboard]::GetImage().Save('%s')"):format(
filepath
)
)
return output.code == 0
end
return false
end
M.get_base64_content = function()
local cmd = M.get_clip_cmd()
---@type vim.SystemCompleted
local output
if cmd == "powershell.exe" then
output = Utils.shell_run(
[[Add-Type -AssemblyName System.Windows.Forms; $ms = New-Object System.IO.MemoryStream;]]
.. [[ [System.Windows.Forms.Clipboard]::GetImage().Save($ms, [System.Drawing.Imaging.ImageFormat]::Png);]]
.. [[ [System.Convert]::ToBase64String($ms.ToArray())]]
)
if output.code == 0 then
return output.stdout:gsub("\r\n", ""):gsub("\n", ""):gsub("\r", "")
end
end
error("Failed to get clipboard content")
end
return M

View File

@@ -1,6 +1,8 @@
---NOTE: user will be merged with defaults and
---we add a default var_accessor for this table to config values.
---
local Utils = require("avante.utils")
---@class avante.CoreConfig: avante.Config
local M = {}
@@ -65,7 +67,7 @@ M.defaults = {
---1. auto_apply_diff_after_generation: Whether to automatically apply diff after LLM response.
--- This would simulate similar behaviour to cursor. Default to false.
---2. auto_set_highlight_group: Whether to automatically set the highlight group for the current line. Default to true.
---3. support_paste_from_clipboard: Whether to support pasting image from clipboard. Note that we will override vim.paste for this. Default to false.
---3. support_paste_from_clipboard: Whether to support pasting image from clipboard. This will be determined automatically based whether img-clip is available or not.
behaviour = {
auto_set_highlight_group = true,
auto_apply_diff_after_generation = false,
@@ -151,7 +153,17 @@ M.hints = {}
---@param opts? avante.Config
function M.setup(opts)
M.options = vim.tbl_deep_extend("force", M.defaults, opts or {})
M.options = vim.tbl_deep_extend(
"force",
M.defaults,
opts or {},
---@type avante.Config
{
behaviour = {
support_paste_from_clipboard = M.support_paste_image(),
},
}
)
M.diff = vim.tbl_deep_extend(
"force",
@@ -168,6 +180,14 @@ function M.setup(opts)
end
end
M.support_paste_image = function()
local supported = Utils.has("img-clip.nvim")
if not supported then
Utils.warn("img-clip.nvim is not installed. Pasting image will be disabled.", { once = true })
end
return supported
end
---@param opts? avante.Config
function M.override(opts)
opts = opts or {}

View File

@@ -96,11 +96,27 @@ M.stream = function(question, code_lang, code_content, selected_content_content,
mode = mode or "planning"
local provider = Config.provider
-- Check if the question contains an image path
local image_path = nil
local original_question = question
if question:match("image: ") then
local lines = vim.split(question, "\n")
for i, line in ipairs(lines) do
if line:match("^image: ") then
image_path = line:gsub("^image: ", "")
table.remove(lines, i)
original_question = table.concat(lines, "\n")
break
end
end
end
---@type AvantePromptOptions
local code_opts = {
base_prompt = mode == "planning" and planning_mode_prompt or editing_mode_prompt,
system_prompt = system_prompt,
question = question,
question = original_question,
image_path = image_path,
code_lang = code_lang,
code_content = code_content,
selected_code_content = selected_content_content,
@@ -188,13 +204,6 @@ M.stream = function(question, code_lang, code_content, selected_content_content,
on_complete("API request failed with status " .. result.status .. ". Body: " .. vim.inspect(result.body))
end
end)
else
vim.schedule(function()
if not completed then
completed = true
on_complete(nil)
end
end)
end
active_job = nil
end,

View File

@@ -39,22 +39,22 @@ M.parse_message = function(opts)
table.insert(message_content, selected_code_obj)
end
table.insert(message_content, {
type = "text",
text = string.format("<question>%s</question>", opts.question),
})
if Config.behaviour.support_paste_from_clipboard and Clipboard.has_content() then
if Clipboard.support_paste_image() and opts.image_path then
table.insert(message_content, {
type = "image",
source = {
type = "base64",
media_type = "image/png",
data = Clipboard.get_base64_content(),
data = Clipboard.get_base64_content(opts.image_path),
},
})
end
table.insert(message_content, {
type = "text",
text = string.format("<question>%s</question>", opts.question),
})
local user_prompt = opts.base_prompt
local user_prompt_obj = {

View File

@@ -12,6 +12,7 @@ local Dressing = require("avante.ui.dressing")
---@field base_prompt AvanteBasePrompt
---@field system_prompt AvanteSystemPrompt
---@field question string
---@field image_path? string
---@field code_lang string
---@field code_content string
---@field selected_code_content? string

View File

@@ -63,12 +63,12 @@ end
M.parse_message = function(opts)
---@type string | OpenAIMessage[]
local user_content
if Config.behaviour.support_paste_from_clipboard and Clipboard.has_content() then
if Config.behaviour.support_paste_from_clipboard and opts.image_path then
user_content = {}
table.insert(user_content, {
type = "image_url",
image_url = {
url = "data:image/png;base64," .. Clipboard.get_base64_content(),
url = "data:image/png;base64," .. Clipboard.get_base64_content(opts.image_path),
},
})
table.insert(user_content, { type = "text", text = M.get_user_message(opts) })