Coverage for haystack/components/retrievers/sentence_window_retriever.py: 99%
112 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +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 of `Document` objects, containing the retrieved documents plus the
196 context documents surrounding them, grouped by window in the order of
197 `retrieved_documents`. Within each window, the documents are sorted by the
198 meta field set in `split_id_meta_field`. Documents shared by overlapping
199 windows appear once per window.
201 """
202 window_size = self.window_size if window_size is None else window_size
203 self._validate_window_size(window_size)
204 self._raise_if_documents_do_not_have_expected_metadata(retrieved_documents)
206 windows = self._get_windows(retrieved_documents)
207 filters = self._build_filters(windows, window_size)
208 fetched_documents = self.document_store.filter_documents(filters) if filters else []
210 return self._assemble_context(retrieved_documents, windows, window_size, fetched_documents)
212 @component.output_types(context_windows=list[str], context_documents=list[Document])
213 async def run_async(self, retrieved_documents: list[Document], window_size: int | None = None) -> dict[str, Any]:
214 """
215 Based on the `source_id` and on the `doc.meta['split_id']` get surrounding documents from the document store.
217 Implements the logic behind the sentence-window technique, retrieving the surrounding documents of a given
218 document from the document store.
220 :param retrieved_documents: List of retrieved documents from the previous retriever.
221 :param window_size: The number of documents to retrieve before and after the relevant one. This will overwrite
222 the `window_size` parameter set in the constructor. It must be greater than 0; values of
223 0 and negative values are rejected.
224 :returns:
225 A dictionary with the following keys:
226 - `context_windows`: A list of strings, where each string represents the concatenated text from the
227 context window of the corresponding document in `retrieved_documents`.
228 - `context_documents`: A list of `Document` objects, containing the retrieved documents plus the
229 context documents surrounding them, grouped by window in the order of
230 `retrieved_documents`. Within each window, the documents are sorted by the
231 meta field set in `split_id_meta_field`. Documents shared by overlapping
232 windows appear once per window.
234 """
235 window_size = self.window_size if window_size is None else window_size
236 self._validate_window_size(window_size)
237 self._raise_if_documents_do_not_have_expected_metadata(retrieved_documents)
239 windows = self._get_windows(retrieved_documents)
240 filters = self._build_filters(windows, window_size)
241 # Ignoring type error because DocumentStore protocol doesn't define filter_documents_async
242 fetched_documents = (
243 await self.document_store.filter_documents_async(filters) # type: ignore[attr-defined]
244 if filters
245 else []
246 )
248 return self._assemble_context(retrieved_documents, windows, window_size, fetched_documents)
250 @staticmethod
251 def _validate_window_size(window_size: int) -> None:
252 if window_size < 1:
253 raise ValueError("The window_size parameter must be greater than 0.")
255 def _raise_if_documents_do_not_have_expected_metadata(self, retrieved_documents: list[Document]) -> None:
256 if (
257 not all(self.split_id_meta_field in doc.meta for doc in retrieved_documents)
258 and self.raise_on_missing_meta_fields
259 ):
260 raise ValueError(f"The retrieved documents must have '{self.split_id_meta_field}' in their metadata.")
262 if (
263 not all(field in doc.meta for doc in retrieved_documents for field in self._source_id_meta_fields)
264 and self.raise_on_missing_meta_fields
265 ):
266 raise ValueError(f"The retrieved documents must have '{self.source_id_meta_field}' in their metadata.")
268 def _get_windows(self, retrieved_documents: list[Document]) -> list[tuple[list[Any], int] | None]:
269 """
270 Extract the source IDs and split ID of each retrieved document, or `None` if the metadata is missing.
271 """
272 windows: list[tuple[list[Any], int] | None] = []
273 for doc in retrieved_documents:
274 source_ids = [doc.meta.get(field) for field in self._source_id_meta_fields]
275 split_id = doc.meta.get(self.split_id_meta_field)
276 if any(source_id is None for source_id in source_ids) or split_id is None:
277 logger.warning(
278 "Document {doc_id} is missing required metadata fields to be used with "
279 "SentenceWindowRetriever: {source_id} or {split_id}. Skipping context retrieval for this document.",
280 doc_id=doc.id,
281 source_id=self._source_id_meta_fields,
282 split_id=self.split_id_meta_field,
283 )
284 windows.append(None)
285 else:
286 windows.append((source_ids, split_id))
287 return windows
289 def _build_filters(self, windows: list[tuple[list[Any], int] | None], window_size: int) -> dict[str, Any] | None:
290 """
291 Combine the windows of all retrieved documents into one filter, so the Document Store is queried only once.
292 """
293 conditions: list[dict[str, Any]] = []
294 for window in windows:
295 if window is None:
296 continue
297 source_ids, split_id = window
298 condition = self._build_filter_conditions(split_id, window_size, source_ids)
299 if condition not in conditions:
300 conditions.append(condition)
301 if not conditions:
302 return None
303 return {"operator": "OR", "conditions": conditions}
305 def _assemble_context(
306 self,
307 retrieved_documents: list[Document],
308 windows: list[tuple[list[Any], int] | None],
309 window_size: int,
310 fetched_documents: list[Document],
311 ) -> dict[str, Any]:
312 """
313 Split the documents fetched with the combined filter back into the context window of each retrieved document.
314 """
315 context_text = []
316 context_documents = []
317 for doc, window in zip(retrieved_documents, windows, strict=True):
318 if window is None:
319 context_text.append(doc.content or "")
320 context_documents.append(doc)
321 continue
323 source_ids, split_id = window
324 context_docs = [
325 fetched
326 for fetched in fetched_documents
327 if self._is_in_window(fetched, source_ids, split_id - window_size, split_id + window_size)
328 ]
329 context_docs_sorted = sorted(context_docs, key=lambda d: d.meta[self.split_id_meta_field])
330 context_text.append(self.merge_documents_text(context_docs_sorted))
331 context_documents.extend(context_docs_sorted)
333 return {"context_windows": context_text, "context_documents": context_documents}
335 def _is_in_window(self, doc: Document, source_ids: list[Any], min_split_id: int, max_split_id: int) -> bool:
336 split_id = doc.meta.get(self.split_id_meta_field)
337 return (
338 split_id is not None
339 and min_split_id <= split_id <= max_split_id
340 and all(
341 doc.meta.get(field) == source_id
342 for field, source_id in zip(self._source_id_meta_fields, source_ids, strict=True)
343 )
344 )
346 def _build_filter_conditions(self, split_id: int, window_size: int, source_ids: list[Any]) -> dict[str, Any]:
347 min_before = split_id - window_size
348 max_after = split_id + window_size
349 source_id_filters = [
350 {"field": f"meta.{source_id_meta_field}", "operator": "==", "value": source_id}
351 for source_id_meta_field, source_id in zip(self._source_id_meta_fields, source_ids, strict=True)
352 ]
353 conditions = [
354 {"field": f"meta.{self.split_id_meta_field}", "operator": ">=", "value": min_before},
355 {"field": f"meta.{self.split_id_meta_field}", "operator": "<=", "value": max_after},
356 *source_id_filters,
357 ]
358 return {"operator": "AND", "conditions": conditions}
360 def close(self) -> None:
361 """
362 Release the synchronous resources of the underlying Document Store.
363 """
364 if hasattr(self.document_store, "close"):
365 self.document_store.close()
367 async def close_async(self) -> None:
368 """
369 Release the asynchronous resources of the underlying Document Store.
370 """
371 if hasattr(self.document_store, "close_async"):
372 await self.document_store.close_async()