Merge pull request #7694 from erics118/main

feat: add kagi search engine
2024-12-28 06:42:47 +00:00 · 2024-12-09 16:12:32 -08:00 · 2024-12-09 16:12:32 -08:00 · d4d6d1e7db
commit d4d6d1e7db
parent a4ed027498 b825947745
4 changed files with 88 additions and 0 deletions
--- a/backend/open_webui/apps/retrieval/main.py
+++ b/backend/open_webui/apps/retrieval/main.py
@ -29,6 +29,7 @@ from open_webui.apps.retrieval.loaders.youtube import YoutubeLoader
 from open_webui.apps.retrieval.web.main import SearchResult
 from open_webui.apps.retrieval.web.utils import get_web_loader
 from open_webui.apps.retrieval.web.brave import search_brave
+from open_webui.apps.retrieval.web.kagi import search_kagi
 from open_webui.apps.retrieval.web.mojeek import search_mojeek
 from open_webui.apps.retrieval.web.duckduckgo import search_duckduckgo
 from open_webui.apps.retrieval.web.google_pse import search_google_pse
@ -54,6 +55,7 @@ from open_webui.apps.retrieval.utils import (
 from open_webui.apps.webui.models.files import Files
 from open_webui.config import (
    BRAVE_SEARCH_API_KEY,
+    KAGI_SEARCH_API_KEY,
    MOJEEK_SEARCH_API_KEY,
    TIKTOKEN_ENCODING_NAME,
    RAG_TEXT_SPLITTER,
@ -184,6 +186,7 @@ app.state.config.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
 app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
 app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
 app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
+app.state.config.KAGI_SEARCH_API_KEY = KAGI_SEARCH_API_KEY
 app.state.config.MOJEEK_SEARCH_API_KEY = MOJEEK_SEARCH_API_KEY
 app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
 app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
@ -484,6 +487,7 @@ async def get_rag_config(user=Depends(get_admin_user)):
                "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
                "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
                "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
+                "kagi_search_api_key": app.state.config.KAGI_SEARCH_API_KEY,
                "mojeek_search_api_key": app.state.config.MOJEEK_SEARCH_API_KEY,
                "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": app.state.config.SERPSTACK_HTTPS,
@ -531,6 +535,7 @@ class WebSearchConfig(BaseModel):
    google_pse_api_key: Optional[str] = None
    google_pse_engine_id: Optional[str] = None
    brave_search_api_key: Optional[str] = None
+    kagi_search_api_key: Optional[str] = None
    mojeek_search_api_key: Optional[str] = None
    serpstack_api_key: Optional[str] = None
    serpstack_https: Optional[bool] = None
@ -603,6 +608,9 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
        app.state.config.BRAVE_SEARCH_API_KEY = (
            form_data.web.search.brave_search_api_key
        )
+        app.state.config.KAGI_SEARCH_API_KEY = (
+            form_data.web.search.kagi_search_api_key
+        )
        app.state.config.MOJEEK_SEARCH_API_KEY = (
            form_data.web.search.mojeek_search_api_key
        )
@ -657,6 +665,7 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
                "google_pse_api_key": app.state.config.GOOGLE_PSE_API_KEY,
                "google_pse_engine_id": app.state.config.GOOGLE_PSE_ENGINE_ID,
                "brave_search_api_key": app.state.config.BRAVE_SEARCH_API_KEY,
+                "kagi_search_api_key": app.state.config.KAGI_SEARCH_API_KEY,
                "mojeek_search_api_key": app.state.config.MOJEEK_SEARCH_API_KEY,
                "serpstack_api_key": app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": app.state.config.SERPSTACK_HTTPS,
@ -1162,6 +1171,7 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
    - SEARXNG_QUERY_URL
    - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
    - BRAVE_SEARCH_API_KEY
+    - KAGI_SEARCH_API_KEY
    - MOJEEK_SEARCH_API_KEY
    - SERPSTACK_API_KEY
    - SERPER_API_KEY
@ -1209,6 +1219,16 @@ def search_web(engine: str, query: str) -> list[SearchResult]:
            )
        else:
            raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
+    elif engine == "kagi":
+        if app.state.config.KAGI_SEARCH_API_KEY:
+            return search_kagi(
+                app.state.config.KAGI_SEARCH_API_KEY,
+                query,
+                app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
+                app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
+            )
+        else:
+            raise Exception("No KAGI_SEARCH_API_KEY found in environment variables")
    elif engine == "mojeek":
        if app.state.config.MOJEEK_SEARCH_API_KEY:
            return search_mojeek(
--- a/backend/open_webui/apps/retrieval/web/kagi.py
+++ b/backend/open_webui/apps/retrieval/web/kagi.py
@ -0,0 +1,50 @@
+import logging
+from typing import Optional
+
+import requests
+from open_webui.apps.retrieval.web.main import SearchResult, get_filtered_results
+from open_webui.env import SRC_LOG_LEVELS
+
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["RAG"])
+
+
+def search_kagi(
+    api_key: str, query: str, count: int, filter_list: Optional[list[str]] = None
+) -> list[SearchResult]:
+    """Search using Kagi's Search API and return the results as a list of SearchResult objects.
+
+    The Search API will inherit the settings in your account, including results personalization and snippet length.
+
+    Args:
+        api_key (str): A Kagi Search API key
+        query (str): The query to search for
+        count (int): The number of results to return
+    """
+    url = "https://kagi.com/api/v0/search"
+    headers = {
+        "Authorization": f"Bot {api_key}",
+    }
+    params = {"q": query, "limit": count}
+
+    response = requests.get(url, headers=headers, params=params)
+    response.raise_for_status()
+    json_response = response.json()
+    search_results = json_response.get("data", [])
+    
+    results = [
+        SearchResult(
+            link=result["url"],
+            title=result["title"],
+            snippet=result.get("snippet")
+        )
+        for result in search_results
+        if result["t"] == 0
+    ]
+    
+    print(results)
+
+    if filter_list:
+        results = get_filtered_results(results, filter_list)
+
+    return results
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@ -1380,6 +1380,12 @@ BRAVE_SEARCH_API_KEY = PersistentConfig(
    os.getenv("BRAVE_SEARCH_API_KEY", ""),
 )

+KAGI_SEARCH_API_KEY = PersistentConfig(
+    "KAGI_SEARCH_API_KEY",
+    "rag.web.search.kagi_search_api_key",
+    os.getenv("KAGI_SEARCH_API_KEY", ""),
+)
+
 MOJEEK_SEARCH_API_KEY = PersistentConfig(
    "MOJEEK_SEARCH_API_KEY",
    "rag.web.search.mojeek_search_api_key",
--- a/src/lib/components/admin/Settings/WebSearch.svelte
+++ b/src/lib/components/admin/Settings/WebSearch.svelte
@ -16,6 +16,7 @@
 		'searxng',
 		'google_pse',
 		'brave',
+		'kagi',
 		'mojeek',
 		'serpstack',
 		'serper',
@ -155,6 +156,17 @@
 									bind:value={webConfig.search.brave_search_api_key}
 								/>
 							</div>
+						{:else if webConfig.search.engine === 'kagi'}
+							<div>
+								<div class=" self-center text-xs font-medium mb-1">
+									{$i18n.t('Kagi Search API Key')}
+								</div>
+
+								<SensitiveInput
+									placeholder={$i18n.t('Enter Kagi Search API Key')}
+									bind:value={webConfig.search.kagi_search_api_key}
+								/>
+							</div>
 						{:else if webConfig.search.engine === 'mojeek'}
 							<div>
 								<div class=" self-center text-xs font-medium mb-1">