Coverage for open_webui/retrieval/web/bocha.py: 24%

1import logging 

2from typing import Optional 

3 

4import requests 

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

6from open_webui.utils.json_codec import JSONCodec 

7 

8log = logging.getLogger(__name__) 

9 

10 

11def _parse_response(response): 

12 results = [] 

13 if 'data' in response: 

14 data = response['data'] 

15 if 'webPages' in data: 

16 webPages = data['webPages'] 

17 if 'value' in webPages: 

18 results = [ 

19 { 

20 'id': item.get('id', ''), 

21 'name': item.get('name', ''), 

22 'url': item.get('url', ''), 

23 'snippet': item.get('snippet', ''), 

24 'summary': item.get('summary', ''), 

25 'siteName': item.get('siteName', ''), 

26 'siteIcon': item.get('siteIcon', ''), 

27 'datePublished': item.get('datePublished', '') or item.get('dateLastCrawled', ''), 

28 } 

29 for item in webPages['value'] 

30 ] 

31 return results 

32 

33 

34def search_bocha(api_key: str, query: str, count: int, filter_list: Optional[list[str]] = None) -> list[SearchResult]: 

35 """Search using Bocha's Search API and return the results as a list of SearchResult objects. 

36 

37 Args: 

38 api_key (str): A Bocha Search API key 

39 query (str): The query to search for 

40 """ 

41 url = 'https://api.bochaai.com/v1/web-search?utm_source=ollama' 

42 headers = {'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'} 

43 

44 payload = JSONCodec.dumps({'query': query, 'summary': True, 'freshness': 'noLimit', 'count': count}) 

45 

46 response = requests.post(url, headers=headers, data=payload, timeout=5) 

47 response.raise_for_status() 

48 results = _parse_response(response.json()) 

49 

50 if filter_list: 

51 results = get_filtered_results(results, filter_list) 

52 

53 return [ 

54 SearchResult(link=result['url'], title=result.get('name'), snippet=result.get('summary')) 

55 for result in results[:count] 

56 ]