Merge remote-tracking branch 'upstream/dev' into playwright

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]  
+    ]
+

backend/open_webui/retrieval/web/google_pse.py

@@ -8,7 +8,6 @@ from open_webui.env import SRC_LOG_LEVELS
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
 
-
 def search_google_pse(
     api_key: str,
     search_engine_id: str,
@@ -17,34 +16,51 @@ def search_google_pse(
     filter_list: Optional[list[str]] = None,
 ) -> list[SearchResult]:
     """Search using Google's Programmable Search Engine API and return the results as a list of SearchResult objects.
+    Handles pagination for counts greater than 10.
 
     Args:
         api_key (str): A Programmable Search Engine API key
         search_engine_id (str): A Programmable Search Engine ID
         query (str): The query to search for
+        count (int): The number of results to return (max 100, as PSE max results per query is 10 and max page is 10)
+        filter_list (Optional[list[str]], optional): A list of keywords to filter out from results. Defaults to None.
+
+    Returns:
+        list[SearchResult]: A list of SearchResult objects.
     """
     url = "https://www.googleapis.com/customsearch/v1"
-
     headers = {"Content-Type": "application/json"}
-    params = {
-        "cx": search_engine_id,
-        "q": query,
-        "key": api_key,
-        "num": count,
-    }
+    all_results = []
+    start_index = 1  # Google PSE start parameter is 1-based
 
-    response = requests.request("GET", url, headers=headers, params=params)
-    response.raise_for_status()
+    while count > 0:
+        num_results_this_page = min(count, 10)  # Google PSE max results per page is 10
+        params = {
+            "cx": search_engine_id,
+            "q": query,
+            "key": api_key,
+            "num": num_results_this_page,
+            "start": start_index,
+        }
+        response = requests.request("GET", url, headers=headers, params=params)
+        response.raise_for_status()
+        json_response = response.json()
+        results = json_response.get("items", [])
+        if results: # check if results are returned. If not, no more pages to fetch.
+            all_results.extend(results)
+            count -= len(results) # Decrement count by the number of results fetched in this page.
+            start_index += 10 # Increment start index for the next page
+        else:
+            break # No more results from Google PSE, break the loop
 
-    json_response = response.json()
-    results = json_response.get("items", [])
     if filter_list:
-        results = get_filtered_results(results, filter_list)
+        all_results = get_filtered_results(all_results, filter_list)
+
     return [
         SearchResult(
             link=result["link"],
             title=result.get("title"),
             snippet=result.get("snippet"),
         )
-        for result in results
+        for result in all_results
     ]

backend/open_webui/retrieval/web/jina_search.py

@@ -25,13 +25,10 @@ def search_jina(api_key: str, query: str, count: int) -> list[SearchResult]:
         "Accept": "application/json",
         "Content-Type": "application/json",
         "Authorization": api_key,
-        "X-Retain-Images": "none"
+        "X-Retain-Images": "none",
     }
 
-    payload = {
-        "q": query,
-        "count": count if count <= 10 else 10
-    }
+    payload = {"q": query, "count": count if count <= 10 else 10}
 
     url = str(URL(jina_search_endpoint))
     response = requests.post(url, headers=headers, json=payload)