Coverage for haystack/components/rankers/llm_ranker.py: 97%
154 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.components.builders import PromptBuilder
9from haystack.components.generators.chat.openai_responses import OpenAIResponsesChatGenerator
10from haystack.components.generators.chat.types import ChatGenerator
11from haystack.components.generators.utils import _trace_chat_generator_run
12from haystack.core.serialization import component_to_dict
13from haystack.dataclasses import ChatMessage
14from haystack.utils import deserialize_chatgenerator_inplace
15from haystack.utils.async_utils import _execute_component_async
16from haystack.utils.misc import _deduplicate_documents, _parse_dict_from_json
18logger = logging.getLogger(__name__)
21def _default_openai_chat_generator() -> ChatGenerator:
22 return OpenAIResponsesChatGenerator(
23 model="gpt-4.1-mini",
24 generation_kwargs={
25 "temperature": 0.0,
26 "store": False,
27 "text": {
28 "format": {
29 "type": "json_schema",
30 "name": "document_ranking",
31 "schema": {
32 "type": "object",
33 "properties": {
34 "documents": {
35 "type": "array",
36 "items": {
37 "type": "object",
38 "properties": {"index": {"type": "integer"}},
39 "required": ["index"],
40 "additionalProperties": False,
41 },
42 }
43 },
44 "required": ["documents"],
45 "additionalProperties": False,
46 },
47 }
48 },
49 },
50 )
53DEFAULT_PROMPT_TEMPLATE = """
54You are ranking retrieved documents for relevance to a query.
56Return valid JSON only, with this structure:
57{
58 "documents": [
59 {"index": 1}
60 ]
61}
63Rules:
64- Rank documents from most relevant to least relevant for answering the query.
65- Only include documents that are relevant to the query.
66- Do not return or rank documents that are not relevant.
67- If none are relevant, return {"documents": []}.
68- Use only document indices from the provided documents.
69- Do not repeat document indices.
70- Do not include explanations or any text outside the JSON object.
72Query:
73{{ query }}
75Documents:
76{% for document in documents %}
77Document {{ loop.index }}:
78content: {{ document.content or "" }}
80{% endfor %}
81""".strip()
84@component
85class LLMRanker:
86 """
87 Ranks documents for a query using a Large Language Model.
89 The LLM is expected to return a JSON object containing ranked document indices. Only the documents the LLM returns
90 are kept, so the output can contain fewer than `top_k` documents, or none if the LLM considers none relevant.
92 If the query is empty, or if generation or parsing fails and `raise_on_failure` is `False`, the ranker falls back
93 to the deduplicated input documents, unranked and in their original order. `top_k` is not applied to them, since
94 cutting off unranked documents could drop relevant ones.
96 Usage example:
98 ```python
99 from haystack import Document
100 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
101 from haystack.components.rankers import LLMRanker
103 chat_generator = OpenAIResponsesChatGenerator(
104 model="gpt-5.6-luna",
105 generation_kwargs={
106 "text": {
107 "format": {
108 "type": "json_schema",
109 "name": "document_ranking",
110 "strict": True,
111 "schema": {
112 "type": "object",
113 "properties": {
114 "documents": {
115 "type": "array",
116 "items": {
117 "type": "object",
118 "properties": {"index": {"type": "integer"}},
119 "required": ["index"],
120 "additionalProperties": False,
121 },
122 }
123 },
124 "required": ["documents"],
125 "additionalProperties": False,
126 },
127 }
128 }
129 },
130 )
132 ranker = LLMRanker(chat_generator=chat_generator)
134 documents = [
135 Document(id="paris", content="Paris is the capital of France."),
136 Document(id="berlin", content="Berlin is the capital of Germany."),
137 ]
139 result = ranker.run(query="capital of Germany", documents=documents)
140 print(result["documents"][0].id)
141 ```
142 """
144 def __init__(
145 self,
146 *,
147 chat_generator: ChatGenerator | None = None,
148 prompt: str = DEFAULT_PROMPT_TEMPLATE,
149 top_k: int = 10,
150 raise_on_failure: bool = False,
151 ) -> None:
152 """
153 Initialize the LLMRanker component.
155 :param chat_generator:
156 The chat generator to use for reranking. If `None`, a default `OpenAIResponsesChatGenerator` configured for
157 JSON output and with `store` set to `False` is used.
158 :param prompt:
159 Custom prompt template for reranking. The prompt must include exactly the variables `query` and
160 `documents` and instruct the LLM to return ranked 1-based document indices as JSON.
161 :param top_k:
162 The maximum number of ranked documents to return. Not applied when the ranker falls back to the input
163 documents.
164 :param raise_on_failure:
165 If `True`, raise when generation or response parsing fails. If `False`, log the failure and return all
166 deduplicated input documents unranked, in their original order, without applying `top_k`.
167 :raises ValueError:
168 If `top_k` is not greater than 0 or if `prompt` doesn't include exactly the variables `query` and
169 `documents`.
170 """
171 if top_k <= 0:
172 raise ValueError(f"top_k must be > 0, but got {top_k}")
174 self.top_k = top_k
175 self.raise_on_failure = raise_on_failure
176 self.prompt = prompt
177 self._prompt_builder = PromptBuilder(template=self.prompt, required_variables=["documents", "query"])
178 if set(self._prompt_builder.variables) != {"documents", "query"}:
179 raise ValueError("prompt must include exactly the variables 'documents' and 'query'.")
181 if chat_generator is None:
182 self._chat_generator = _default_openai_chat_generator()
183 else:
184 self._chat_generator = chat_generator
186 def warm_up(self) -> None:
187 """Warm up the underlying chat generator."""
188 if hasattr(self._chat_generator, "warm_up"):
189 self._chat_generator.warm_up()
191 async def warm_up_async(self) -> None:
192 """Warm up the underlying chat generator on the serving event loop."""
193 if hasattr(self._chat_generator, "warm_up_async"):
194 await self._chat_generator.warm_up_async()
195 elif hasattr(self._chat_generator, "warm_up"):
196 self._chat_generator.warm_up()
198 def close(self) -> None:
199 """Release the underlying chat generator's resources."""
200 if hasattr(self._chat_generator, "close"):
201 self._chat_generator.close()
203 async def close_async(self) -> None:
204 """Release the underlying chat generator's async resources."""
205 if hasattr(self._chat_generator, "close_async"):
206 await self._chat_generator.close_async()
207 elif hasattr(self._chat_generator, "close"):
208 self._chat_generator.close()
210 def to_dict(self) -> dict[str, Any]:
211 """
212 Serialize this component to a dictionary.
214 :returns:
215 Dictionary with serialized data.
216 """
217 return default_to_dict(
218 self,
219 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
220 prompt=self.prompt,
221 top_k=self.top_k,
222 raise_on_failure=self.raise_on_failure,
223 )
225 @classmethod
226 def from_dict(cls, data: dict[str, Any]) -> "LLMRanker":
227 """
228 Deserialize this component from a dictionary.
230 :param data:
231 The dictionary representation of the component.
232 :returns:
233 The deserialized component instance.
234 """
235 init_params = data.get("init_parameters", {})
236 if init_params.get("chat_generator"):
237 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
238 return default_from_dict(cls, data)
240 @component.output_types(documents=list[Document])
241 def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
242 """
243 Rank documents for a query using an LLM.
245 Before ranking, duplicate documents are removed.
247 :param query:
248 The query used for reranking. If empty, the documents are returned without ranking.
249 :param documents:
250 Candidate documents to rerank.
251 :param top_k:
252 The maximum number of ranked documents to return. Overrides the instance's `top_k` if provided. Not
253 applied when the ranker falls back to the input documents.
254 :returns:
255 A dictionary with the following key:
256 - `documents`: Up to `top_k` documents the LLM considers relevant, from most to least relevant. If
257 `query` is empty, or if generation or parsing fails and `raise_on_failure` is `False`, all deduplicated
258 input documents unranked, in their original order.
259 :raises ValueError:
260 If `top_k` is not greater than 0.
261 """
262 if top_k is not None and top_k <= 0:
263 raise ValueError(f"top_k must be > 0, but got {top_k}")
265 if not documents:
266 return {"documents": []}
268 top_k = self.top_k if top_k is None else top_k
269 deduplicated_documents = _deduplicate_documents(documents)
270 fallback_documents = deduplicated_documents
272 if not isinstance(query, str) or not query.strip():
273 logger.warning("Empty query provided to LLMRanker. Returning the deduplicated input documents unranked.")
274 return {"documents": fallback_documents}
276 self.warm_up()
278 prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents)
279 messages = [ChatMessage.from_user(prompt["prompt"])]
281 try:
282 with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
283 result = self._chat_generator.run(messages=messages)
284 span.set_content_tag("haystack.component.output", result)
285 except Exception as exc:
286 if self.raise_on_failure:
287 raise
288 logger.warning(
289 "LLMRanker failed during chat generation. Returning the deduplicated input documents unranked. "
290 "Error: {error}",
291 error=exc,
292 )
293 return {"documents": fallback_documents}
295 try:
296 reply_text = self._get_reply_text(result)
297 ranked_documents = self._rank_documents_from_reply(reply_text=reply_text, documents=deduplicated_documents)
298 except (TypeError, ValueError) as exc:
299 if self.raise_on_failure:
300 raise
301 logger.warning(
302 "LLMRanker failed while processing the chat response. Returning the deduplicated input documents "
303 "unranked. Error: {error}",
304 error=exc,
305 )
306 return {"documents": fallback_documents}
308 return {"documents": ranked_documents[:top_k]}
310 @component.output_types(documents=list[Document])
311 async def run_async(
312 self, query: str, documents: list[Document], top_k: int | None = None
313 ) -> dict[str, list[Document]]:
314 """
315 Asynchronously rank documents for a query using an LLM.
317 Before ranking, duplicate documents are removed.
319 This is the asynchronous version of the `run` method. It has the same parameters and return values
320 but can be used with `await` in an async code. If the chat generator only implements a synchronous
321 `run` method, it is executed in a thread to avoid blocking the event loop.
323 :param query:
324 The query used for reranking. If empty, the documents are returned without ranking.
325 :param documents:
326 Candidate documents to rerank.
327 :param top_k:
328 The maximum number of ranked documents to return. Overrides the instance's `top_k` if provided. Not
329 applied when the ranker falls back to the input documents.
330 :returns:
331 A dictionary with the following key:
332 - `documents`: Up to `top_k` documents the LLM considers relevant, from most to least relevant. If
333 `query` is empty, or if generation or parsing fails and `raise_on_failure` is `False`, all deduplicated
334 input documents unranked, in their original order.
335 :raises ValueError:
336 If `top_k` is not greater than 0.
337 """
338 if top_k is not None and top_k <= 0:
339 raise ValueError(f"top_k must be > 0, but got {top_k}")
341 if not documents:
342 return {"documents": []}
344 top_k = self.top_k if top_k is None else top_k
345 deduplicated_documents = _deduplicate_documents(documents)
346 fallback_documents = deduplicated_documents
348 if not isinstance(query, str) or not query.strip():
349 logger.warning("Empty query provided to LLMRanker. Returning the deduplicated input documents unranked.")
350 return {"documents": fallback_documents}
352 await self.warm_up_async()
354 prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents)
355 messages = [ChatMessage.from_user(prompt["prompt"])]
357 try:
358 with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
359 result = await _execute_component_async(self._chat_generator, messages=messages)
360 span.set_content_tag("haystack.component.output", result)
361 except Exception as exc:
362 if self.raise_on_failure:
363 raise
364 logger.warning(
365 "LLMRanker failed during chat generation. Returning the deduplicated input documents unranked. "
366 "Error: {error}",
367 error=exc,
368 )
369 return {"documents": fallback_documents}
371 try:
372 reply_text = self._get_reply_text(result)
373 ranked_documents = self._rank_documents_from_reply(reply_text=reply_text, documents=deduplicated_documents)
374 except (TypeError, ValueError) as exc:
375 if self.raise_on_failure:
376 raise
377 logger.warning(
378 "LLMRanker failed while processing the chat response. Returning the deduplicated input documents "
379 "unranked. Error: {error}",
380 error=exc,
381 )
382 return {"documents": fallback_documents}
384 return {"documents": ranked_documents[:top_k]}
386 @staticmethod
387 def _get_reply_text(result: dict[str, Any]) -> str:
388 replies = result.get("replies") or []
389 if not replies:
390 raise ValueError("ChatGenerator returned no replies.")
392 reply_text = replies[0].text
393 if reply_text is None:
394 raise ValueError("ChatGenerator returned a reply without text.")
396 return reply_text
398 @staticmethod
399 def _rank_documents_from_reply(reply_text: str, documents: list[Document]) -> list[Document]:
400 parsed_response = _parse_dict_from_json(reply_text, expected_keys=["documents"], raise_on_failure=True)
401 ranked_entries = parsed_response["documents"]
403 if not isinstance(ranked_entries, list):
404 raise TypeError("Expected 'documents' in ranking response to be a list.")
406 if not ranked_entries:
407 return []
409 ranked_documents: list[Document] = []
411 for entry in ranked_entries:
412 if not isinstance(entry, dict):
413 raise TypeError("Expected each ranked document entry to be a JSON object.")
415 document_index = entry.get("index")
416 if document_index is None:
417 continue
419 try:
420 # LLMs can return numeric indices as strings even when asked for integers.
421 document_index = int(document_index)
422 except (TypeError, ValueError):
423 continue
425 # Jinja's `loop.index` is 1-based:
426 # https://jinja.palletsprojects.com/en/stable/templates/#for
427 if document_index < 1 or document_index > len(documents):
428 continue
430 document = documents[document_index - 1]
431 ranked_documents.append(document)
433 if not ranked_documents:
434 raise ValueError("Ranking response did not contain any valid document indices.")
436 return ranked_documents