feat(web_search_engine): adds google as a provider (#1226)

This commit is contained in:
Amin Roosta 2025-02-10 00:21:55 -05:00 committed by GitHub
parent 59fdbd72a4
commit c07d396e30
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 52 additions and 4 deletions

View File

@ -550,15 +550,16 @@ For more information, see [Custom Providers](https://github.com/yetone/avante.nv
## Web Search Engines
Avante's tools include some web search engines, currently support [tavily](https://tavily.com/) and [serpapi](https://serpapi.com/). The default is tavily, and can be changed through configuring `Config.web_search_engine.provider`:
Avante's tools include some web search engines, currently support [tavily](https://tavily.com/), [serpapi](https://serpapi.com/) and google's [programmable search engine](https://developers.google.com/custom-search/v1/overview). The default is tavily, and can be changed through configuring `Config.web_search_engine.provider`:
```lua
web_search_engine = {
provider = "tavily", -- tavily or serpapi
provider = "tavily", -- tavily, serpapi or google
}
```
You need to set the environment variable `TAVILY_API_KEY` or `SERPAPI_API_KEY` to use tavily or serpapi.
You need to set the environment variable `TAVILY_API_KEY` , `SERPAPI_API_KEY` to use tavily or serpapi.
To use google, set the `GOOGLE_SEARCH_API_KEY` as the [API key](https://developers.google.com/custom-search/v1/overview), and `GOOGLE_SEARCH_ENGINE_ID` as the [search engine](https://programmablesearchengine.google.com) ID.
## Disable Tools

View File

@ -55,8 +55,33 @@ M._defaults = {
}
end
)
:take(5)
:totable()
return vim.json.encode(jsn), nil
end
return "", nil
end,
},
google = {
api_key_name = "GOOGLE_SEARCH_API_KEY",
engine_id_name = "GOOGLE_SEARCH_ENGINE_ID",
extra_request_body = {},
---@type WebSearchEngineProviderResponseBodyFormatter
format_response_body = function(body)
if body.items ~= nil then
local jsn = vim
.iter(body.items)
:map(
function(result)
return {
title = result.title,
link = result.link,
snippet = result.snippet,
}
end
)
:take(5)
:totable()
if #jsn > 5 then jsn = vim.list_slice(jsn, 1, 5) end
return vim.json.encode(jsn), nil
end
return "", nil

View File

@ -327,6 +327,28 @@ function M.web_search(opts, on_log)
if resp.status ~= 200 then return nil, "Error: " .. resp.body end
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)
if engine_id == nil or engine_id == "" then
return nil, "Environment variable " .. search_engine.engine_id_namee .. " is not set"
end
local query_params = vim.tbl_deep_extend("force", {
key = api_key,
cx = engine_id,
q = opts.query,
}, search_engine.extra_request_body)
local query_string = ""
for key, value in pairs(query_params) do
query_string = query_string .. key .. "=" .. vim.uri_encode(value) .. "&"
end
local resp = curl.get("https://www.googleapis.com/customsearch/v1?" .. query_string, {
headers = {
["Content-Type"] = "application/json",
},
})
if resp.status ~= 200 then return nil, "Error: " .. resp.body end
local jsn = vim.json.decode(resp.body)
return search_engine.format_response_body(jsn)
end
end