Coverage for open_webui/retrieval/web/brave.py: 28%

1from __future__ import annotations 

2 

3import asyncio 

4import logging 

5 

6from open_webui.retrieval.web.main import SearchResult, get_filtered_results 

7from open_webui.utils.session_pool import get_session 

8 

9log = logging.getLogger(__name__) 

10 

11# Brave free-tier rate limit: 1 request per second. 

12_RATE_LIMIT_RETRY_DELAY = 1.0 

13 

14 

15async def search_brave( 

16 api_key: str, 

17 query: str, 

18 count: int, 

19 filter_list: list[str] | None = None, 

20) -> list[SearchResult]: 

21 """Query the Brave Web Search API and return normalised results. 

22 

23 Retries once on HTTP 429 (rate-limit) after a short delay. 

24 """ 

25 url = 'https://api.search.brave.com/res/v1/web/search' 

26 headers = { 

27 'Accept': 'application/json', 

28 'Accept-Encoding': 'gzip', 

29 'X-Subscription-Token': api_key, 

30 } 

31 params = {'q': query, 'count': count} 

32 

33 session = await get_session() 

34 async with session.get(url, headers=headers, params=params) as response: 

35 if response.status == 429: 

36 log.info('Brave Search rate-limited (429); retrying after %.1fs', _RATE_LIMIT_RETRY_DELAY) 

37 await asyncio.sleep(_RATE_LIMIT_RETRY_DELAY) 

38 async with session.get(url, headers=headers, params=params) as retry_resp: 

39 retry_resp.raise_for_status() 

40 payload = await retry_resp.json() 

41 else: 

42 response.raise_for_status() 

43 payload = await response.json() 

44 

45 web_results = payload.get('web', {}).get('results', []) 

46 if filter_list: 

47 web_results = get_filtered_results(web_results, filter_list) 

48 

49 return [ 

50 SearchResult( 

51 link=item.get('url', ''), 

52 title=item.get('title'), 

53 snippet=item.get('description'), 

54 ) 

55 for item in web_results[:count] 

56 ]