fix: retry when rate limited (#1417)

This commit is contained in:
yetone
2025-02-27 12:52:58 +08:00
committed by GitHub
parent e69e7eceeb
commit 5abe579019
3 changed files with 25 additions and 1 deletions

View File

@@ -198,6 +198,13 @@ function M._stream(opts)
end
return handle_next_tool_use(sorted_tool_use_list, 1, old_tool_histories)
end
if stop_opts.reason == "rate_limit" then
local msg = "Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds ..."
opts.on_chunk("\n*[" .. msg .. "]*\n")
Utils.info("Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds", { title = "Avante" })
vim.defer_fn(function() M._stream(opts) end, stop_opts.retry_after * 1000)
return
end
return opts.on_stop(stop_opts)
end,
}
@@ -301,6 +308,21 @@ function M._stream(opts)
else
Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
end
if result.status == 429 then
local headers_map = vim.iter(result.headers):fold({}, function(acc, value)
local pieces = vim.split(value, ":")
local key = pieces[1]
local remain = vim.list_slice(pieces, 2)
if not remain then return acc end
local val = Utils.trim_spaces(table.concat(remain, ":"))
acc[key] = val
return acc
end)
local retry_after = 10
if headers_map["retry-after"] then retry_after = tonumber(headers_map["retry-after"]) or 10 end
handler_opts.on_stop({ reason = "rate_limit", retry_after = retry_after })
return
end
vim.schedule(function()
if not completed then
completed = true

View File

@@ -257,6 +257,7 @@ function M.parse_curl_args(provider, prompt_opts)
end
function M.on_error(result)
if result.status == 429 then return end
if not result.body then
return Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
end

View File

@@ -231,10 +231,11 @@ vim.g.avante_login = vim.g.avante_login
---@field usage? AvanteLLMUsage
---
---@class AvanteLLMStopCallbackOptions
---@field reason "complete" | "tool_use" | "error"
---@field reason "complete" | "tool_use" | "error" | "rate_limit"
---@field error? string | table
---@field usage? AvanteLLMUsage
---@field tool_use_list? AvanteLLMToolUse[]
---@field retry_after? integer
---
---@alias AvanteStreamParser fun(line: string, handler_opts: AvanteHandlerOptions): nil
---@alias AvanteLLMStartCallback fun(opts: AvanteLLMStartCallbackOptions): nil