feat: serpapi (#1198)

This commit is contained in:
yetone 2025-02-06 20:10:55 +08:00 committed by GitHub
parent 402c7f9665
commit 152d0c609c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 99 additions and 13 deletions

View File

@ -513,6 +513,35 @@ Avante provides a set of default providers, but users can also create their own
For more information, see [Custom Providers](https://github.com/yetone/avante.nvim/wiki/Custom-providers)
## 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`:
```lua
web_search_engine = {
provider = "tavily", -- tavily or serpapi
}
```
You need to set the environment variable `TAVILY_API_KEY` or `SERPAPI_API_KEY` to use tavily or serpapi.
## Disable Tools
Avante enables tools by default, but some LLM models do not support tools. You can disable tools by setting `disable_tools = true` for the provider. For example:
```lua
{
claude = {
endpoint = "https://api.anthropic.com",
model = "claude-3-5-sonnet-20241022",
timeout = 30000, -- Timeout in milliseconds
temperature = 0,
max_tokens = 4096,
disable_tools = true, -- disable tools!
},
}
```
## Custom prompts
By default, `avante.nvim` provides three different modes to interact with: `planning`, `editing`, and `suggesting`, followed with three different prompts per mode.

View File

@ -1,6 +1,8 @@
---NOTE: user will be merged with defaults and
---we add a default var_accessor for this table to config values.
---@alias WebSearchEngineProviderResponseBodyFormatter fun(body: table): (string, string?)
local Utils = require("avante.utils")
---@class avante.CoreConfig: avante.Config
@ -22,11 +24,45 @@ M._defaults = {
tokenizer = "tiktoken",
web_search_engine = {
provider = "tavily",
providers = {
tavily = {
api_key_name = "TAVILY_API_KEY",
provider_opts = {
extra_request_body = {
time_range = "d",
include_answer = "basic",
},
---@type WebSearchEngineProviderResponseBodyFormatter
format_response_body = function(body) return body.anwser, nil end,
},
serpapi = {
api_key_name = "SERPAPI_API_KEY",
extra_request_body = {
engine = "google",
google_domain = "google.com",
},
---@type WebSearchEngineProviderResponseBodyFormatter
format_response_body = function(body)
if body.answer_box ~= nil then return body.answer_box.result, nil end
if body.organic_results ~= nil then
local jsn = vim
.iter(body.organic_results)
:map(
function(result)
return {
title = result.title,
link = result.link,
snippet = result.snippet,
}
end
)
:totable()
if #jsn > 5 then jsn = vim.list_slice(jsn, 1, 5) end
return vim.json.encode(jsn), nil
end
return "", nil
end,
},
},
},
---@type AvanteSupportedProvider
openai = {

View File

@ -285,14 +285,18 @@ end
---@return string|nil result
---@return string|nil error
function M.web_search(opts, on_log)
local provider_type = Config.web_search_engine.provider
if provider_type == nil then return nil, "Search engine provider is not set" end
if on_log then on_log("provider: " .. provider_type) end
if on_log then on_log("query: " .. opts.query) end
local search_engine = Config.web_search_engine
if search_engine.provider == "tavily" then
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)
if api_key == nil or api_key == "" then
return nil, "Environment variable " .. search_engine.api_key_name .. " is not set"
end
if provider_type == "tavily" then
local resp = curl.post("https://api.tavily.com/search", {
headers = {
["Content-Type"] = "application/json",
@ -300,11 +304,28 @@ function M.web_search(opts, on_log)
},
body = vim.json.encode(vim.tbl_deep_extend("force", {
query = opts.query,
}, search_engine.provider_opts)),
}, search_engine.extra_request_body)),
})
if resp.status ~= 200 then return nil, "Error: " .. resp.body end
local jsn = vim.json.decode(resp.body)
return jsn.anwser, nil
return search_engine.format_response_body(jsn)
elseif provider_type == "serpapi" then
local query_params = vim.tbl_deep_extend("force", {
api_key = api_key,
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://serpapi.com/search?" .. 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