Merge pull request #9803 from BochaAI/main

add Bocha
2025-03-27 02:02:31 +01:00 · 2025-02-11 21:01:04 -08:00 · 2025-02-11 21:01:04 -08:00 · 8906a2e260
commit 8906a2e260
parent 8199281eee 31360fe991
6 changed files with 112 additions and 0 deletions
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@ -1732,6 +1732,12 @@ MOJEEK_SEARCH_API_KEY = PersistentConfig(
    os.getenv("MOJEEK_SEARCH_API_KEY", ""),
 )

+BOCHA_SEARCH_API_KEY = PersistentConfig(
+    "BOCHA_SEARCH_API_KEY",
+    "rag.web.search.bocha_search_api_key",
+    os.getenv("BOCHA_SEARCH_API_KEY", ""),
+)
+
 SERPSTACK_API_KEY = PersistentConfig(
    "SERPSTACK_API_KEY",
    "rag.web.search.serpstack_api_key",
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@ -188,6 +188,7 @@ from open_webui.config import (
    EXA_API_KEY,
    KAGI_SEARCH_API_KEY,
    MOJEEK_SEARCH_API_KEY,
+    BOCHA_SEARCH_API_KEY,
    GOOGLE_PSE_API_KEY,
    GOOGLE_PSE_ENGINE_ID,
    GOOGLE_DRIVE_CLIENT_ID,
@ -526,6 +527,7 @@ 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.BOCHA_SEARCH_API_KEY = BOCHA_SEARCH_API_KEY
 app.state.config.SERPSTACK_API_KEY = SERPSTACK_API_KEY
 app.state.config.SERPSTACK_HTTPS = SERPSTACK_HTTPS
 app.state.config.SERPER_API_KEY = SERPER_API_KEY
--- a/backend/open_webui/retrieval/web/bocha.py
+++ b/backend/open_webui/retrieval/web/bocha.py
@ -0,0 +1,72 @@
+import logging
+from typing import Optional
+
+import requests
+import json
+from open_webui.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 _parse_response(response):
+    result = {}
+    if "data" in response:
+        data = response["data"]
+        if "webPages" in data:
+            webPages = data["webPages"]
+            if "value" in webPages:
+                result["webpage"] = [
+                    {
+                        "id": item.get("id", ""),
+                        "name": item.get("name", ""),
+                        "url": item.get("url", ""),
+                        "snippet": item.get("snippet", ""),
+                        "summary": item.get("summary", ""),
+                        "siteName": item.get("siteName", ""),
+                        "siteIcon": item.get("siteIcon", ""),
+                        "datePublished": item.get("datePublished", "") or item.get("dateLastCrawled", ""),
+                    }
+                    for item in webPages["value"]
+                ]
+    return result
+
+
+def search_bocha(
+    api_key: str, query: str, count: int, filter_list: Optional[list[str]] = None
+) -> list[SearchResult]:
+    """Search using Bocha's Search API and return the results as a list of SearchResult objects.
+
+    Args:
+        api_key (str): A Bocha Search API key
+        query (str): The query to search for
+    """
+    url = "https://api.bochaai.com/v1/web-search?utm_source=ollama"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    
+    payload = json.dumps({
+        "query": query,
+        "summary": True,
+        "freshness": "noLimit",
+        "count": count
+    })
+
+    response = requests.post(url, headers=headers, data=payload, timeout=5)
+    response.raise_for_status()
+    results = _parse_response(response.json())
+    print(results)
+    if filter_list:
+        results = get_filtered_results(results, filter_list)
+
+    return [
+        SearchResult(
+            link=result["url"], 
+            title=result.get("name"), 
+            snippet=result.get("summary")
+        )
+        for result in results.get("webpage", [])[:count]  
+    ]
+
--- a/backend/open_webui/routers/retrieval.py
+++ b/backend/open_webui/routers/retrieval.py
@ -45,6 +45,7 @@ from open_webui.retrieval.web.utils import get_web_loader
 from open_webui.retrieval.web.brave import search_brave
 from open_webui.retrieval.web.kagi import search_kagi
 from open_webui.retrieval.web.mojeek import search_mojeek
+from open_webui.retrieval.web.bocha import search_bocha
 from open_webui.retrieval.web.duckduckgo import search_duckduckgo
 from open_webui.retrieval.web.google_pse import search_google_pse
 from open_webui.retrieval.web.jina_search import search_jina
@ -379,6 +380,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
                "brave_search_api_key": request.app.state.config.BRAVE_SEARCH_API_KEY,
                "kagi_search_api_key": request.app.state.config.KAGI_SEARCH_API_KEY,
                "mojeek_search_api_key": request.app.state.config.MOJEEK_SEARCH_API_KEY,
+                "bocha_search_api_key": request.app.state.config.BOCHA_SEARCH_API_KEY,
                "serpstack_api_key": request.app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": request.app.state.config.SERPSTACK_HTTPS,
                "serper_api_key": request.app.state.config.SERPER_API_KEY,
@ -429,6 +431,7 @@ class WebSearchConfig(BaseModel):
    brave_search_api_key: Optional[str] = None
    kagi_search_api_key: Optional[str] = None
    mojeek_search_api_key: Optional[str] = None
+    bocha_search_api_key: Optional[str] = None
    serpstack_api_key: Optional[str] = None
    serpstack_https: Optional[bool] = None
    serper_api_key: Optional[str] = None
@ -525,6 +528,9 @@ async def update_rag_config(
        request.app.state.config.MOJEEK_SEARCH_API_KEY = (
            form_data.web.search.mojeek_search_api_key
        )
+        request.app.state.config.BOCHA_SEARCH_API_KEY = (
+            form_data.web.search.bocha_search_api_key
+        )
        request.app.state.config.SERPSTACK_API_KEY = (
            form_data.web.search.serpstack_api_key
        )
@ -591,6 +597,7 @@ async def update_rag_config(
                "brave_search_api_key": request.app.state.config.BRAVE_SEARCH_API_KEY,
                "kagi_search_api_key": request.app.state.config.KAGI_SEARCH_API_KEY,
                "mojeek_search_api_key": request.app.state.config.MOJEEK_SEARCH_API_KEY,
+                "bocha_search_api_key": request.app.state.config.BOCHA_SEARCH_API_KEY,
                "serpstack_api_key": request.app.state.config.SERPSTACK_API_KEY,
                "serpstack_https": request.app.state.config.SERPSTACK_HTTPS,
                "serper_api_key": request.app.state.config.SERPER_API_KEY,
@ -1113,6 +1120,7 @@ def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
    - BRAVE_SEARCH_API_KEY
    - KAGI_SEARCH_API_KEY
    - MOJEEK_SEARCH_API_KEY
+    - BOCHA_SEARCH_API_KEY
    - SERPSTACK_API_KEY
    - SERPER_API_KEY
    - SERPLY_API_KEY
@ -1180,6 +1188,16 @@ def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
            )
        else:
            raise Exception("No MOJEEK_SEARCH_API_KEY found in environment variables")
+    elif engine == "bocha":
+        if request.app.state.config.BOCHA_SEARCH_API_KEY:
+            return search_bocha(
+                request.app.state.config.BOCHA_SEARCH_API_KEY,
+                query,
+                request.app.state.config.RAG_WEB_SEARCH_RESULT_COUNT,
+                request.app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST,
+            )
+        else:
+            raise Exception("No BOCHA_SEARCH_API_KEY found in environment variables")
    elif engine == "serpstack":
        if request.app.state.config.SERPSTACK_API_KEY:
            return search_serpstack(
--- a/src/lib/components/admin/Settings/WebSearch.svelte
+++ b/src/lib/components/admin/Settings/WebSearch.svelte
@ -18,6 +18,7 @@
 		'brave',
 		'kagi',
 		'mojeek',
+		'bocha',
 		'serpstack',
 		'serper',
 		'serply',
@ -195,6 +196,17 @@
 									bind:value={webConfig.search.mojeek_search_api_key}
 								/>
 							</div>
+						{:else if webConfig.search.engine === 'bocha'}
+							<div>
+								<div class=" self-center text-xs font-medium mb-1">
+									{$i18n.t('Bocha Search API Key')}
+								</div>
+
+								<SensitiveInput
+									placeholder={$i18n.t('Enter Bocha Search API Key')}
+									bind:value={webConfig.search.bocha_search_api_key}
+								/>
+							</div>
 						{:else if webConfig.search.engine === 'serpstack'}
 							<div>
 								<div class=" self-center text-xs font-medium mb-1">
--- a/src/lib/i18n/locales/zh-CN/translation.json
+++ b/src/lib/i18n/locales/zh-CN/translation.json
@ -363,6 +363,7 @@
 	"Enter Model ID": "输入模型 ID",
 	"Enter model tag (e.g. {{modelTag}})": "输入模型标签 (例如：{{modelTag}})",
 	"Enter Mojeek Search API Key": "输入 Mojeek Search API 密钥",
+	"Enter Bocha Search API Key": "输入 Bocha Search API 密钥",
 	"Enter Number of Steps (e.g. 50)": "输入步骤数 (Steps) (例如：50)",
 	"Enter proxy URL (e.g. https://user:password@host:port)": "输入代理 URL (例如：https://用户名:密码@主机名:端口）",
 	"Enter reasoning effort": "设置推理努力",
@ -632,6 +633,7 @@
 	"Models Access": "访问模型列表",
 	"Models configuration saved successfully": "模型配置保存成功",
 	"Mojeek Search API Key": "Mojeek Search API 密钥",
+	"Bocha Search API Key": "Bocha Search API 密钥",
 	"more": "更多",
 	"More": "更多",
 	"Name": "名称",