refac

backend/open_webui/retrieval/web/serper.py

@@ -3,36 +3,39 @@ from __future__ import annotations
 import json
 import logging
 
-import requests
 from open_webui.retrieval.web.main import SearchResult, get_filtered_results
+from open_webui.utils.session_pool import get_session
 
 log = logging.getLogger(__name__)
 
 
-def search_serper(api_key: str, query: str, count: int, filter_list: list[str | None] = None) -> list[SearchResult]:
-    """Search using serper.dev's API and return the results as a list of SearchResult objects.
+async def search_serper(
+    api_key: str,
+    query: str,
+    count: int,
+    filter_list: list[str | None] | None = None,
+) -> list[SearchResult]:
+    """Query the serper.dev Google Search API and return normalised results.
 
-    Args:
-        api_key (str): A serper.dev API key
-        query (str): The query to search for
+    Results are sorted by their position field before truncation.
     """
     url = 'https://google.serper.dev/search'
-
-    payload = json.dumps({'q': query})
     headers = {'X-API-KEY': api_key, 'Content-Type': 'application/json'}
 
-    response = requests.request('POST', url, headers=headers, data=payload)
-    response.raise_for_status()
+    session = await get_session()
+    async with session.post(url, headers=headers, data=json.dumps({'q': query})) as response:
+        response.raise_for_status()
+        payload = await response.json()
 
-    json_response = response.json()
-    results = sorted(json_response.get('organic', []), key=lambda x: x.get('position', 0))
+    organic = sorted(payload.get('organic', []), key=lambda item: item.get('position', 0))
     if filter_list:
-        results = get_filtered_results(results, filter_list)
+        organic = get_filtered_results(organic, filter_list)
+
     return [
         SearchResult(
-            link=result['link'],
-            title=result.get('title'),
-            snippet=result.get('snippet'),
+            link=item.get('link', ''),
+            title=item.get('title'),
+            snippet=item.get('snippet'),
         )
-        for result in results[:count]
+        for item in organic[:count]
     ]