From fd6eea3e8489feab8a72e41401170d351c77f204 Mon Sep 17 00:00:00 2001 From: n4ze3m Date: Fri, 13 Dec 2024 20:03:52 +0530 Subject: [PATCH] feat: Add Brave Search API support --- src/assets/locale/da/settings.json | 4 + src/assets/locale/de/settings.json | 7 +- src/assets/locale/en/settings.json | 4 + src/assets/locale/es/settings.json | 4 + src/assets/locale/fa/settings.json | 4 + src/assets/locale/fr/settings.json | 4 + src/assets/locale/it/settings.json | 4 + src/assets/locale/ja-JP/settings.json | 4 + src/assets/locale/ko/settings.json | 4 + src/assets/locale/ml/settings.json | 4 + src/assets/locale/no/settings.json | 4 + src/assets/locale/pt-BR/settings.json | 4 + src/assets/locale/ru/settings.json | 4 + src/assets/locale/sv/settings.json | 4 + src/assets/locale/uk/settings.json | 4 + src/assets/locale/zh/settings.json | 4 + .../Option/Playground/PlaygroundChat.tsx | 2 +- .../Option/Settings/search-mode.tsx | 24 +++- src/services/search.ts | 30 +++- src/utils/search-provider.ts | 4 + src/web/search-engines/brave-api.ts | 128 ++++++++++++++++++ src/web/web.ts | 3 + 22 files changed, 248 insertions(+), 10 deletions(-) create mode 100644 src/web/search-engines/brave-api.ts diff --git a/src/assets/locale/da/settings.json b/src/assets/locale/da/settings.json index b612c68..1712a45 100644 --- a/src/assets/locale/da/settings.json +++ b/src/assets/locale/da/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave API Nøgle", + "placeholder": "Indtast din Brave API nøgle" } }, "system": { diff --git a/src/assets/locale/de/settings.json b/src/assets/locale/de/settings.json index 0e2ac78..581cc28 100644 --- a/src/assets/locale/de/settings.json +++ b/src/assets/locale/de/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG-URL" } + }, + "braveApi": { + "label": "Brave API-Schlüssel", + "placeholder": "Geben Sie Ihren Brave API-Schlüssel ein" } }, "system": { @@ -78,7 +82,8 @@ "label": "System zurücksetzen", "button": "Alles zurücksetzen", "confirm": "Sind Sie sicher, dass Sie einen Systemreset durchführen möchten? Dies löscht alle Daten und kann nicht rückgängig gemacht werden." - }, "export": { + }, + "export": { "label": "Chatverlauf, Wissensbasis und Prompts exportieren", "button": "Daten exportieren", "success": "Export erfolgreich" diff --git a/src/assets/locale/en/settings.json b/src/assets/locale/en/settings.json index f48b46b..d32dccf 100644 --- a/src/assets/locale/en/settings.json +++ b/src/assets/locale/en/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave API Key", + "placeholder": "Enter your Brave API key" } }, "system": { diff --git a/src/assets/locale/es/settings.json b/src/assets/locale/es/settings.json index b225c54..3809ff1 100644 --- a/src/assets/locale/es/settings.json +++ b/src/assets/locale/es/settings.json @@ -70,6 +70,10 @@ "url": { "label": "URL de SearXNG" } + }, + "braveApi": { + "label": "Clave API de Brave", + "placeholder": "Ingrese su clave API de Brave" } }, "system": { diff --git a/src/assets/locale/fa/settings.json b/src/assets/locale/fa/settings.json index 26be070..bb037ff 100644 --- a/src/assets/locale/fa/settings.json +++ b/src/assets/locale/fa/settings.json @@ -67,6 +67,10 @@ "url": { "label": "آدرس SearXNG" } + }, + "braveApi": { + "label": "کلید API بریو", + "placeholder": "کلید API بریو خود را وارد کنید" } }, "system": { diff --git a/src/assets/locale/fr/settings.json b/src/assets/locale/fr/settings.json index 97dec81..93ba685 100644 --- a/src/assets/locale/fr/settings.json +++ b/src/assets/locale/fr/settings.json @@ -70,6 +70,10 @@ "url": { "label": "URL SearXNG" } + }, + "braveApi": { + "label": "Clé API Brave", + "placeholder": "Entrez votre clé API Brave" } }, "system": { diff --git a/src/assets/locale/it/settings.json b/src/assets/locale/it/settings.json index 4aa58b9..d7d7007 100644 --- a/src/assets/locale/it/settings.json +++ b/src/assets/locale/it/settings.json @@ -70,6 +70,10 @@ "url": { "label": "URL SearXNG" } + }, + "braveApi": { + "label": "Chiave API Brave", + "placeholder": "Inserisci la tua chiave API Brave" } }, "system": { diff --git a/src/assets/locale/ja-JP/settings.json b/src/assets/locale/ja-JP/settings.json index 533aa1e..62363a5 100644 --- a/src/assets/locale/ja-JP/settings.json +++ b/src/assets/locale/ja-JP/settings.json @@ -73,6 +73,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave APIキー", + "placeholder": "Brave APIキーを入力してください" } }, "system": { diff --git a/src/assets/locale/ko/settings.json b/src/assets/locale/ko/settings.json index fe31b5f..f4f019a 100644 --- a/src/assets/locale/ko/settings.json +++ b/src/assets/locale/ko/settings.json @@ -73,6 +73,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave API 키", + "placeholder": "Brave API 키를 입력하세요" } }, "system": { diff --git a/src/assets/locale/ml/settings.json b/src/assets/locale/ml/settings.json index 640fc4c..e18ef5e 100644 --- a/src/assets/locale/ml/settings.json +++ b/src/assets/locale/ml/settings.json @@ -73,6 +73,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "ബ്രേവ് API കീ", + "placeholder": "നിങ്ങളുടെ ബ്രേവ് API കീ നൽകുക" } }, "system": { diff --git a/src/assets/locale/no/settings.json b/src/assets/locale/no/settings.json index 1b994a8..cd96712 100644 --- a/src/assets/locale/no/settings.json +++ b/src/assets/locale/no/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave API Nøkkel", + "placeholder": "Skriv inn din Brave API nøkkel" } }, "system": { diff --git a/src/assets/locale/pt-BR/settings.json b/src/assets/locale/pt-BR/settings.json index e7adfb7..6dbd407 100644 --- a/src/assets/locale/pt-BR/settings.json +++ b/src/assets/locale/pt-BR/settings.json @@ -70,6 +70,10 @@ "url": { "label": "URL do SearXNG" } + }, + "braveApi": { + "label": "Chave da API do Brave", + "placeholder": "Digite sua chave da API do Brave" } }, "system": { diff --git a/src/assets/locale/ru/settings.json b/src/assets/locale/ru/settings.json index 710ffdc..c71c037 100644 --- a/src/assets/locale/ru/settings.json +++ b/src/assets/locale/ru/settings.json @@ -71,6 +71,10 @@ "url": { "label": "URL-адрес SearXNG" } + }, + "braveApi": { + "label": "API-ключ Brave", + "placeholder": "Введите ваш API-ключ Brave" } }, "system": { diff --git a/src/assets/locale/sv/settings.json b/src/assets/locale/sv/settings.json index 7ed494a..712e9a3 100644 --- a/src/assets/locale/sv/settings.json +++ b/src/assets/locale/sv/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG URL" } + }, + "braveApi": { + "label": "Brave API-nyckel", + "placeholder": "Ange din Brave API-nyckel" } }, "system": { diff --git a/src/assets/locale/uk/settings.json b/src/assets/locale/uk/settings.json index f1526c3..7462317 100644 --- a/src/assets/locale/uk/settings.json +++ b/src/assets/locale/uk/settings.json @@ -70,6 +70,10 @@ "url": { "label": "SearXNG URL-адреса" } + }, + "braveApi": { + "label": "Ключ API Brave", + "placeholder": "Введіть ваш ключ API Brave" } }, "system": { diff --git a/src/assets/locale/zh/settings.json b/src/assets/locale/zh/settings.json index 24fa548..2557c44 100644 --- a/src/assets/locale/zh/settings.json +++ b/src/assets/locale/zh/settings.json @@ -73,6 +73,10 @@ "url": { "label": "SearXNG 网址" } + }, + "braveApi": { + "label": "Brave API 密钥", + "placeholder": "输入您的 Brave API 密钥" } }, "system": { diff --git a/src/components/Option/Playground/PlaygroundChat.tsx b/src/components/Option/Playground/PlaygroundChat.tsx index aaf3462..ea5361f 100644 --- a/src/components/Option/Playground/PlaygroundChat.tsx +++ b/src/components/Option/Playground/PlaygroundChat.tsx @@ -27,7 +27,7 @@ export const PlaygroundChat = () => { <>
+ className="custom-scrollbar grow flex flex-col md:translate-x-0 transition-transform duration-300 ease-in-out overflow-y-auto h-[calc(100vh-160px)]"> {messages.length === 0 && (
diff --git a/src/components/Option/Settings/search-mode.tsx b/src/components/Option/Settings/search-mode.tsx index 69ae112..932632b 100644 --- a/src/components/Option/Settings/search-mode.tsx +++ b/src/components/Option/Settings/search-mode.tsx @@ -2,7 +2,7 @@ import { SaveButton } from "@/components/Common/SaveButton" import { getSearchSettings, setSearchSettings } from "@/services/search" import { SUPPORTED_SERACH_PROVIDERS } from "@/utils/search-provider" import { useForm } from "@mantine/form" -import { useQuery, useQueryClient } from "@tanstack/react-query" +import { useQuery } from "@tanstack/react-query" import { Select, Skeleton, Switch, InputNumber, Input } from "antd" import { useTranslation } from "react-i18next" @@ -16,7 +16,8 @@ export const SearchModeSettings = () => { totalSearchResults: 0, visitSpecificWebsite: false, searxngURL: "", - searxngJSONMode: false + searxngJSONMode: false, + braveApiKey: "", } }) @@ -81,6 +82,25 @@ export const SearchModeSettings = () => {
)} + {form.values.searchProvider === "brave-api" && ( + <> +
+ + {t("generalSettings.webSearch.braveApi.label")} + +
+ +
+
+ + )}
{t("generalSettings.webSearch.searchMode.label")} diff --git a/src/services/search.ts b/src/services/search.ts index 66fb5cf..73b3327 100644 --- a/src/services/search.ts +++ b/src/services/search.ts @@ -1,6 +1,9 @@ import { Storage } from "@plasmohq/storage" const storage = new Storage() +const storage2 = new Storage({ + area: "local" +}) const TOTAL_SEARCH_RESULTS = 2 const DEFAULT_PROVIDER = "google" @@ -80,10 +83,20 @@ export const setSearxngURL = async (searxngURL: string) => { await storage.set("searxngURL", searxngURL) } +export const getBraveApiKey = async () => { + const braveApiKey = await storage2.get("braveApiKey") + return braveApiKey || "" +} + +export const setBraveApiKey = async (braveApiKey: string) => { + await storage2.set("braveApiKey", braveApiKey) +} + export const getSearchSettings = async () => { const [isSimpleInternetSearch, searchProvider, totalSearchResult, visitSpecificWebsite, searxngURL, - searxngJSONMode + searxngJSONMode, + braveApiKey ] = await Promise.all([ getIsSimpleInternetSearch(), @@ -91,7 +104,8 @@ export const getSearchSettings = async () => { totalSearchResults(), getIsVisitSpecificWebsite(), getSearxngURL(), - isSearxngJSONMode() + isSearxngJSONMode(), + getBraveApiKey() ]) return { @@ -100,7 +114,8 @@ export const getSearchSettings = async () => { totalSearchResults: totalSearchResult, visitSpecificWebsite, searxngURL, - searxngJSONMode + searxngJSONMode, + braveApiKey } } @@ -110,14 +125,16 @@ export const setSearchSettings = async ({ totalSearchResults, visitSpecificWebsite, searxngJSONMode, - searxngURL + searxngURL, + braveApiKey }: { isSimpleInternetSearch: boolean searchProvider: string totalSearchResults: number visitSpecificWebsite: boolean searxngURL: string - searxngJSONMode: boolean + searxngJSONMode: boolean, + braveApiKey: string }) => { await Promise.all([ setIsSimpleInternetSearch(isSimpleInternetSearch), @@ -125,6 +142,7 @@ export const setSearchSettings = async ({ setTotalSearchResults(totalSearchResults), setIsVisitSpecificWebsite(visitSpecificWebsite), setSearxngJSONMode(searxngJSONMode), - setSearxngURL(searxngURL) + setSearxngURL(searxngURL), + setBraveApiKey(braveApiKey) ]) } diff --git a/src/utils/search-provider.ts b/src/utils/search-provider.ts index 404c13f..24004a4 100644 --- a/src/utils/search-provider.ts +++ b/src/utils/search-provider.ts @@ -18,5 +18,9 @@ export const SUPPORTED_SERACH_PROVIDERS = [ { label: "Searxng", value: "searxng" + }, + { + label: "Brave Search API", + value: "brave-api" } ] \ No newline at end of file diff --git a/src/web/search-engines/brave-api.ts b/src/web/search-engines/brave-api.ts new file mode 100644 index 0000000..5e13312 --- /dev/null +++ b/src/web/search-engines/brave-api.ts @@ -0,0 +1,128 @@ +import { cleanUrl } from "~/libs/clean-url" +import { getIsSimpleInternetSearch, totalSearchResults, getBraveApiKey } from "@/services/search" +import { pageAssistEmbeddingModel } from "@/models/embedding" +import type { Document } from "@langchain/core/documents" +import { RecursiveCharacterTextSplitter } from "langchain/text_splitter" +import { MemoryVectorStore } from "langchain/vectorstores/memory" +import { PageAssistHtmlLoader } from "~/loader/html" +import { + defaultEmbeddingChunkOverlap, + defaultEmbeddingChunkSize, + defaultEmbeddingModelForRag, + getOllamaURL +} from "~/services/ollama" + +interface BraveAPIResult { + title: string + url: string + description: string +} + +interface BraveAPIResponse { + web: { + results: BraveAPIResult[] + } +} + +export const braveAPISearch = async (query: string) => { + const braveApiKey = await getBraveApiKey() + if (!braveApiKey || braveApiKey.trim() === "") { + throw new Error("Brave API key not configured") + } + const results = await apiBraveSearch(braveApiKey, query) + const TOTAL_SEARCH_RESULTS = await totalSearchResults() + + const searchResults = results.slice(0, TOTAL_SEARCH_RESULTS) + + const isSimpleMode = await getIsSimpleInternetSearch() + + if (isSimpleMode) { + await getOllamaURL() + return searchResults.map((result) => { + return { + url: result.link, + content: result.content + } + }) + } + + const docs: Document>[] = [] + try { + for (const result of searchResults) { + const loader = new PageAssistHtmlLoader({ + html: "", + url: result.link + }) + + const documents = await loader.loadByURL() + documents.forEach((doc) => { + docs.push(doc) + }) + } + } catch (error) { + console.error(error) + } + + const ollamaUrl = await getOllamaURL() + const embeddingModel = await defaultEmbeddingModelForRag() + const ollamaEmbedding = await pageAssistEmbeddingModel({ + model: embeddingModel || "", + baseUrl: cleanUrl(ollamaUrl) + }) + + const chunkSize = await defaultEmbeddingChunkSize() + const chunkOverlap = await defaultEmbeddingChunkOverlap() + const textSplitter = new RecursiveCharacterTextSplitter({ + chunkSize, + chunkOverlap + }) + + const chunks = await textSplitter.splitDocuments(docs) + const store = new MemoryVectorStore(ollamaEmbedding) + await store.addDocuments(chunks) + + const resultsWithEmbeddings = await store.similaritySearch(query, 3) + + const searchResult = resultsWithEmbeddings.map((result) => { + return { + url: result.metadata.url, + content: result.pageContent + } + }) + + return searchResult +} + +const apiBraveSearch = async (braveApiKey: string, query: string) => { + const TOTAL_SEARCH_RESULTS = await totalSearchResults() + + const searchURL = `https://api.search.brave.com/res/v1/web/search?q=${query}&count=${TOTAL_SEARCH_RESULTS}` + + const abortController = new AbortController() + setTimeout(() => abortController.abort(), 20000) + + try { + const response = await fetch(searchURL, { + signal: abortController.signal, + headers: { + "X-Subscription-Token": braveApiKey, + Accept: "application/json", + } + }) + + if (!response.ok) { + return [] + } + + const data = await response.json() as BraveAPIResponse + + return data?.web?.results.map(result => ({ + title: result.title, + link: result.url, + content: result.description + })) + } catch (error) { + console.error('Brave API search failed:', error) + return [] + } +} diff --git a/src/web/web.ts b/src/web/web.ts index 6c48f1b..42d8ca9 100644 --- a/src/web/web.ts +++ b/src/web/web.ts @@ -6,6 +6,7 @@ import { webSogouSearch } from "./search-engines/sogou" import { webBraveSearch } from "./search-engines/brave" import { getWebsiteFromQuery, processSingleWebsite } from "./website" import { searxngSearch } from "./search-engines/searxng" +import { braveAPISearch } from "./search-engines/brave-api" const getHostName = (url: string) => { try { @@ -26,6 +27,8 @@ const searchWeb = (provider: string, query: string) => { return webBraveSearch(query) case "searxng": return searxngSearch(query) + case "brave-api": + return braveAPISearch(query) default: return webGoogleSearch(query) }