feat: serpapi (#1198)

This commit is contained in:
yetone
2025-02-06 20:10:55 +08:00
committed by GitHub
parent 402c7f9665
commit 152d0c609c
3 changed files with 99 additions and 13 deletions

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,10 +24,44 @@ M._defaults = {
tokenizer = "tiktoken",
web_search_engine = {
provider = "tavily",
api_key_name = "TAVILY_API_KEY",
provider_opts = {
time_range = "d",
include_answer = "basic",
providers = {
tavily = {
api_key_name = "TAVILY_API_KEY",
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

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
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
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