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

1import logging 

2import time 

3from typing import Optional 

4 

5import requests 

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

7 

8log = logging.getLogger(__name__) 

9 

10 

11def search_brave_llm_context( 

12 api_key: str, 

13 query: str, 

14 count: int, 

15 filter_list: Optional[list[str]] = None, 

16 context_tokens: int = 8192, 

17) -> list[SearchResult]: 

18 """Search using Brave's LLM Context API and return pre-extracted, relevance-scored 

19 page content ready for LLM consumption. 

20 

21 Uses /res/v1/llm/context instead of /res/v1/web/search. Same API key, same pricing. 

22 Returns full extracted passages per URL rather than short snippets, eliminating 

23 the need for post-search scraping. 

24 

25 Args: 

26 api_key (str): A Brave Search API key (same key as web search) 

27 query (str): The query to search for 

28 count (int): Maximum number of results to return 

29 filter_list (list[str], optional): Domain filter list 

30 context_tokens (int): Maximum total tokens to retrieve (1024–32768, default 8192) 

31 """ 

32 url = 'https://api.search.brave.com/res/v1/llm/context' 

33 headers = { 

34 'Accept': 'application/json', 

35 'Accept-Encoding': 'gzip', 

36 'X-Subscription-Token': api_key, 

37 } 

38 params = { 

39 'q': query, 

40 'count': count, 

41 'maximum_number_of_tokens': context_tokens, 

42 } 

43 

44 response = requests.get(url, headers=headers, params=params) 

45 

46 # Handle 429 rate limiting - same rate limits as web search 

47 if response.status_code == 429: 

48 log.info('Brave LLM Context API rate limited (429), retrying after 1 second...') 

49 time.sleep(1) 

50 response = requests.get(url, headers=headers, params=params) 

51 

52 response.raise_for_status() 

53 

54 json_response = response.json() 

55 results = json_response.get('grounding', {}).get('generic', []) 

56 if filter_list: 

57 results = get_filtered_results(results, filter_list) 

58 

59 return [ 

60 SearchResult( 

61 link=result['url'], 

62 title=result.get('title'), 

63 snippet='\n\n'.join(result.get('snippets', [])), 

64 ) 

65 for result in results[:count] 

66 ]