feat: add Dual Boost Mode for Enhanced LLM Responses with Multi-Provi…

…der Support (yetone#854)
tmerse · Nov 17, 2024 · 3051bfd · 3051bfd
1 parent 9891b03
commit 3051bfd
Show file tree

Hide file tree

Showing 3 changed files with 147 additions and 25 deletions.
diff --git a/README.md b/README.md
@@ -208,6 +208,22 @@ _See [config.lua#L9](./lua/avante/config.lua) for the full config_
     temperature = 0,
     max_tokens = 4096,
   },
+  ---Specify the special dual_boost mode
+  ---1. enabled: Whether to enable dual_boost mode. Default to false.
+  ---2. first_provider: The first provider to generate response. Default to "openai".
+  ---3. second_provider: The second provider to generate response. Default to "claude".
+  ---4. prompt: The prompt to generate response based on the two reference outputs.
+  ---5. timeout: Timeout in milliseconds. Default to 60000.
+  ---Whow it works:
+  --- When dual_boost is enabled, avante will generate two responses from the first_provider and second_provider respectively. Then use the response from the first_provider as provider1_output and the response from the second_provider as provider2_output. Finally, avante will generate a response based on the prompt and the two reference outputs, with the default Provider as normal.
+  ---Note: This is an experimental feature and may not work as expected.
+  dual_boost = {
+    enabled = false,
+    first_provider = "openai",
+    second_provider = "claude",
+    prompt = "Based on the two reference outputs below, generate a response that incorporates elements from both but reflects your own judgment and unique perspective. Do not provide any explanation, just give the response directly. Reference Output 1: [{{provider1_output}}], Reference Output 2: [{{provider2_output}}]",
+    timeout = 60000, -- Timeout in milliseconds
+  },
   behaviour = {
     auto_suggestions = false, -- Experimental stage
     auto_set_highlight_group = true,

diff --git a/lua/avante/config.lua b/lua/avante/config.lua
@@ -97,6 +97,22 @@ M.defaults = {
       max_tokens = 8000,
     },
   },
+  ---Specify the special dual_boost mode
+  ---1. enabled: Whether to enable dual_boost mode. Default to false.
+  ---2. first_provider: The first provider to generate response. Default to "openai".
+  ---3. second_provider: The second provider to generate response. Default to "claude".
+  ---4. prompt: The prompt to generate response based on the two reference outputs.
+  ---5. timeout: Timeout in milliseconds. Default to 60000.
+  ---Whow it works:
+  --- When dual_boost is enabled, avante will generate two responses from the first_provider and second_provider respectively. Then use the response from the first_provider as provider1_output and the response from the second_provider as provider2_output. Finally, avante will generate a response based on the prompt and the two reference outputs, with the default Provider as normal.
+  ---Note: This is an experimental feature and may not work as expected.
+  dual_boost = {
+    enabled = false,
+    first_provider = "openai",
+    second_provider = "claude",
+    prompt = "Based on the two reference outputs below, generate a response that incorporates elements from both but reflects your own judgment and unique perspective. Do not provide any explanation, just give the response directly. Reference Output 1: [{{provider1_output}}], Reference Output 2: [{{provider2_output}}]",
+    timeout = 60000, -- Timeout in milliseconds
+  },
   ---Specify the behaviour of avante.nvim
   ---1. auto_apply_diff_after_generation: Whether to automatically apply diff after LLM response.
   ---                                     This would simulate similar behaviour to cursor. Default to false.

diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
@@ -18,32 +18,10 @@ M.CANCEL_PATTERN = "AvanteLLMEscape"
 
 local group = api.nvim_create_augroup("avante_llm", { clear = true })
 
----@alias LlmMode "planning" | "editing" | "suggesting"
----
----@class TemplateOptions
----@field use_xml_format boolean
----@field ask boolean
----@field question string
----@field code_lang string
----@field file_content string
----@field selected_code string | nil
----@field project_context string | nil
----@field history_messages AvanteLLMMessage[]
----
----@class StreamOptions: TemplateOptions
----@field ask boolean
----@field bufnr integer
----@field instructions string
----@field mode LlmMode
----@field provider AvanteProviderFunctor | nil
----@field on_chunk AvanteChunkParser
----@field on_complete AvanteCompleteParser
-
----@param opts StreamOptions
-M.stream = function(opts)
+M._stream = function(opts, Provider)
+  -- print opts
   local mode = opts.mode or "planning"
   ---@type AvanteProviderFunctor
-  local Provider = opts.provider or P[Config.provider]
   local _, body_opts = P.parse_config(Provider)
   local max_tokens = body_opts.max_tokens or 4096
 
@@ -126,7 +104,6 @@ M.stream = function(opts)
     messages = messages,
     image_paths = image_paths,
   }
