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

1import base64 

2import io 

3import logging 

4import os 

5from typing import List, Optional 

6from xml.etree import ElementTree as ET 

7from xml.etree.ElementTree import Element 

8 

9import requests 

10from fastapi import Request 

11from open_webui.env import FORWARD_SESSION_INFO_HEADER_CHAT_ID 

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

13from open_webui.utils.headers import include_user_info_headers 

14from open_webui.utils.json_codec import JSONCodec 

15 

16log = logging.getLogger(__name__) 

17 

18 

19def xml_element_contents_to_string(element: Element) -> str: 

20 if element is None: 

21 return '' 

22 buffer = [element.text if element.text else ''] 

23 

24 for child in element: 

25 buffer.append(xml_element_contents_to_string(child)) 

26 

27 buffer.append(element.tail if element.tail else '') 

28 

29 return ''.join(buffer) 

30 

31 

32def search_yandex( 

33 request: Request, 

34 yandex_search_url: str, 

35 yandex_search_api_key: str, 

36 yandex_search_config: str, 

37 query: str, 

38 count: int, 

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

40 user=None, 

41) -> List[SearchResult]: 

42 try: 

43 headers = { 

44 # LICENSE covers this Open WebUI user-agent identifier. 

45 # Do not alter, remove, obscure, or replace it except as LICENSE permits: 

46 # https://docs.openwebui.com/license. 

47 'User-Agent': 'Open WebUI (https://github.com/open-webui/open-webui) RAG Bot', 

48 'Authorization': f'Api-Key {yandex_search_api_key}', 

49 } 

50 

51 if user is not None: 

52 headers = include_user_info_headers(headers, user) 

53 

54 chat_id = getattr(request.state, 'chat_id', None) 

55 if chat_id: 

56 headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = str(chat_id) 

57 

58 payload = {} if yandex_search_config == '' else JSONCodec.loads(yandex_search_config) 

59 

60 if type(payload.get('query', None)) != dict: 

61 payload['query'] = {} 

62 

63 if 'searchType' not in payload['query']: 

64 payload['query']['searchType'] = 'SEARCH_TYPE_RU' 

65 

66 payload['query']['queryText'] = query 

67 

68 if type(payload.get('groupSpec', None)) != dict: 

69 payload['groupSpec'] = {} 

70 

71 if 'groupMode' not in payload['groupSpec']: 

72 payload['groupSpec']['groupMode'] = 'GROUP_MODE_DEEP' 

73 

74 payload['groupSpec']['groupsOnPage'] = count 

75 payload['groupSpec']['docsInGroup'] = 1 

76 

77 response = requests.post( 

78 ('https://searchapi.api.cloud.yandex.net/v2/web/search' if yandex_search_url == '' else yandex_search_url), 

79 headers=headers, 

80 json=payload, 

81 ) 

82 

83 response.raise_for_status() 

84 

85 response_body = response.json() 

86 if 'rawData' not in response_body: 

87 raise Exception(f'No `rawData` in response body: {response_body}') 

88 

89 search_result_body_bytes = base64.decodebytes(bytes(response_body['rawData'], 'utf-8')) 

90 

91 doc_root = ET.parse(io.BytesIO(search_result_body_bytes)) 

92 

93 results = [] 

94 

95 for group in doc_root.findall('response/results/grouping/group'): 

96 results.append( 

97 { 

98 'url': xml_element_contents_to_string(group.find('doc/url')).strip('\n'), 

99 'title': xml_element_contents_to_string(group.find('doc/title')).strip('\n'), 

100 'snippet': xml_element_contents_to_string(group.find('doc/passages/passage')), 

101 } 

102 ) 

103 

104 results = get_filtered_results(results, filter_list) 

105 

106 results = [ 

107 SearchResult( 

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

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

110 snippet=result.get('snippet'), 

111 ) 

112 for result in results[:count] 

113 ] 

114 

115 log.info('Yandex search results: %s', results) 

116 

117 return results 

118 except Exception as e: 

119 log.error(f'Error in search: {e}') 

120 

121 return [] 

122 

123 

124if __name__ == '__main__': 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true

125 from fastapi import FastAPI 

126 from starlette.datastructures import Headers 

127 

128 result = search_yandex( 

129 Request( 

130 { 

131 'type': 'http', 

132 'asgi.version': '3.0', 

133 'asgi.spec_version': '2.0', 

134 'method': 'GET', 

135 'path': '/internal', 

136 'query_string': b'', 

137 'headers': Headers({}).raw, 

138 'client': ('127.0.0.1', 12345), 

139 'server': ('127.0.0.1', 80), 

140 'scheme': 'http', 

141 'app': FastAPI(), 

142 }, 

143 None, 

144 ), 

145 os.environ.get('YANDEX_WEB_SEARCH_URL', ''), 

146 os.environ.get('YANDEX_WEB_SEARCH_API_KEY', ''), 

147 os.environ.get('YANDEX_WEB_SEARCH_CONFIG', '{"query": {"searchType": "SEARCH_TYPE_COM"}}'), 

148 'TOP movies of the past year', 

149 3, 

150 ) 

151 

152 print(result)