feat(clipboard): initial support (yetone#279)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
SsSniperSkinZo · Aug 27, 2024 · cf68572 · cf68572
1 parent 77551ce
commit cf68572
Show file tree

Hide file tree

Showing 10 changed files with 244 additions and 10 deletions.
diff --git a/lua/avante/clipboard/darwin.lua b/lua/avante/clipboard/darwin.lua
@@ -0,0 +1,59 @@
+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.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.get_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
diff --git a/lua/avante/clipboard/init.lua b/lua/avante/clipboard/init.lua
@@ -0,0 +1,27 @@
+---NOTE: this module is inspired by https://github.com/HakonHarnes/img-clip.nvim/tree/main
+
+local Utils = require("avante.utils")
+
+---@class AvanteClipboard
+---@field clip_cmd string
+---@field get_clip_cmd fun(): string
+---@field has_content fun(): boolean
+---@field get_content fun(): string
+---
+---@class avante.Clipboard: AvanteClipboard
+local M = {}
+
+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]
+    elseif t[k] ~= nil then
+      return t[k]
+    else
+      error("Failed to find clipboard implementation for " .. os_mapping)
+    end
+  end,
+})
diff --git a/lua/avante/clipboard/linux.lua b/lua/avante/clipboard/linux.lua
@@ -0,0 +1,59 @@
+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.get_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
diff --git a/lua/avante/clipboard/windows.lua b/lua/avante/clipboard/windows.lua
@@ -0,0 +1,51 @@
+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.get_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
diff --git a/lua/avante/config.lua b/lua/avante/config.lua
@@ -75,9 +75,11 @@ 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.
   behaviour = {
     auto_set_highlight_group = true,
     auto_apply_diff_after_generation = false,
+    support_paste_from_clipboard = false,
   },
   history = {
     storage_path = vim.fn.stdpath("state") .. "/avante",

diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
@@ -27,7 +27,8 @@ Your primary task is to suggest code modifications with precise line number rang
 2. When suggesting modifications:
    a. Use the language in the question to reply. If there are non-English parts in the question, use the language of those parts.
    b. Explain why the change is necessary or beneficial.
-   c. Provide the exact code snippet to be replaced using this format:
+   c. If an image is provided, make sure to use the image in conjunction with the code snippet.
+   d. Provide the exact code snippet to be replaced using this format:
 
 Replace lines: {{start_line}}-{{end_line}}
 ```{{language}}
@@ -96,6 +97,8 @@ M.stream = function(question, code_lang, code_content, selected_content_content,
   ---@type AvanteCurlOutput
   local spec = Provider.parse_curl_args(Provider, code_opts)
 
+  Utils.debug(spec)
+
   ---@param line string
   local function parse_stream_data(line)
     local event = line:match("^event: (.+)$")

diff --git a/lua/avante/providers/claude.lua b/lua/avante/providers/claude.lua
@@ -1,5 +1,6 @@
+local Config = require("avante.config")
 local Utils = require("avante.utils")
-local Tokens = require("avante.utils.tokens")
+local Clipboard = require("avante.clipboard")
 local P = require("avante.providers")
 
 ---@class AvanteProviderFunctor
@@ -13,7 +14,7 @@ M.parse_message = function(opts)
     text = string.format("<code>```%s\n%s```</code>", opts.code_lang, opts.code_content),
   }
 
-  if Tokens.calculate_tokens(code_prompt_obj.text) > 1024 then
+  if Utils.tokens.calculate_tokens(code_prompt_obj.text) > 1024 then
     code_prompt_obj.cache_control = { type = "ephemeral" }
   end
 
@@ -31,7 +32,7 @@ M.parse_message = function(opts)
       text = string.format("<code>```%s\n%s```</code>", opts.code_lang, opts.selected_code_content),
     }
 
-    if Tokens.calculate_tokens(selected_code_obj.text) > 1024 then
+    if Utils.tokens.calculate_tokens(selected_code_obj.text) > 1024 then
       selected_code_obj.cache_control = { type = "ephemeral" }
     end
 
@@ -43,14 +44,25 @@ M.parse_message = function(opts)
     text = string.format("<question>%s</question>", opts.question),
   })
 
+  if Config.behaviour.support_paste_from_clipboard and Clipboard.has_content() then
+    table.insert(message_content, {
+      type = "image",
+      source = {
+        type = "base64",
+        media_type = "image/png",
+        data = Clipboard.get_content(),
+      },
+    })
+  end
+
   local user_prompt = opts.base_prompt
 
   local user_prompt_obj = {
     type = "text",
     text = user_prompt,
   }
 
-  if Tokens.calculate_tokens(user_prompt_obj.text) > 1024 then
+  if Utils.tokens.calculate_tokens(user_prompt_obj.text) > 1024 then
     user_prompt_obj.cache_control = { type = "ephemeral" }
   end
 

diff --git a/lua/avante/providers/init.lua b/lua/avante/providers/init.lua
@@ -63,7 +63,7 @@ local Dressing = require("avante.ui.dressing")
 ---@field parse_response_data AvanteResponseParser
 ---@field parse_curl_args? AvanteCurlArgsParser
 ---@field parse_stream_data? AvanteStreamParser
----@field parse_api_key fun(): string | nil
+---@field parse_api_key? fun(): string | nil
 ---
 ---@class AvanteProviderFunctor
 ---@field parse_message AvanteMessageParser

diff --git a/lua/avante/utils.lua → lua/avante/utils/init.lua b/lua/avante/utils.lua → lua/avante/utils/init.lua
@@ -1,6 +1,7 @@
 local api = vim.api
 
 ---@class avante.utils: LazyUtilCore
+---@field tokens avante.utils.tokens
 local M = {}
 
 setmetatable(M, {
@@ -43,11 +44,29 @@ end
 
 --- This function will run given shell command synchronously.
 ---@param input_cmd string
----@return integer, string?, string?
+---@return vim.SystemCompleted
 M.shell_run = function(input_cmd)
-  local output =
-    vim.system(vim.split("sh -c " .. vim.fn.shellescape(input_cmd), " ", { trimempty = true }), { text = true }):wait()
-  return output.code, output.stderr, output.stdout
+  local shell = vim.o.shell:lower()
+  ---@type string
+  local cmd
+
+  -- powershell then we can just run the cmd
+  if shell:match("powershell") or shell:match("pwsh") then
+    cmd = input_cmd
+  elseif vim.fn.has("wsl") > 0 then
+    -- wsl: powershell.exe -Command 'command "/path"'
+    cmd = "powershell.exe -NoProfile -Command '" .. input_cmd:gsub("'", '"') .. "'"
+  elseif vim.fn.has("win32") > 0 then
+    cmd = 'powershell.exe -NoProfile -Command "' .. input_cmd:gsub('"', "'") .. '"'
+  else
+    -- linux and macos we wil just do sh -c
+    cmd = "sh -c " .. vim.fn.shellescape(input_cmd)
+  end
+
+  local output = vim.fn.system(cmd)
+  local code = vim.v.shell_error
+
+  return { stdout = output, code = code }
 end
 
 ---@alias _ToggleSet fun(state: boolean): nil

diff --git a/lua/avante/utils/tokens.lua b/lua/avante/utils/tokens.lua
@@ -1,5 +1,7 @@
 --Taken from https://github.com/jackMort/ChatGPT.nvim/blob/main/lua/chatgpt/flows/chat/tokens.lua
 local Tiktoken = require("avante.tiktoken")
+
+---@class avante.utils.tokens
 local Tokens = {}
 
 ---@type table<[string], number>