diff --git a/README.md b/README.md index 5ed469b..004966b 100644 --- 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 index 120d757..5794ec2 100644 --- 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 index dbf9e23..20e1703 100644 --- 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