feat: support cmd: loading api keys for google search APIs (#1532)

refactor environment loading logic into Utils and use it for existing
provider loading, and web search where appropriate.
This commit is contained in:
adasauce
2025-03-09 04:01:49 -03:00
committed by GitHub
parent 868c136574
commit 02fb470140
4 changed files with 70 additions and 53 deletions

View File

@@ -298,7 +298,7 @@ function M.web_search(opts, on_log)
local search_engine = Config.web_search_engine.providers[provider_type]
if search_engine == nil then return nil, "No search engine found: " .. provider_type end
if search_engine.api_key_name == "" then return nil, "No API key provided" end
local api_key = os.getenv(search_engine.api_key_name)
local api_key = Utils.environment.parse(search_engine.api_key_name)
if api_key == nil or api_key == "" then
return nil, "Environment variable " .. search_engine.api_key_name .. " is not set"
end
@@ -350,7 +350,7 @@ function M.web_search(opts, on_log)
local jsn = vim.json.decode(resp.body)
return search_engine.format_response_body(jsn)
elseif provider_type == "google" then
local engine_id = os.getenv(search_engine.engine_id_name)
local engine_id = Utils.environment.parse(search_engine.engine_id_name)
if engine_id == nil or engine_id == "" then
return nil, "Environment variable " .. search_engine.engine_id_name .. " is not set"
end

View File

@@ -30,61 +30,13 @@ E.cache = {}
---@param Opts AvanteSupportedProvider | AvanteProviderFunctor | AvanteBedrockProviderFunctor
---@return string | nil
function E.parse_envvar(Opts)
local api_key_name = Opts.api_key_name
if api_key_name == nil then error("Requires api_key_name") end
local cache_key = type(api_key_name) == "table" and table.concat(api_key_name, "__") or api_key_name
if E.cache[cache_key] ~= nil then return E.cache[cache_key] end
local cmd = type(api_key_name) == "table" and api_key_name or api_key_name:match("^cmd:(.*)")
local value = nil
if cmd ~= nil then
-- NOTE: in case api_key_name is cmd, and users still set envvar
-- We will try to get envvar first
if Opts._shellenv ~= nil and Opts._shellenv ~= "" then
value = os.getenv(Opts._shellenv)
if value ~= nil then
E.cache[cache_key] = value
vim.g.avante_login = true
return value
end
end
if type(cmd) == "string" then cmd = vim.split(cmd, " ", { trimempty = true }) end
Utils.debug("running command:", cmd)
local exit_codes = { 0 }
local ok, job_or_err = pcall(vim.system, cmd, { text = true }, function(result)
Utils.debug("command result:", result)
local code = result.code
local stderr = result.stderr or ""
local stdout = result.stdout and vim.split(result.stdout, "\n") or {}
if vim.tbl_contains(exit_codes, code) then
value = stdout[1]
E.cache[cache_key] = value
vim.g.avante_login = true
else
Utils.error("Failed to get API key: (error code" .. code .. ")\n" .. stderr, { once = true, title = "Avante" })
end
end)
if not ok then
error("failed to run command: " .. cmd .. "\n" .. job_or_err)
return
end
else
value = os.getenv(api_key_name)
end
local value = Utils.environment.parse(Opts.api_key_name, Opts._shellenv)
if value ~= nil then
E.cache[cache_key] = value
vim.g.avante_login = true
return value
end
return value
return nil
end
--- initialize the environment variable for current neovim session.

View File

@@ -0,0 +1,64 @@
local Utils = require("avante.utils")
---@class avante.utils.environment
local M = {}
---@private
---@type table<string, string>
M.cache = {}
---Parse environment variable using optional cmd: feature with an override fallback
---@param key_name string
---@param override? string
---@return string | nil
function M.parse(key_name, override)
if key_name == nil then error("Requires key_name") end
local cache_key = type(key_name) == "table" and table.concat(key_name, "__") or key_name
if M.cache[cache_key] ~= nil then return M.cache[cache_key] end
local cmd = type(key_name) == "table" and key_name or key_name:match("^cmd:(.*)")
local value = nil
if cmd ~= nil then
if override ~= nil and override ~= "" then
value = os.getenv(override)
if value ~= nil then
M.cache[cache_key] = value
return value
end
end
if type(cmd) == "string" then cmd = vim.split(cmd, " ", { trimempty = true }) end
Utils.debug("running command:", cmd)
local exit_codes = { 0 }
local ok, job_or_err = pcall(vim.system, cmd, { text = true }, function(result)
Utils.debug("command result:", result)
local code = result.code
local stderr = result.stderr or ""
local stdout = result.stdout and vim.split(result.stdout, "\n") or {}
if vim.tbl_contains(exit_codes, code) then
value = stdout[1]
M.cache[cache_key] = value
else
Utils.error("failed to get key: (error code" .. code .. ")\n" .. stderr, { once = true, title = "Avante" })
end
end)
if not ok then
Utils.error("failed to run command: " .. cmd .. "\n" .. job_or_err)
return
end
else
value = os.getenv(key_name)
end
if value ~= nil then M.cache[cache_key] = value end
return value
end
return M

View File

@@ -9,6 +9,7 @@ local lsp = vim.lsp
---@field root avante.utils.root
---@field file avante.utils.file
---@field history avante.utils.history
---@field environment avante.utils.environment
local M = {}
setmetatable(M, {