Merge pull request #9803 from BochaAI/main

add Bocha
2025-12-16 03:47:49 +01:00 · 2025-02-11 21:01:04 -08:00
parent 8199281eee 31360fe991
commit 8906a2e260
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(