Coverage for haystack/components/retrievers/in_memory/embedding_retriever.py: 98%
53 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from typing import Any
7from haystack import Document, component, default_from_dict, default_to_dict
8from haystack.document_stores.in_memory import InMemoryDocumentStore
9from haystack.document_stores.types import FilterPolicy, apply_filter_policy
12@component
13class InMemoryEmbeddingRetriever:
14 """
15 Retrieves documents that are most semantically similar to the query.
17 Use this retriever with the InMemoryDocumentStore.
19 When using this retriever, make sure it has query and document embeddings available.
20 In indexing pipelines, use a DocumentEmbedder to embed documents.
21 In query pipelines, use a TextEmbedder to embed queries and send them to the retriever.
23 ### Usage example
24 ```python
25 from haystack import Document
26 from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
27 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
28 from haystack.document_stores.in_memory import InMemoryDocumentStore
30 docs = [
31 Document(content="Python is a popular programming language"),
32 Document(content="python ist eine beliebte Programmiersprache"),
33 ]
34 doc_embedder = OpenAIDocumentEmbedder()
35 docs_with_embeddings = doc_embedder.run(docs)["documents"]
37 doc_store = InMemoryDocumentStore()
38 doc_store.write_documents(docs_with_embeddings)
39 retriever = InMemoryEmbeddingRetriever(doc_store)
41 query="Programmiersprache"
42 text_embedder = OpenAITextEmbedder()
43 query_embedding = text_embedder.run(query)["embedding"]
45 result = retriever.run(query_embedding=query_embedding)
47 print(result["documents"])
48 ```
49 """
51 def __init__(
52 self,
53 document_store: InMemoryDocumentStore,
54 filters: dict[str, Any] | None = None,
55 top_k: int = 10,
56 scale_score: bool = False,
57 return_embedding: bool = False,
58 filter_policy: FilterPolicy = FilterPolicy.REPLACE,
59 ) -> None:
60 """
61 Create the InMemoryEmbeddingRetriever component.
63 :param document_store:
64 An instance of InMemoryDocumentStore where the retriever should search for relevant documents.
65 :param filters:
66 A dictionary with filters to narrow down the retriever's search space in the document store.
67 :param top_k:
68 The maximum number of documents to retrieve. Must be greater than 0.
69 :param scale_score:
70 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
71 When `False`, uses raw similarity scores.
72 :param return_embedding:
73 When `True`, returns the embedding of the retrieved documents.
74 When `False`, returns just the documents, without their embeddings.
75 :param filter_policy: The filter policy to apply during retrieval.
76 Filter policy determines how filters are applied when retrieving documents. You can choose:
77 - `REPLACE` (default): Overrides the initialization filters with the filters specified at runtime.
78 Use this policy to dynamically change filtering for specific queries.
79 - `MERGE`: Combines runtime filters with initialization filters to narrow down the search.
80 :raises TypeError: If the document_store is not an instance of InMemoryDocumentStore.
81 :raises ValueError:
82 If `top_k` is not greater than 0.
83 """
84 if not isinstance(document_store, InMemoryDocumentStore):
85 raise TypeError("document_store must be an instance of InMemoryDocumentStore")
87 self.document_store = document_store
89 if top_k <= 0:
90 raise ValueError(f"top_k must be greater than 0, but got {top_k}")
92 self.filters = filters
93 self.top_k = top_k
94 self.scale_score = scale_score
95 self.return_embedding = return_embedding
96 self.filter_policy = filter_policy
98 def _get_telemetry_data(self) -> dict[str, Any]:
99 """
100 Data that is sent to Posthog for usage analytics.
101 """
102 return {"document_store": type(self.document_store).__name__}
104 def to_dict(self) -> dict[str, Any]:
105 """
106 Serializes the component to a dictionary.
108 :returns:
109 Dictionary with serialized data.
110 """
111 return default_to_dict(
112 self,
113 document_store=self.document_store,
114 filters=self.filters,
115 top_k=self.top_k,
116 scale_score=self.scale_score,
117 return_embedding=self.return_embedding,
118 filter_policy=self.filter_policy.value,
119 )
121 @classmethod
122 def from_dict(cls, data: dict[str, Any]) -> "InMemoryEmbeddingRetriever":
123 """
124 Deserializes the component from a dictionary.
126 :param data:
127 The dictionary to deserialize from.
128 :returns:
129 The deserialized component.
130 """
131 init_params = data.get("init_parameters", {})
132 if "filter_policy" in init_params:
133 init_params["filter_policy"] = FilterPolicy.from_str(init_params["filter_policy"])
134 return default_from_dict(cls, data)
136 @component.output_types(documents=list[Document])
137 def run(
138 self,
139 query_embedding: list[float],
140 filters: dict[str, Any] | None = None,
141 top_k: int | None = None,
142 scale_score: bool | None = None,
143 return_embedding: bool | None = None,
144 ) -> dict[str, list[Document]]:
145 """
146 Run the InMemoryEmbeddingRetriever on the given input data.
148 :param query_embedding:
149 Embedding of the query.
150 :param filters:
151 A dictionary with filters to narrow down the search space when retrieving documents.
152 :param top_k:
153 The maximum number of documents to return. If 0, no documents are returned.
154 :param scale_score:
155 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
156 When `False`, uses raw similarity scores.
157 :param return_embedding:
158 When `True`, returns the embedding of the retrieved documents.
159 When `False`, returns just the documents, without their embeddings.
160 :returns:
161 The retrieved documents.
163 :raises ValueError:
164 If the specified DocumentStore is not found or is not an InMemoryDocumentStore instance,
165 or if `top_k` is negative.
166 """
167 filters = apply_filter_policy(self.filter_policy, self.filters, filters)
168 if top_k is None:
169 top_k = self.top_k
170 if top_k < 0:
171 raise ValueError(f"top_k must be greater than or equal to 0, but got {top_k}")
172 if scale_score is None:
173 scale_score = self.scale_score
174 if return_embedding is None:
175 return_embedding = self.return_embedding
177 docs = self.document_store.embedding_retrieval(
178 query_embedding=query_embedding,
179 filters=filters,
180 top_k=top_k,
181 scale_score=scale_score,
182 return_embedding=return_embedding,
183 )
185 return {"documents": docs}
187 @component.output_types(documents=list[Document])
188 async def run_async(
189 self,
190 query_embedding: list[float],
191 filters: dict[str, Any] | None = None,
192 top_k: int | None = None,
193 scale_score: bool | None = None,
194 return_embedding: bool | None = None,
195 ) -> dict[str, list[Document]]:
196 """
197 Run the InMemoryEmbeddingRetriever on the given input data.
199 :param query_embedding:
200 Embedding of the query.
201 :param filters:
202 A dictionary with filters to narrow down the search space when retrieving documents.
203 :param top_k:
204 The maximum number of documents to return. If 0, no documents are returned.
205 :param scale_score:
206 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
207 When `False`, uses raw similarity scores.
208 :param return_embedding:
209 When `True`, returns the embedding of the retrieved documents.
210 When `False`, returns just the documents, without their embeddings.
211 :returns:
212 The retrieved documents.
214 :raises ValueError:
215 If the specified DocumentStore is not found or is not an InMemoryDocumentStore instance,
216 or if `top_k` is negative.
217 """
218 filters = apply_filter_policy(self.filter_policy, self.filters, filters)
219 if top_k is None:
220 top_k = self.top_k
221 if top_k < 0:
222 raise ValueError(f"top_k must be greater than or equal to 0, but got {top_k}")
223 if scale_score is None:
224 scale_score = self.scale_score
225 if return_embedding is None:
226 return_embedding = self.return_embedding
228 docs = await self.document_store.embedding_retrieval_async(
229 query_embedding=query_embedding,
230 filters=filters,
231 top_k=top_k,
232 scale_score=scale_score,
233 return_embedding=return_embedding,
234 )
236 return {"documents": docs}