Coverage for open_webui/retrieval/web/google_pse.py: 16%

1from __future__ import annotations 

2 

3import logging 

4 

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

6from open_webui.utils.session_pool import get_session 

7 

8log = logging.getLogger(__name__) 

9 

10 

11async def search_google_pse( 

12 api_key: str, 

13 search_engine_id: str, 

14 query: str, 

15 count: int, 

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

17 referer: str | None = None, 

18) -> list[SearchResult]: 

19 """Query Google Programmable Search Engine with automatic pagination. 

20 

21 The PSE API returns at most 10 results per request, so this function 

22 issues multiple requests when ``count > 10``. 

23 """ 

24 url = 'https://www.googleapis.com/customsearch/v1' 

25 headers: dict[str, str] = {'Content-Type': 'application/json'} 

26 if referer: 

27 headers['Referer'] = referer 

28 

29 all_items: list[dict] = [] 

30 start_index = 1 # PSE uses 1-based pagination 

31 

32 session = await get_session() 

33 remaining = count 

34 while remaining > 0: 

35 page_size = min(remaining, 10) 

36 params = { 

37 'cx': search_engine_id, 

38 'q': query, 

39 'key': api_key, 

40 'num': str(page_size), 

41 'start': str(start_index), 

42 } 

43 

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

45 response.raise_for_status() 

46 payload = await response.json() 

47 

48 items = payload.get('items', []) 

49 if not items: 

50 break 

51 

52 all_items.extend(items) 

53 remaining -= len(items) 

54 start_index += 10 

55 

56 if filter_list: 

57 all_items = get_filtered_results(all_items, filter_list) 

58 

59 return [ 

60 SearchResult( 

61 link=item.get('link', ''), 

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

63 snippet=item.get('snippet'), 

64 ) 

65 for item in all_items 

66 ]