diff --git a/README.md b/README.md index a9979de..f9d79ba 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/lua/avante/config.lua b/lua/avante/config.lua index 123ac11..0364d02 100644 --- a/lua/avante/config.lua +++ b/lua/avante/config.lua @@ -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 diff --git a/lua/avante/llm_tools.lua b/lua/avante/llm_tools.lua index b2b0dce..5117a05 100644 --- a/lua/avante/llm_tools.lua +++ b/lua/avante/llm_tools.lua @@ -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