feat(templates): avanterules filetype support (closes #254) (#466)

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>

lua/avante/path.lua

@@ -1,5 +1,7 @@
 local fn, api = vim.fn, vim.api
+local Utils = require("avante.utils")
 local Path = require("plenary.path")
+local Scan = require("plenary.scandir")
 local Config = require("avante.config")
 
 ---@class avante.Path
@@ -7,9 +9,10 @@ local Config = require("avante.config")
 ---@field cache_path Path
 local P = {}
 
-local M = {}
-
+-- Helpers
 local H = {}
+
+-- Get a chat history file name given a buffer
 ---@param bufnr integer
 ---@return string
 H.filename = function(bufnr)
@@ -20,12 +23,23 @@ H.filename = function(bufnr)
   return fn.substitute(path_with_separators, "[^A-Za-z0-9._]", "_", "g") .. ".json"
 end
 
+-- Given a mode, return the file name for the custom prompt.
+---@param mode LlmMode
+H.get_mode_file = function(mode)
+  return string.format("custom.%s.avanterules", mode)
+end
+
+-- History path
+local M = {}
+
+-- Returns the Path to the chat history file for the given buffer.
 ---@param bufnr integer
 ---@return Path
 M.get = function(bufnr)
   return Path:new(Config.history.storage_path):joinpath(H.filename(bufnr))
 end
 
+-- Loads the chat history for the given buffer.
 ---@param bufnr integer
 M.load = function(bufnr)
   local history_file = M.get(bufnr)
@@ -36,6 +50,7 @@ M.load = function(bufnr)
   return {}
 end
 
+-- Saves the chat history for the given buffer.
 ---@param bufnr integer
 ---@param history table
 M.save = function(bufnr, history)
@@ -45,6 +60,83 @@ end
 
 P.history = M
 
+-- Prompt path
+local N = {}
+
+---@class AvanteTemplates
+---@field initialize fun(directory: string): nil
+---@field render fun(template: string, context: TemplateOptions): string
+local templates = nil
+
+N.templates = { planning = nil, editing = nil, suggesting = nil }
+
+-- Creates a directory in the cache path for the given buffer and copies the custom prompts to it.
+-- We need to do this beacuse the prompt template engine requires a given directory to load all required files.
+-- PERF: Hmm instead of copy to cache, we can also load in globals context, but it requires some work on bindings. (eh maybe?)
+---@param bufnr number
+---@return string the resulted cache_directory to be loaded with avante_templates
+N.get = function(bufnr)
+  if not P.available() then
+    error("Make sure to build avante (missing avante_templates)", 2)
+  end
+
+  -- get root directory of given bufnr
+  local directory = Path:new(Utils.root.get({ buf = bufnr }))
+  ---@cast directory Path
+  ---@type Path
+  local cache_prompt_dir = P.cache_path:joinpath(directory)
+  if not cache_prompt_dir:exists() then
+    cache_prompt_dir:mkdir({ parents = true })
+  end
+
+  local scanner = Scan.scan_dir(directory:absolute(), { depth = 1, add_dirs = true })
+  for _, entry in ipairs(scanner) do
+    local file = Path:new(entry)
+    if entry:find("planning") and N.templates.planning == nil then
+      N.templates.planning = file:read()
+    elseif entry:find("editing") and N.templates.editing == nil then
+      N.templates.editing = file:read()
+    elseif entry:find("suggesting") and N.templates.suggesting == nil then
+      N.templates.suggesting = file:read()
+    end
+  end
+
+  Path:new(debug.getinfo(1).source:match("@?(.*/)"):gsub("/lua/avante/path.lua$", "") .. "templates")
+    :copy({ destination = cache_prompt_dir, recursive = true })
+
+  vim
+    .iter(N.templates)
+    :filter(function(_, v)
+      return v ~= nil
+    end)
+    :each(function(k, v)
+      local f = cache_prompt_dir:joinpath(H.get_mode_file(k))
+      f:write(v, "w")
+    end)
+
+  return cache_prompt_dir:absolute()
+end
+
+---@param mode LlmMode
+N.get_file = function(mode)
+  if N.templates[mode] ~= nil then
+    return H.get_mode_file(mode)
+  end
+  return string.format("%s.avanterules", mode)
+end
+
+---@param mode LlmMode
+---@param opts TemplateOptions
+N.render = function(mode, opts)
+  return templates.render(N.get_file(mode), opts)
+end
+
+N.initialize = function(directory)
+  templates.initialize(directory)
+end
+
+P.prompts = N
+
 P.setup = function()
   local history_path = Path:new(Config.history.storage_path)
   if not history_path:exists() then
@@ -57,6 +149,26 @@ P.setup = function()
     cache_path:mkdir({ parents = true })
   end
   P.cache_path = cache_path
+
+  vim.defer_fn(function()
+    local ok, module = pcall(require, "avante_templates")
+    ---@cast module AvanteTemplates
+    ---@cast ok boolean
+    if not ok then
+      return
+    end
+    if templates == nil then
+      templates = module
+    end
+  end, 1000)
+end
+
+P.available = function()
+  return templates ~= nil
+end
+
+P.clear = function()
+  P.cache_path:rm({ recursive = true })
 end
 
 return P