refac

backend/open_webui/retrieval/web/brave.py

@@ -1,20 +1,26 @@
 from __future__ import annotations
 
+import asyncio
 import logging
-import time
 
-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__)
 
+# Brave free-tier rate limit: 1 request per second.
+_RATE_LIMIT_RETRY_DELAY = 1.0
 
-def search_brave(api_key: str, query: str, count: int, filter_list: list[str | None] = None) -> list[SearchResult]:
-    """Search using Brave's Search API and return the results as a list of SearchResult objects.
 
-    Args:
-        api_key (str): A Brave Search API key
-        query (str): The query to search for
+async def search_brave(
+    api_key: str,
+    query: str,
+    count: int,
+    filter_list: list[str | None] | None = None,
+) -> list[SearchResult]:
+    """Query the Brave Web Search API and return normalised results.
+
+    Retries once on HTTP 429 (rate-limit) after a short delay.
     """
     url = 'https://api.search.brave.com/res/v1/web/search'
     headers = {
@@ -24,27 +30,27 @@ def search_brave(api_key: str, query: str, count: int, filter_list: list[str | N
     }
     params = {'q': query, 'count': count}
 
-    response = requests.get(url, headers=headers, params=params)
+    session = await get_session()
+    async with session.get(url, headers=headers, params=params) as response:
+        if response.status == 429:
+            log.info('Brave Search rate-limited (429); retrying after %.1fs', _RATE_LIMIT_RETRY_DELAY)
+            await asyncio.sleep(_RATE_LIMIT_RETRY_DELAY)
+            async with session.get(url, headers=headers, params=params) as retry_resp:
+                retry_resp.raise_for_status()
+                payload = await retry_resp.json()
+        else:
+            response.raise_for_status()
+            payload = await response.json()
 
-    # Handle 429 rate limiting - Brave free tier allows 1 request/second
-    # If rate limited, wait 1 second and retry once before failing
-    if response.status_code == 429:
-        log.info('Brave Search API rate limited (429), retrying after 1 second...')
-        time.sleep(1)
-        response = requests.get(url, headers=headers, params=params)
-
-    response.raise_for_status()
-
-    json_response = response.json()
-    results = json_response.get('web', {}).get('results', [])
+    web_results = payload.get('web', {}).get('results', [])
     if filter_list:
-        results = get_filtered_results(results, filter_list)
+        web_results = get_filtered_results(web_results, filter_list)
 
     return [
         SearchResult(
-            link=result['url'],
-            title=result.get('title'),
-            snippet=result.get('description'),
+            link=item.get('url', ''),
+            title=item.get('title'),
+            snippet=item.get('description'),
         )
-        for result in results[:count]
+        for item in web_results[:count]
     ]