Coverage for apis_core/utils/autocomplete.py: 0%

1import json 

2import logging 

3 

4import requests 

5from django.template.loader import render_to_string 

6 

7logger = logging.getLogger(__name__) 

8 

9 

10class ExternalAutocomplete: 

11 """ 

12 This is a helper base class for implementing external 

13 autocomplete classes. <Modelname>ExernalAutocomplete classes 

14 are expected to have a `get_results(self, q)` method that 

15 returns a list of results usable by the autocomplete view. 

16 This base class implements this `get_results` method in a 

17 way that you can inherit from it and just define a list of 

18 `adapters`. Those adapters are then used one by one to 

19 add external autocomplete search results. 

20 """ 

21 

22 session = requests.Session() 

23 adapters = [] 

24 

25 def get_results(self, q): 

26 results = [] 

27 for adapter in self.adapters: 

28 results.extend(adapter.get_results(q, self.session)) 

29 return results 

30 

31 

32class ExternalAutocompleteAdapter: 

33 """ 

34 Base class for ExternalAutocompleteAdapters. It provides 

35 the methods used for templating the autocomplete results. 

36 You can pass a `template` name to initialization, which 

37 is then used to style the results. 

38 """ 

39 

40 template = None 

41 

42 def __init__(self, *args, **kwargs): 

43 self.template = kwargs.get("template", None) 

44 

45 def default_template(self, result): 

46 return f'{result["label"]} <a href="{result["id"]}">{result["id"]}</a>' 

47 

48 def get_result_label(self, result): 

49 if self.template: 

50 return render_to_string(self.template, {"result": result}) 

51 return self.default_template(result) 

52 

53 

54class TypeSenseAutocompleteAdapter(ExternalAutocompleteAdapter): 

55 """ 

56 This autocomplete adapters queries typesense collections on a 

57 typesense server. The `collections` variable can either be a 

58 string or a list - if its a string, that collection is queried 

59 directly, if its a list, the adapter uses typesense `multi_search` 

60 endpoint. 

61 """ 

62 

63 collections = None 

64 token = None 

65 server = None 

66 

67 def __init__(self, *args, **kwargs): 

68 super().__init__(*args, **kwargs) 

69 self.collections = kwargs.get("collections", None) 

70 self.token = kwargs.get("token", None) 

71 self.server = kwargs.get("server", None) 

72 

73 def default_template(self, result): 

74 return super().default_template(result["document"]) 

75 

76 def extract(self, res): 

77 if res.get("document"): 

78 return { 

79 "id": res["document"]["id"], 

80 "text": self.get_result_label(res), 

81 "selected_text": self.get_result_label(res), 

82 } 

83 logger.error( 

84 "Could not parse result from typesense collection %s: %s", 

85 self.collections, 

86 res, 

87 ) 

88 return False 

89 

90 def get_results(self, q, session=requests.Session()): 

91 headers = {"X-TYPESENSE-API-KEY": self.token} 

92 res = None 

93 if self.token and self.server: 

94 match self.collections: 

95 # if there is only on collection configured, we hit that collection directly 

96 case str() as collection: 

97 url = f"{self.server}/collections/{collection}/documents/search?q={q}&query_by=description&query_by=label" 

98 res = session.get(url, headers=headers) 

99 # if there are multiple collections configured, we use the `multi_search` endpoint 

100 case list() as collectionlist: 

101 url = f"{self.server}/multi_search?q={q}&query_by=description&query_by=label" 

102 data = {"searches": []} 

103 for collection in collectionlist: 

104 data["searches"].append({"collection": collection}) 

105 res = session.post(url, data=json.dumps(data), headers=headers) 

106 case unknown: 

107 logger.error("Don't know what to do with collection %s", unknown) 

108 

109 if res: 

110 data = res.json() 

111 hits = data.get("hits", []) 

112 for result in data.get("results", []): 

113 hits.extend(result["hits"]) 

114 return list(filter(bool, map(self.extract, hits))) 

115 return [] 

116 

117 

118class LobidAutocompleteAdapter(ExternalAutocompleteAdapter): 

119 """ 

120 This autocomplete adapters queries the lobid autocomplete apis. 

121 See https://lobid.org/gnd/api for details 

122 You can pass a `lobid_params` dict which will then be use as GET 

123 request parameters. 

124 """ 

125 

126 params = {} 

127 

128 def __init__(self, *args, **kwargs): 

129 super().__init__(*args, **kwargs) 

130 self.params = kwargs.get("params", {}) 

131 

132 def extract(self, res): 

133 return { 

134 "id": res["id"], 

135 "text": self.get_result_label(res), 

136 "selected_text": self.get_result_label(res), 

137 } 

138 

139 def get_results(self, q, session=requests.Session()): 

140 endpoint = "https://lobid.org/gnd/search?" 

141 self.params["q"] = q 

142 res = session.get(endpoint, params=self.params) 

143 if res: 

144 return list(filter(bool, map(self.extract, res.json()))) 

145 return []