-
   ---@type string
   local current_event_state = nil
 
@@ -265,6 +242,119 @@ M.stream = function(opts)
   return active_job
 end
 
+local function _merge_response(first_response, second_response, opts, Provider)
+  local prompt = "\n" .. Config.dual_boost.prompt
+  prompt = prompt
+    :gsub("{{[%s]*provider1_output[%s]*}}", first_response)
+    :gsub("{{[%s]*provider2_output[%s]*}}", second_response)
+
+  prompt = prompt .. "\n"
+
+  -- append this reference prompt to the code_opts messages at last
+  opts.instructions = opts.instructions .. prompt
+
+  M._stream(opts, Provider)
+end
+
+local function _collector_process_responses(collector, opts, Provider)
+  if not collector[1] or not collector[2] then
+    Utils.error("One or both responses failed to complete")
+    return
+  end
+  _merge_response(collector[1], collector[2], opts, Provider)
+end
+
+local function _collector_add_response(collector, index, response, opts, Provider)
+  collector[index] = response
+  collector.count = collector.count + 1
+
+  if collector.count == 2 then
+    collector.timer:stop()
+    _collector_process_responses(collector, opts, Provider)
+  end
+end
+
+M._dual_boost_stream = function(opts, Provider, Provider1, Provider2)
+  Utils.debug("Starting Dual Boost Stream")
+
+  local collector = {
+    count = 0,
+    responses = {},
+    timer = uv.new_timer(),
+    timeout_ms = Config.dual_boost.timeout,
+  }
+
+  -- Setup timeout
+  collector.timer:start(
+    collector.timeout_ms,
+    0,
+    vim.schedule_wrap(function()
+      if collector.count < 2 then
+        Utils.warn("Dual boost stream timeout reached")
+        collector.timer:stop()
+        -- Process whatever responses we have
+        _collector_process_responses(collector, opts, Provider)
+      end
+    end)
+  )
+
+  -- Create options for both streams
+  local function create_stream_opts(index)
+    local response = ""
+    return vim.tbl_extend("force", opts, {
+      on_chunk = function(chunk)
+        if chunk then response = response .. chunk end
+      end,
+      on_complete = function(err)
+        if err then
+          Utils.error(string.format("Stream %d failed: %s", index, err))
+          return
+        end
+        Utils.debug(string.format("Response %d completed", index))
+        _collector_add_response(collector, index, response, opts, Provider)
+      end,
+    })
+  end
+
+  -- Start both streams
+  local success, err = xpcall(function()
+    M._stream(create_stream_opts(1), Provider1)
+    M._stream(create_stream_opts(2), Provider2)
+  end, function(err) return err end)
+  if not success then Utils.error("Failed to start dual_boost streams: " .. tostring(err)) end
+end
+
+---@alias LlmMode "planning" | "editing" | "suggesting"
+---
+---@class TemplateOptions
+---@field use_xml_format boolean
+---@field ask boolean
+---@field question string
+---@field code_lang string
+---@field file_content string
+---@field selected_code string | nil
+---@field project_context string | nil
+---@field history_messages AvanteLLMMessage[]
+---
+---@class StreamOptions: TemplateOptions
+---@field ask boolean
+---@field bufnr integer
+---@field instructions string
+---@field mode LlmMode
+---@field provider AvanteProviderFunctor | nil
+---@field on_chunk AvanteChunkParser
+---@field on_complete AvanteCompleteParser
+
+---@param opts StreamOptions
+M.stream = function(opts)
+  local Provider = opts.provider or P[Config.provider]
+  if Config.dual_boost.enabled then
+    M._dual_boost_stream(opts, Provider, P[Config.dual_boost.first_provider], P[Config.dual_boost.second_provider])
+  else
+    M._stream(opts, Provider)
+  end
+end
+
 function M.cancel_inflight_request() api.nvim_exec_autocmds("User", { pattern = M.CANCEL_PATTERN }) end
 
 return M