Coverage for haystack/components/retrievers/filter_retriever.py: 97%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from haystack import Document, component, default_from_dict, default_to_dict 

8from haystack.document_stores.types import DocumentStore 

9 

10 

11@component 

12class FilterRetriever: 

13 """ 

14 Retrieves documents that match the provided filters. 

15 

16 ### Usage example 

17 

18 ```python 

19 from haystack import Document 

20 from haystack.components.retrievers import FilterRetriever 

21 from haystack.document_stores.in_memory import InMemoryDocumentStore 

22 

23 docs = [ 

24 Document(content="Python is a popular programming language", meta={"lang": "en"}), 

25 Document(content="python ist eine beliebte Programmiersprache", meta={"lang": "de"}), 

26 ] 

27 

28 doc_store = InMemoryDocumentStore() 

29 doc_store.write_documents(docs) 

30 retriever = FilterRetriever(doc_store, filters={"field": "lang", "operator": "==", "value": "en"}) 

31 

32 # if passed in the run method, filters override those provided at initialization 

33 result = retriever.run(filters={"field": "lang", "operator": "==", "value": "de"}) 

34 

35 print(result["documents"]) 

36 ``` 

37 """ 

38 

39 def __init__(self, document_store: DocumentStore, filters: dict[str, Any] | None = None) -> None: 

40 """ 

41 Create the FilterRetriever component. 

42 

43 :param document_store: 

44 An instance of a Document Store to use with the Retriever. 

45 :param filters: 

46 A dictionary with filters to narrow down the search space. 

47 """ 

48 self.document_store = document_store 

49 self.filters = filters 

50 

51 def _get_telemetry_data(self) -> dict[str, Any]: 

52 """ 

53 Data that is sent to Posthog for usage analytics. 

54 """ 

55 return {"document_store": type(self.document_store).__name__} 

56 

57 def to_dict(self) -> dict[str, Any]: 

58 """ 

59 Serializes the component to a dictionary. 

60 

61 :returns: 

62 Dictionary with serialized data. 

63 """ 

64 return default_to_dict(self, document_store=self.document_store, filters=self.filters) 

65 

66 @classmethod 

67 def from_dict(cls, data: dict[str, Any]) -> "FilterRetriever": 

68 """ 

69 Deserializes the component from a dictionary. 

70 

71 :param data: 

72 The dictionary to deserialize from. 

73 :returns: 

74 The deserialized component. 

75 """ 

76 return default_from_dict(cls, data) 

77 

78 @component.output_types(documents=list[Document]) 

79 def run(self, filters: dict[str, Any] | None = None) -> dict[str, Any]: 

80 """ 

81 Run the FilterRetriever on the given input data. 

82 

83 :param filters: 

84 A dictionary with filters to narrow down the search space. 

85 If not specified, the FilterRetriever uses the values provided at initialization. 

86 :returns: 

87 A list of retrieved documents. 

88 """ 

89 resolved_filters = filters if filters is not None else self.filters 

90 return {"documents": self.document_store.filter_documents(filters=resolved_filters)} 

91 

92 @component.output_types(documents=list[Document]) 

93 async def run_async(self, filters: dict[str, Any] | None = None) -> dict[str, Any]: 

94 """ 

95 Asynchronously run the FilterRetriever on the given input data. 

96 

97 :param filters: 

98 A dictionary with filters to narrow down the search space. 

99 If not specified, the FilterRetriever uses the values provided at initialization. 

100 :returns: 

101 A list of retrieved documents. 

102 """ 

103 # 'ignore' since filter_documents_async is not defined in the Protocol but exists in the implementations 

104 resolved_filters = filters if filters is not None else self.filters 

105 out_documents = await self.document_store.filter_documents_async( # type: ignore[attr-defined] 

106 filters=resolved_filters 

107 ) 

108 return {"documents": out_documents} 

109 

110 def close(self) -> None: 

111 """ 

112 Release the synchronous resources of the underlying Document Store. 

113 """ 

114 if hasattr(self.document_store, "close"): 

115 self.document_store.close() 

116 

117 async def close_async(self) -> None: 

118 """ 

119 Release the asynchronous resources of the underlying Document Store. 

120 """ 

121 if hasattr(self.document_store, "close_async"): 

122 await self.document_store.close_async()