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

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 loguru 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 self._lock = Lock() 

23 

24 def register(self, name: str, retriever: BaseRetriever) -> None: 

25 """ 

26 Register a retriever with a given name. 

27 

28 Args: 

29 name: Name to register the retriever under 

30 retriever: LangChain BaseRetriever instance 

31 """ 

32 with self._lock: 

33 self._retrievers[name] = retriever 

34 logger.info( 

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

36 ) 

37 

38 def register_multiple(self, retrievers: Dict[str, BaseRetriever]) -> None: 

39 """ 

40 Register multiple retrievers at once. 

41 

42 Args: 

43 retrievers: Dictionary of {name: retriever} pairs 

44 """ 

45 with self._lock: 

46 for name, retriever in retrievers.items(): 

47 self._retrievers[name] = retriever 

48 logger.info( 

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

50 ) 

51 

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

53 """ 

54 Get a registered retriever by name. 

55 

56 Args: 

57 name: Name of the retriever 

58 

59 Returns: 

60 The retriever if found, None otherwise 

61 """ 

62 with self._lock: 

63 return self._retrievers.get(name) 

64 

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

66 """ 

67 Remove a registered retriever. 

68 

69 Args: 

70 name: Name of the retriever to remove 

71 """ 

72 with self._lock: 

73 if name in self._retrievers: 73 ↛ exitline 73 didn't jump to the function exit

74 del self._retrievers[name] 

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

76 

77 def clear(self) -> None: 

78 """Clear all registered retrievers.""" 

79 with self._lock: 

80 count = len(self._retrievers) 

81 self._retrievers.clear() 

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

83 

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

85 """ 

86 Check if a retriever is registered. 

87 

88 Args: 

89 name: Name of the retriever 

90 

91 Returns: 

92 True if registered, False otherwise 

93 """ 

94 with self._lock: 

95 return name in self._retrievers 

96 

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

98 """ 

99 Get list of all registered retriever names. 

100 

101 Returns: 

102 List of retriever names 

103 """ 

104 with self._lock: 

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

106 

107 

108# Global registry instance 

109retriever_registry = RetrieverRegistry()