From ab63b52ffbbf3fd7014cc901e1393e4ba678b996 Mon Sep 17 00:00:00 2001
From: Limbo Peng <iwinux@gmail.com>
Date: Tue, 4 Mar 2025 23:47:04 +0800
Subject: [PATCH] feat: add Brave Search as web search engine provider (#1481)

* feat: add Brave Search as web search engine provider

* docs: update README
---
 README.md                | 26 +++++++++++++++++---------
 lua/avante/config.lua    | 22 ++++++++++++++++++++++
 lua/avante/llm_tools.lua | 17 +++++++++++++++++
 3 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index a6afb30..dd64181 100644
--- a/README.md
+++ b/README.md
@@ -673,13 +673,14 @@ Additionally, RAG Service also depends on Docker! (For macOS users, OrbStack is
 
 Avante's tools include some web search engines, currently support:
 
-- [tavily](https://tavily.com/)
-- [serpapi](https://serpapi.com/)
-- [searchapi](https://www.searchapi.io/)
-- google's [programmable search engine](https://developers.google.com/custom-search/v1/overview)
-- [kagi](https://help.kagi.com/kagi/api/search.html)
+- [Tavily](https://tavily.com/)
+- [SerpApi](https://serpapi.com/)
+- [SearchAPI](https://www.searchapi.io/)
+- Google's [Programmable Search Engine](https://developers.google.com/custom-search/v1/overview)
+- [Kagi](https://help.kagi.com/kagi/api/search.html)
+- [Brave Search](https://api-dashboard.search.brave.com/app/documentation/web-search/get-started)
 
-The default is tavily, and can be changed through configuring `Config.web_search_engine.provider`:
+The default is Tavily, and can be changed through configuring `Config.web_search_engine.provider`:
 
 ```lua
 web_search_engine = {
@@ -687,9 +688,16 @@ web_search_engine = {
 }
 ```
 
-You need to set the environment variable `TAVILY_API_KEY` , `SERPAPI_API_KEY`, `SEARCHAPI_API_KEY` to use tavily or serpapi or searchapi.
-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.
-To use kagi, set the `KAGI_API_KEY` as the [API Token](https://kagi.com/settings?p=api).
+Environment variables required for providers:
+
+- Tavily: `TAVILY_API_KEY`
+- SerpApi: `SERPAPI_API_KEY`
+- SearchAPI: `SEARCHAPI_API_KEY`
+- Google:
+  - `GOOGLE_SEARCH_API_KEY` as the [API key](https://developers.google.com/custom-search/v1/overview)
+  - `GOOGLE_SEARCH_ENGINE_ID` as the [search engine](https://programmablesearchengine.google.com) ID
+- Kagi: `KAGI_API_KEY` as the [API Token](https://kagi.com/settings?p=api)
+- Brave Search: `BRAVE_API_KEY` as the [API key](https://api-dashboard.search.brave.com/app/keys)
 
 ## Disable Tools
 
diff --git a/lua/avante/config.lua b/lua/avante/config.lua
index c865004..34bfa94 100644
--- a/lua/avante/config.lua
+++ b/lua/avante/config.lua
@@ -161,6 +161,28 @@ M._defaults = {
           return "", nil
         end,
       },
+      brave = {
+        api_key_name = "BRAVE_API_KEY",
+        extra_request_body = {
+          count = "10",
+          result_filter = "web",
+        },
+        format_response_body = function(body)
+          if body.web == nil then return "", nil end
+
+          local jsn = vim.iter(body.web.results):map(
+            function(result)
+              return {
+                title = result.title,
+                url = result.url,
+                snippet = result.description,
+              }
+            end
+          )
+
+          return vim.json.encode(jsn), nil
+        end,
+      },
     },
   },
   ---@type AvanteSupportedProvider
diff --git a/lua/avante/llm_tools.lua b/lua/avante/llm_tools.lua
index 5380460..6b4dd9f 100644
--- a/lua/avante/llm_tools.lua
+++ b/lua/avante/llm_tools.lua
@@ -378,6 +378,23 @@ 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 == "brave" then
+    local query_params = vim.tbl_deep_extend("force", {
+      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://api.search.brave.com/res/v1/web/search?" .. query_string, {
+      headers = {
+        ["Content-Type"] = "application/json",
+        ["X-Subscription-Token"] = api_key,
+      },
+    })
+    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