Coverage for src/local_deep_research/web_search_engines/retriever_registry.py: 91%

1""" 

2Registry for dynamically registering LangChain retrievers as search engines. 

3""" 

4 

5from threading import Lock 

6from typing import Dict, Optional 

7 

8from langchain_core.retrievers import BaseRetriever 

9from ..security.secure_logging import logger 

10 

11 

12class RetrieverRegistry: 

13 """ 

14 Thread-safe registry for LangChain retrievers. 

15 

16 This allows users to register retrievers programmatically and use them 

17 as search engines within LDR. 

18 """ 

19 

20 def __init__(self): 

21 self._retrievers: Dict[str, BaseRetriever] = {} 

22 # Parallel map of name -> classification metadata, e.g. 

23 # {"is_local": True}. Kept separate from _retrievers so get() 

24 # keeps returning the bare retriever object (callers + tests 

25 # depend on `get() is retriever`). 

26 self._metadata: Dict[str, dict] = {} 

27 self._lock = Lock() 

28 

29 def register( 

30 self, 

31 name: str, 

32 retriever: BaseRetriever, 

33 is_local: bool = True, 

34 ) -> None: 

35 """ 

36 Register a retriever with a given name. 

37 

38 Args: 

39 name: Name to register the retriever under 

40 retriever: LangChain BaseRetriever instance 

41 is_local: Whether this retriever reads local/private data 

42 (a user's vector store / KB) rather than the public 

43 internet. Defaults to True because programmatically 

44 registered retrievers are almost always private corpora; 

45 the egress policy uses this to decide whether the 

46 retriever may run under PUBLIC_ONLY / PRIVATE_ONLY scopes. 

47 """ 

48 with self._lock: 

49 self._retrievers[name] = retriever 

50 self._metadata[name] = {"is_local": is_local} 

51 logger.info( 

52 f"Registered retriever '{name}' of type {type(retriever).__name__} " 

53 f"(is_local={is_local})" 

54 ) 

55 

56 def register_multiple( 

57 self, 

58 retrievers: Dict[str, "BaseRetriever | dict"], 

59 is_local: bool = True, 

60 ) -> None: 

61 """ 

62 Register multiple retrievers at once. 

63 

64 Args: 

65 retrievers: Either ``{name: retriever}`` (uses ``is_local`` 

66 arg as a single default) OR a richer 

67 ``{name: {"retriever": …, "is_local": bool}}`` mapping 

68 so different retrievers can carry different 

69 classifications in one call. Mixing both shapes inside 

70 the same dict is allowed. 

71 is_local: Default classification when an entry uses the 

72 bare-retriever shape. The public programmatic API 

73 (api/research_functions.py) calls this; we default to 

74 True here ONLY because the old API contract was True. 

75 CLI / programmatic callers passing a cloud-backed 

76 retriever MUST use the dict shape with 

77 ``is_local=False`` or they will egress under 

78 PRIVATE_ONLY (silent classification lie). 

79 """ 

80 with self._lock: 

81 for name, entry in retrievers.items(): 

82 if isinstance(entry, dict): 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true

83 retriever = entry.get("retriever") 

84 entry_is_local = bool(entry.get("is_local", is_local)) 

85 else: 

86 retriever = entry 

87 entry_is_local = is_local 

88 if retriever is None: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true

89 logger.warning( 

90 f"Retriever entry '{name}' has no retriever — skipped" 

91 ) 

92 continue 

93 self._retrievers[name] = retriever 

94 self._metadata[name] = {"is_local": entry_is_local} 

95 logger.info( 

96 f"Registered retriever '{name}' of type {type(retriever).__name__} " 

97 f"(is_local={entry_is_local})" 

98 ) 

99 

100 def get(self, name: str) -> Optional[BaseRetriever]: 

101 """ 

102 Get a registered retriever by name. 

103 

104 Args: 

105 name: Name of the retriever 

106 

107 Returns: 

108 The retriever if found, None otherwise 

109 """ 

110 with self._lock: 

111 return self._retrievers.get(name) 

112 

113 def unregister(self, name: str) -> None: 

114 """ 

115 Remove a registered retriever. 

116 

117 Args: 

118 name: Name of the retriever to remove 

119 """ 

120 with self._lock: 

121 if name in self._retrievers: 

122 del self._retrievers[name] 

123 self._metadata.pop(name, None) 

124 logger.info(f"Unregistered retriever '{name}'") 

125 

126 def clear(self) -> None: 

127 """Clear all registered retrievers.""" 

128 with self._lock: 

129 count = len(self._retrievers) 

130 self._retrievers.clear() 

131 self._metadata.clear() 

132 logger.info(f"Cleared {count} registered retrievers") 

133 

134 def is_registered(self, name: str) -> bool: 

135 """ 

136 Check if a retriever is registered. 

137 

138 Args: 

139 name: Name of the retriever 

140 

141 Returns: 

142 True if registered, False otherwise 

143 """ 

144 with self._lock: 

145 return name in self._retrievers 

146 

147 def list_registered(self) -> list[str]: 

148 """ 

149 Get list of all registered retriever names. 

150 

151 Returns: 

152 List of retriever names 

153 """ 

154 with self._lock: 

155 return list(self._retrievers.keys()) 

156 

157 def get_metadata( 

158 self, name: str, username: Optional[str] = None 

159 ) -> Optional[Dict]: 

160 """Return policy-relevant metadata for a registered retriever. 

161 

162 Returns a dict like ``{"is_local": True}`` for a registered 

163 retriever, or ``None`` when the retriever is unknown. A 

164 registered retriever with no recorded classification yields an 

165 empty dict, which ``evaluate_retriever`` treats as "unclassified" 

166 and fails closed under any non-BOTH scope. 

167 

168 The egress policy's ``evaluate_retriever`` consults this hook to 

169 decide whether a retriever may run under the active scope. 

170 

171 ``username`` is accepted so future per-user retriever isolation 

172 can plug in without changing the call sites. 

173 """ 

174 with self._lock: 

175 if name not in self._retrievers: 

176 return None 

177 # Copy so callers can't mutate the stored classification. 

178 return dict(self._metadata.get(name, {})) 

179 

180 

181# Global registry instance 

182retriever_registry = RetrieverRegistry()