Coverage for open_webui/retrieval/web/sougou.py: 21%

1import logging 

2from typing import List, Optional 

3 

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

5from open_webui.utils.json_codec import JSONCodec 

6 

7log = logging.getLogger(__name__) 

8 

9 

10def search_sougou( 

11 sougou_api_sid: str, 

12 sougou_api_sk: str, 

13 query: str, 

14 count: int, 

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

16) -> List[SearchResult]: 

17 from tencentcloud.common import credential 

18 from tencentcloud.common.common_client import CommonClient 

19 from tencentcloud.common.exception.tencent_cloud_sdk_exception import ( 

20 TencentCloudSDKException, 

21 ) 

22 from tencentcloud.common.profile.client_profile import ClientProfile 

23 from tencentcloud.common.profile.http_profile import HttpProfile 

24 

25 try: 

26 cred = credential.Credential(sougou_api_sid, sougou_api_sk) 

27 http_profile = HttpProfile() 

28 http_profile.endpoint = 'tms.tencentcloudapi.com' 

29 client_profile = ClientProfile() 

30 client_profile.http_profile = http_profile 

31 params = JSONCodec.dumps({'Query': query, 'Cnt': 20}) 

32 common_client = CommonClient('tms', '2020-12-29', cred, '', profile=client_profile) 

33 results = [ 

34 JSONCodec.loads(page) 

35 for page in common_client.call_json('SearchPro', JSONCodec.loads(params))['Response']['Pages'] 

36 ] 

37 sorted_results = sorted(results, key=lambda x: x.get('scour', 0.0), reverse=True) 

38 if filter_list: 

39 sorted_results = get_filtered_results(sorted_results, filter_list) 

40 

41 return [ 

42 SearchResult( 

43 link=result.get('url'), 

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

45 snippet=result.get('passage'), 

46 ) 

47 for result in sorted_results[:count] 

48 ] 

49 except TencentCloudSDKException as err: 

50 log.error(f'Error in Sougou search: {err}') 

51 return []