Coverage for haystack/components/rankers/lost_in_the_middle.py: 95%
43 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
6from haystack import Document, component
7from haystack.utils.misc import _deduplicate_documents
10@component
11class LostInTheMiddleRanker:
12 """
13 A LostInTheMiddle Ranker.
15 Ranks documents based on the 'lost in the middle' order so that the most relevant documents are either at the
16 beginning or end, while the least relevant are in the middle.
18 LostInTheMiddleRanker assumes that some prior component in the pipeline has already ranked documents by relevance
19 and requires no query as input but only documents. It is typically used as the last component before building a
20 prompt for an LLM to prepare the input context for the LLM.
22 Lost in the Middle ranking lays out document contents into LLM context so that the most relevant contents are at
23 the beginning or end of the input context, while the least relevant is in the middle of the context. See the
24 paper ["Lost in the Middle: How Language Models Use Long Contexts"](https://arxiv.org/abs/2307.03172) for more
25 details.
27 Usage example:
28 ```python
29 from haystack.components.rankers import LostInTheMiddleRanker
30 from haystack import Document
32 ranker = LostInTheMiddleRanker()
33 docs = [Document(content="Paris"), Document(content="Berlin"), Document(content="Madrid")]
34 result = ranker.run(documents=docs)
35 for doc in result["documents"]:
36 print(doc.content)
37 ```
38 """
40 def __init__(self, word_count_threshold: int | None = None, top_k: int | None = None) -> None:
41 """
42 Initialize the LostInTheMiddleRanker.
44 If 'word_count_threshold' is specified, this ranker includes all documents up until the point where adding
45 another document would exceed the 'word_count_threshold'. The last document that causes the threshold to
46 be breached will be included in the resulting list of documents, but all subsequent documents will be
47 discarded.
49 :param word_count_threshold: The maximum total number of words across all documents selected by the ranker.
50 :param top_k: The maximum number of documents to return.
51 """
52 if isinstance(word_count_threshold, int) and word_count_threshold <= 0:
53 raise ValueError(
54 f"Invalid value for word_count_threshold: {word_count_threshold}. word_count_threshold must be > 0."
55 )
56 if isinstance(top_k, int) and top_k <= 0:
57 raise ValueError(f"top_k must be > 0, but got {top_k}")
59 self.word_count_threshold = word_count_threshold
60 self.top_k = top_k
62 @component.output_types(documents=list[Document])
63 def run(
64 self, documents: list[Document], top_k: int | None = None, word_count_threshold: int | None = None
65 ) -> dict[str, list[Document]]:
66 """
67 Reranks documents based on the "lost in the middle" order.
69 Before ranking, documents are deduplicated by their id, retaining only the document with the highest score
70 if a score is present.
72 :param documents: List of Documents to reorder.
73 :param top_k: The maximum number of documents to return.
74 :param word_count_threshold: The maximum total number of words across all documents selected by the ranker.
75 :returns:
76 A dictionary with the following keys:
77 - `documents`: Reranked list of Documents
79 :raises ValueError:
80 If any of the documents is not textual.
81 """
82 if isinstance(word_count_threshold, int) and word_count_threshold <= 0:
83 raise ValueError(
84 f"Invalid value for word_count_threshold: {word_count_threshold}. word_count_threshold must be > 0."
85 )
86 if isinstance(top_k, int) and top_k <= 0:
87 raise ValueError(f"top_k must be > 0, but got {top_k}")
89 if not documents:
90 return {"documents": []}
92 top_k = top_k or self.top_k
93 word_count_threshold = word_count_threshold or self.word_count_threshold
95 deduplicated_documents = _deduplicate_documents(documents)
96 documents_to_reorder = deduplicated_documents[:top_k] if top_k else deduplicated_documents
98 # If there's only one document, return it as is
99 if len(documents_to_reorder) == 1:
100 return {"documents": documents_to_reorder}
102 # Raise an error if any document is not textual
103 if any(doc.content is None for doc in documents_to_reorder):
104 raise ValueError("Some provided documents are not textual; LostInTheMiddleRanker can process only text.")
106 # Initialize word count and indices for the "lost in the middle" order
107 word_count = 0
108 document_index = list(range(len(documents_to_reorder)))
109 lost_in_the_middle_indices = [0]
111 # If word count threshold is set and the first document has content, calculate word count for the first document
112 if word_count_threshold and documents_to_reorder[0].content:
113 word_count = len(documents_to_reorder[0].content.split())
115 # If the first document already meets the word count threshold, return it
116 if word_count >= word_count_threshold:
117 return {"documents": [documents_to_reorder[0]]}
119 # Start from the second document and create "lost in the middle" order
120 for doc_idx in document_index[1:]:
121 # Calculate the index at which the current document should be inserted
122 insertion_index = len(lost_in_the_middle_indices) // 2 + len(lost_in_the_middle_indices) % 2
124 # Insert the document index at the calculated position
125 lost_in_the_middle_indices.insert(insertion_index, doc_idx)
127 # If word count threshold is set and the document has content, calculate the total word count
128 if word_count_threshold and documents_to_reorder[doc_idx].content:
129 word_count += len(documents_to_reorder[doc_idx].content.split()) # type: ignore[union-attr]
131 # If the total word count meets the threshold, stop processing further documents
132 if word_count >= word_count_threshold:
133 break
135 # Documents in the "lost in the middle" order
136 ranked_docs = [documents_to_reorder[idx] for idx in lost_in_the_middle_indices]
137 return {"documents": ranked_docs}