Coverage for haystack/components/retrievers/in_memory/embedding_retriever.py: 98%
49 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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.
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 the specified top_k is not > 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. Currently, top_k is {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.
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 """
166 filters = apply_filter_policy(self.filter_policy, self.filters, filters)
167 if top_k is None:
168 top_k = self.top_k
169 if scale_score is None:
170 scale_score = self.scale_score
171 if return_embedding is None:
172 return_embedding = self.return_embedding
174 docs = self.document_store.embedding_retrieval(
175 query_embedding=query_embedding,
176 filters=filters,
177 top_k=top_k,
178 scale_score=scale_score,
179 return_embedding=return_embedding,
180 )
182 return {"documents": docs}
184 @component.output_types(documents=list[Document])
185 async def run_async(
186 self,
187 query_embedding: list[float],
188 filters: dict[str, Any] | None = None,
189 top_k: int | None = None,
190 scale_score: bool | None = None,
191 return_embedding: bool | None = None,
192 ) -> dict[str, list[Document]]:
193 """
194 Run the InMemoryEmbeddingRetriever on the given input data.
196 :param query_embedding:
197 Embedding of the query.
198 :param filters:
199 A dictionary with filters to narrow down the search space when retrieving documents.
200 :param top_k:
201 The maximum number of documents to return.
202 :param scale_score:
203 When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
204 When `False`, uses raw similarity scores.
205 :param return_embedding:
206 When `True`, returns the embedding of the retrieved documents.
207 When `False`, returns just the documents, without their embeddings.
208 :returns:
209 The retrieved documents.
211 :raises ValueError:
212 If the specified DocumentStore is not found or is not an InMemoryDocumentStore instance.
213 """
214 filters = apply_filter_policy(self.filter_policy, self.filters, filters)
215 if top_k is None:
216 top_k = self.top_k
217 if scale_score is None:
218 scale_score = self.scale_score
219 if return_embedding is None:
220 return_embedding = self.return_embedding
222 docs = await self.document_store.embedding_retrieval_async(
223 query_embedding=query_embedding,
224 filters=filters,
225 top_k=top_k,
226 scale_score=scale_score,
227 return_embedding=return_embedding,
228 )
230 return {"documents": docs}