Coverage for open_webui/retrieval/web/serply.py: 32%

1from __future__ import annotations 

2 

3import logging 

4from urllib.parse import urlencode 

5 

6import requests 

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

8 

9log = logging.getLogger(__name__) 

10 

11 

12def search_serply( 

13 api_key: str, 

14 query: str, 

15 count: int, 

16 hl: str = 'us', 

17 limit: int = 10, 

18 device_type: str = 'desktop', 

19 proxy_location: str = 'US', 

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

21) -> list[SearchResult]: 

22 """Search using serper.dev's API and return the results as a list of SearchResult objects. 

23 

24 Args: 

25 api_key (str): A serply.io API key 

26 query (str): The query to search for 

27 hl (str): Host Language code to display results in (reference https://developers.google.com/custom-search/docs/xml_results?hl=en#wsInterfaceLanguages) 

28 limit (int): The maximum number of results to return [10-100, defaults to 10] 

29 """ 

30 log.info('Searching with Serply') 

31 

32 url = 'https://api.serply.io/v1/search/' 

33 

34 query_payload = { 

35 'q': query, 

36 'language': 'en', 

37 'num': limit, 

38 'gl': proxy_location.upper(), 

39 'hl': hl.lower(), 

40 } 

41 

42 url = f'{url}{urlencode(query_payload)}' 

43 headers = { 

44 'X-API-KEY': api_key, 

45 'X-User-Agent': device_type, 

46 'User-Agent': 'open-webui', 

47 'X-Proxy-Location': proxy_location, 

48 } 

49 

50 response = requests.request('GET', url, headers=headers) 

51 response.raise_for_status() 

52 

53 json_response = response.json() 

54 log.info('results from serply search: %s', json_response) 

55 

56 results = sorted(json_response.get('results', []), key=lambda x: x.get('realPosition', 0)) 

57 if filter_list: 

58 results = get_filtered_results(results, filter_list) 

59 return [ 

60 SearchResult( 

61 link=result['link'], 

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

63 snippet=result.get('description'), 

64 ) 

65 for result in results[:count] 

66 ]