Coverage for haystack/components/retrievers/sentence_window_retriever.py: 95%
103 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, logging
8from haystack.document_stores.types import DocumentStore
10logger = logging.getLogger(__name__)
13@component
14class SentenceWindowRetriever:
15 """
16 Retrieves neighboring documents from a DocumentStore to provide context for query results.
18 This component is intended to be used after a Retriever (e.g., BM25Retriever, EmbeddingRetriever).
19 It enhances retrieved results by fetching adjacent document chunks to give
20 additional context for the user.
22 The documents must include metadata indicating their origin and position:
23 - `source_id` is used to group sentence chunks belonging to the same original document.
24 - `split_id` represents the position/order of the chunk within the document.
26 The number of adjacent documents to include on each side of the retrieved document can be configured using the
27 `window_size` parameter. You can also specify which metadata fields to use for source and split ID
28 via `source_id_meta_field` and `split_id_meta_field`.
30 The SentenceWindowRetriever is compatible with the following DocumentStores:
31 - [Astra](https://docs.haystack.deepset.ai/docs/astradocumentstore)
32 - [Elasticsearch](https://docs.haystack.deepset.ai/docs/elasticsearch-document-store)
33 - [OpenSearch](https://docs.haystack.deepset.ai/docs/opensearch-document-store)
34 - [Pgvector](https://docs.haystack.deepset.ai/docs/pgvectordocumentstore)
35 - [Pinecone](https://docs.haystack.deepset.ai/docs/pinecone-document-store)
36 - [Qdrant](https://docs.haystack.deepset.ai/docs/qdrant-document-store)
38 ### Usage example
40 ```python
41 from haystack import Document, Pipeline
42 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
43 from haystack.components.retrievers import SentenceWindowRetriever
44 from haystack.components.preprocessors import DocumentSplitter
45 from haystack.document_stores.in_memory import InMemoryDocumentStore
47 splitter = DocumentSplitter(split_length=10, split_overlap=5, split_by="word")
48 text = (
49 "This is a text with some words. There is a second sentence. And there is also a third sentence. "
50 "It also contains a fourth sentence. And a fifth sentence. And a sixth sentence. And a seventh sentence"
51 )
52 doc = Document(content=text)
53 docs = splitter.run([doc])
54 doc_store = InMemoryDocumentStore()
55 doc_store.write_documents(docs["documents"])
58 rag = Pipeline()
59 rag.add_component("bm25_retriever", InMemoryBM25Retriever(doc_store, top_k=1))
60 rag.add_component("sentence_window_retriever", SentenceWindowRetriever(document_store=doc_store, window_size=2))
61 rag.connect("bm25_retriever", "sentence_window_retriever")
63 rag.run({'bm25_retriever': {"query":"third"}})
65 # >> {'sentence_window_retriever': {'context_windows': ['some words. There is a second sentence.
66 # >> And there is also a third sentence. It also contains a fourth sentence. And a fifth sentence. And a sixth
67 # >> sentence. And a'], 'context_documents': [[Document(id=..., content: 'some words. There is a second sentence.
68 # >> And there is ', meta: {'source_id': '...', 'page_number': 1, 'split_id': 1, 'split_idx_start': 20,
69 # >> '_split_overlap': [{'doc_id': '...', 'range': (20, 43)}, {'doc_id': '...', 'range': (0, 30)}]}),
70 # >> Document(id=..., content: 'second sentence. And there is also a third sentence. It ',
71 # >> meta: {'source_id': '74ea87deb38012873cf8c07e...f19d01a26a098447113e1d7b83efd30c02987114', 'page_number': 1,
72 # >> 'split_id': 2, 'split_idx_start': 43, '_split_overlap': [{'doc_id': '...', 'range': (23, 53)}, {'doc_id': '.',
73 # >> 'range': (0, 26)}]}), Document(id=..., content: 'also a third sentence. It also contains a fourth sentence. ',
74 # >> meta: {'source_id': '...', 'page_number': 1, 'split_id': 3, 'split_idx_start': 73, '_split_overlap':
75 # >> [{'doc_id': '...', 'range': (30, 56)}, {'doc_id': '...', 'range': (0, 33)}]}), Document(id=..., content:
76 # >> 'also contains a fourth sentence. And a fifth sentence. And ', meta: {'source_id': '...', 'page_number': 1,
77 # >> 'split_id': 4, 'split_idx_start': 99, '_split_overlap': [{'doc_id': '...', 'range': (26, 59)},
78 # >> {'doc_id': '...', 'range': (0, 26)}]}), Document(id=..., content: 'And a fifth sentence. And a sixth sentence.
79 # >> And a ', meta: {'source_id': '...', 'page_number': 1, 'split_id': 5, 'split_idx_start': 132,
80 # >> '_split_overlap': [{'doc_id': '...', 'range': (33, 59)}, {'doc_id': '...', 'range': (0, 24)}]})]]}}}}
81 ```
82 """
84 def __init__(
85 self,
86 document_store: DocumentStore,
87 window_size: int = 3,
88 *,
89 source_id_meta_field: str | list[str] = "source_id",
90 split_id_meta_field: str = "split_id",
91 raise_on_missing_meta_fields: bool = True,
92 ) -> None:
93 """
94 Creates a new SentenceWindowRetriever component.
96 :param document_store: The Document Store to retrieve the surrounding documents from.
97 :param window_size: The number of documents to retrieve before and after the relevant one.
98 For example, `window_size: 2` fetches 2 preceding and 2 following documents.
99 :param source_id_meta_field: The metadata field that contains the source ID of the document.
100 This can be a single field or a list of fields. If multiple fields are provided, the retriever will
101 consider the document as part of the same source if all the fields match.
102 :param split_id_meta_field: The metadata field that contains the split ID of the document.
103 :param raise_on_missing_meta_fields: If True, raises an error if the documents do not contain the required
104 metadata fields. If False, it will skip retrieving the context for documents that are missing
105 the required metadata fields, but will still include the original document in the results.
106 """
107 self._validate_window_size(window_size)
109 self.window_size = window_size
110 self.document_store = document_store
111 self.source_id_meta_field = source_id_meta_field
112 # Use this to have an attribute that is always a list of source id meta fields.
113 self._source_id_meta_fields = (
114 source_id_meta_field if isinstance(source_id_meta_field, list) else [source_id_meta_field]
115 )
116 self.split_id_meta_field = split_id_meta_field
117 self.raise_on_missing_meta_fields = raise_on_missing_meta_fields
119 @staticmethod
120 def merge_documents_text(documents: list[Document]) -> str:
121 """
122 Merge a list of document text into a single string.
124 This functions concatenates the textual content of a list of documents into a single string, eliminating any
125 overlapping content.
127 :param documents: List of Documents to merge.
128 """
129 if any("split_idx_start" not in doc.meta for doc in documents):
130 # If any of the documents is missing the 'split_idx_start' metadata we just concatenate their content.
131 return "".join(doc.content for doc in documents if doc.content)
133 sorted_docs = sorted(documents, key=lambda doc: doc.meta["split_idx_start"])
134 merged_text = ""
135 last_idx_end = 0
136 for doc in sorted_docs:
137 if doc.content is None:
138 continue
140 start = doc.meta.get("split_idx_start", 0) # start of the current content
142 # if the start of the current content is before the end of the last appended content, adjust it
143 start = max(start, last_idx_end)
145 # append the non-overlapping part to the merged text
146 merged_text += doc.content[start - int(doc.meta["split_idx_start"]) :]
148 # update the last end index
149 last_idx_end = int(doc.meta["split_idx_start"]) + len(doc.content)
151 return merged_text
153 def to_dict(self) -> dict[str, Any]:
154 """
155 Serializes the component to a dictionary.
157 :returns:
158 Dictionary with serialized data.
159 """
160 return default_to_dict(
161 self,
162 document_store=self.document_store,
163 window_size=self.window_size,
164 source_id_meta_field=self.source_id_meta_field,
165 split_id_meta_field=self.split_id_meta_field,
166 raise_on_missing_meta_fields=self.raise_on_missing_meta_fields,
167 )
169 @classmethod
170 def from_dict(cls, data: dict[str, Any]) -> "SentenceWindowRetriever":
171 """
172 Deserializes the component from a dictionary.
174 :returns:
175 Deserialized component.
176 """
177 return default_from_dict(cls, data)
179 @component.output_types(context_windows=list[str], context_documents=list[Document])
180 def run(self, retrieved_documents: list[Document], window_size: int | None = None) -> dict[str, Any]:
181 """
182 Based on the `source_id` and on the `doc.meta['split_id']` get surrounding documents from the document store.
184 Implements the logic behind the sentence-window technique, retrieving the surrounding documents of a given
185 document from the document store.
187 :param retrieved_documents: List of retrieved documents from the previous retriever.
188 :param window_size: The number of documents to retrieve before and after the relevant one. This will overwrite
189 the `window_size` parameter set in the constructor. It must be greater than 0; values of
190 0 and negative values are rejected.
191 :returns:
192 A dictionary with the following keys:
193 - `context_windows`: A list of strings, where each string represents the concatenated text from the
194 context window of the corresponding document in `retrieved_documents`.
195 - `context_documents`: A list `Document` objects, containing the retrieved documents plus the context
196 document surrounding them. The documents are sorted by the `split_idx_start`
197 meta field.
199 """
200 window_size = self.window_size if window_size is None else window_size
201 self._validate_window_size(window_size)
202 self._raise_if_documents_do_not_have_expected_metadata(retrieved_documents)
204 context_text = []
205 context_documents = []
206 for doc in retrieved_documents:
207 text, docs = self._retrieve_context_for_document(doc, window_size)
208 context_text.append(text)
209 context_documents.extend(docs)
211 return {"context_windows": context_text, "context_documents": context_documents}
213 @component.output_types(context_windows=list[str], context_documents=list[Document])
214 async def run_async(self, retrieved_documents: list[Document], window_size: int | None = None) -> dict[str, Any]:
215 """
216 Based on the `source_id` and on the `doc.meta['split_id']` get surrounding documents from the document store.
218 Implements the logic behind the sentence-window technique, retrieving the surrounding documents of a given
219 document from the document store.
221 :param retrieved_documents: List of retrieved documents from the previous retriever.
222 :param window_size: The number of documents to retrieve before and after the relevant one. This will overwrite
223 the `window_size` parameter set in the constructor. It must be greater than 0; values of
224 0 and negative values are rejected.
225 :returns:
226 A dictionary with the following keys:
227 - `context_windows`: A list of strings, where each string represents the concatenated text from the
228 context window of the corresponding document in `retrieved_documents`.
229 - `context_documents`: A list `Document` objects, containing the retrieved documents plus the context
230 document surrounding them. The documents are sorted by the `split_idx_start`
231 meta field.
233 """
234 window_size = self.window_size if window_size is None else window_size
235 self._validate_window_size(window_size)
236 self._raise_if_documents_do_not_have_expected_metadata(retrieved_documents)
238 context_text = []
239 context_documents = []
240 for doc in retrieved_documents:
241 text, docs = await self._retrieve_context_for_document_async(doc, window_size)
242 context_text.append(text)
243 context_documents.extend(docs)
245 return {"context_windows": context_text, "context_documents": context_documents}
247 @staticmethod
248 def _validate_window_size(window_size: int) -> None:
249 if window_size < 1:
250 raise ValueError("The window_size parameter must be greater than 0.")
252 def _raise_if_documents_do_not_have_expected_metadata(self, retrieved_documents: list[Document]) -> None:
253 if (
254 not all(self.split_id_meta_field in doc.meta for doc in retrieved_documents)
255 and self.raise_on_missing_meta_fields
256 ):
257 raise ValueError(f"The retrieved documents must have '{self.split_id_meta_field}' in their metadata.")
259 if (
260 not all(field in doc.meta for doc in retrieved_documents for field in self._source_id_meta_fields)
261 and self.raise_on_missing_meta_fields
262 ):
263 raise ValueError(f"The retrieved documents must have '{self.source_id_meta_field}' in their metadata.")
265 def _retrieve_context_for_document(self, doc: Document, window_size: int) -> tuple[str, list[Document]]:
266 source_ids = [doc.meta.get(field) for field in self._source_id_meta_fields]
267 split_id = doc.meta.get(self.split_id_meta_field)
269 if any(source_id is None for source_id in source_ids) or split_id is None:
270 logger.warning(
271 "Document {doc_id} is missing required metadata fields to be used with "
272 "SentenceWindowRetriever: {source_id} or {split_id}. Skipping context retrieval for this document.",
273 doc_id=doc.id,
274 source_id=self._source_id_meta_fields,
275 split_id=self.split_id_meta_field,
276 )
277 return doc.content or "", [doc]
279 assert split_id is not None
280 filter_conditions = self._build_filter_conditions(split_id, window_size, source_ids)
281 context_docs = self.document_store.filter_documents(filter_conditions)
282 context_text = self.merge_documents_text(context_docs)
283 context_docs_sorted = sorted(context_docs, key=lambda doc: doc.meta[self.split_id_meta_field])
285 return context_text, context_docs_sorted
287 async def _retrieve_context_for_document_async(self, doc: Document, window_size: int) -> tuple[str, list[Document]]:
288 source_ids = [doc.meta.get(field) for field in self._source_id_meta_fields]
289 split_id = doc.meta.get(self.split_id_meta_field)
291 if any(source_id is None for source_id in source_ids) or split_id is None:
292 logger.warning(
293 "Document {doc_id} is missing required metadata fields to be used with "
294 "SentenceWindowRetriever: {source_id} or {split_id}. Skipping context retrieval for this document.",
295 doc_id=doc.id,
296 source_id=self._source_id_meta_fields,
297 split_id=self.split_id_meta_field,
298 )
299 return doc.content or "", [doc]
301 assert split_id is not None
302 filter_conditions = self._build_filter_conditions(split_id, window_size, source_ids)
303 # Ignoring type error because DocumentStore protocol doesn't define filter_documents_async
304 context_docs = await self.document_store.filter_documents_async(filter_conditions) # type: ignore[attr-defined]
305 context_text = self.merge_documents_text(context_docs)
306 context_docs_sorted = sorted(context_docs, key=lambda doc: doc.meta[self.split_id_meta_field])
308 return context_text, context_docs_sorted
310 def _build_filter_conditions(self, split_id: int, window_size: int, source_ids: list[Any]) -> dict[str, Any]:
311 min_before = split_id - window_size
312 max_after = split_id + window_size
313 source_id_filters = [
314 {"field": f"meta.{source_id_meta_field}", "operator": "==", "value": source_id}
315 for source_id_meta_field, source_id in zip(self._source_id_meta_fields, source_ids, strict=True)
316 ]
317 conditions = [
318 {"field": f"meta.{self.split_id_meta_field}", "operator": ">=", "value": min_before},
319 {"field": f"meta.{self.split_id_meta_field}", "operator": "<=", "value": max_after},
320 *source_id_filters,
321 ]
322 return {"operator": "AND", "conditions": conditions}
324 def close(self) -> None:
325 """
326 Release the synchronous resources of the underlying Document Store.
327 """
328 if hasattr(self.document_store, "close"):
329 self.document_store.close()
331 async def close_async(self) -> None:
332 """
333 Release the asynchronous resources of the underlying Document Store.
334 """
335 if hasattr(self.document_store, "close_async"):
336 await self.document_store.close_async()