Coverage for haystack/components/rankers/llm_ranker.py: 95%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from haystack import Document, component, default_from_dict, default_to_dict, logging 

8from haystack.components.builders import PromptBuilder 

9from haystack.components.generators.chat.openai import OpenAIChatGenerator 

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 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21def _default_openai_chat_generator() -> ChatGenerator: 

22 return OpenAIChatGenerator( 

23 model="gpt-4.1-mini", 

24 generation_kwargs={ 

25 "temperature": 0.0, 

26 "response_format": { 

27 "type": "json_schema", 

28 "json_schema": { 

29 "name": "document_ranking", 

30 "schema": { 

31 "type": "object", 

32 "properties": { 

33 "documents": { 

34 "type": "array", 

35 "items": { 

36 "type": "object", 

37 "properties": {"index": {"type": "integer"}}, 

38 "required": ["index"], 

39 "additionalProperties": False, 

40 }, 

41 } 

42 }, 

43 "required": ["documents"], 

44 "additionalProperties": False, 

45 }, 

46 }, 

47 }, 

48 }, 

49 ) 

50 

51 

52DEFAULT_PROMPT_TEMPLATE = """ 

53You are ranking retrieved documents for relevance to a query. 

54 

55Return valid JSON only, with this structure: 

56{ 

57 "documents": [ 

58 {"index": 1} 

59 ] 

60} 

61 

62Rules: 

63- Rank documents from most relevant to least relevant for answering the query. 

64- Only include documents that are relevant to the query. 

65- Do not return or rank documents that are not relevant. 

66- If none are relevant, return {"documents": []}. 

67- Use only document indices from the provided documents. 

68- Do not repeat document indices. 

69- Do not include explanations or any text outside the JSON object. 

70 

71Query: 

72{{ query }} 

73 

74Documents: 

75{% for document in documents %} 

76Document {{ loop.index }}: 

77content: {{ document.content or "" }} 

78 

79{% endfor %} 

80""".strip() 

81 

82 

83@component 

84class LLMRanker: 

85 """ 

86 Ranks documents for a query using a Large Language Model. 

87 

88 The LLM is expected to return a JSON object containing ranked document indices. 

89 

90 Usage example: 

91 

92 ```python 

93 from haystack import Document 

94 from haystack.components.generators.chat import OpenAIChatGenerator 

95 from haystack.components.rankers import LLMRanker 

96 

97 chat_generator = OpenAIChatGenerator( 

98 model="gpt-4.1-mini", 

99 generation_kwargs={ 

100 "temperature": 0.0, 

101 "response_format": { 

102 "type": "json_schema", 

103 "json_schema": { 

104 "name": "document_ranking", 

105 "schema": { 

106 "type": "object", 

107 "properties": { 

108 "documents": { 

109 "type": "array", 

110 "items": { 

111 "type": "object", 

112 "properties": {"index": {"type": "integer"}}, 

113 "required": ["index"], 

114 "additionalProperties": False, 

115 }, 

116 } 

117 }, 

118 "required": ["documents"], 

119 "additionalProperties": False, 

120 }, 

121 }, 

122 }, 

123 }, 

124 ) 

125 

126 ranker = LLMRanker(chat_generator=chat_generator) 

127 

128 documents = [ 

129 Document(id="paris", content="Paris is the capital of France."), 

130 Document(id="berlin", content="Berlin is the capital of Germany."), 

131 ] 

132 

133 result = ranker.run(query="capital of Germany", documents=documents) 

134 print(result["documents"][0].id) 

135 ``` 

136 """ 

137 

138 def __init__( 

139 self, 

140 *, 

141 chat_generator: ChatGenerator | None = None, 

142 prompt: str = DEFAULT_PROMPT_TEMPLATE, 

143 top_k: int = 10, 

144 raise_on_failure: bool = False, 

145 ) -> None: 

146 """ 

147 Initialize the LLMRanker component. 

148 

149 :param chat_generator: 

150 The chat generator to use for reranking. If `None`, a default `OpenAIChatGenerator` configured for JSON 

151 output is used. 

152 :param prompt: 

153 Custom prompt template for reranking. The prompt must include exactly the variables `query` and 

154 `documents` and instruct the LLM to return ranked 1-based document indices as JSON. 

155 :param top_k: 

156 The maximum number of documents to return. 

157 :param raise_on_failure: 

158 If `True`, raise when generation or response parsing fails. If `False`, log the failure and return the 

159 input documents in fallback order. 

160 """ 

161 if top_k <= 0: 

162 raise ValueError(f"top_k must be > 0, but got {top_k}") 

163 

164 self.top_k = top_k 

165 self.raise_on_failure = raise_on_failure 

166 self.prompt = prompt 

167 self._prompt_builder = PromptBuilder(template=self.prompt, required_variables=["documents", "query"]) 

168 if set(self._prompt_builder.variables) != {"documents", "query"}: 

169 raise ValueError("prompt must include exactly the variables 'documents' and 'query'.") 

170 

171 if chat_generator is None: 

172 self._chat_generator = _default_openai_chat_generator() 

173 else: 

174 self._chat_generator = chat_generator 

175 

176 def warm_up(self) -> None: 

177 """Warm up the underlying chat generator.""" 

178 if hasattr(self._chat_generator, "warm_up"): 

179 self._chat_generator.warm_up() 

180 

181 async def warm_up_async(self) -> None: 

182 """Warm up the underlying chat generator on the serving event loop.""" 

183 if hasattr(self._chat_generator, "warm_up_async"): 

184 await self._chat_generator.warm_up_async() 

185 elif hasattr(self._chat_generator, "warm_up"): 

186 self._chat_generator.warm_up() 

187 

188 def close(self) -> None: 

189 """Release the underlying chat generator's resources.""" 

190 if hasattr(self._chat_generator, "close"): 

191 self._chat_generator.close() 

192 

193 async def close_async(self) -> None: 

194 """Release the underlying chat generator's async resources.""" 

195 if hasattr(self._chat_generator, "close_async"): 

196 await self._chat_generator.close_async() 

197 elif hasattr(self._chat_generator, "close"): 

198 self._chat_generator.close() 

199 

200 def to_dict(self) -> dict[str, Any]: 

201 """ 

202 Serialize this component to a dictionary. 

203 

204 :returns: 

205 Dictionary with serialized data. 

206 """ 

207 return default_to_dict( 

208 self, 

209 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"), 

210 prompt=self.prompt, 

211 top_k=self.top_k, 

212 raise_on_failure=self.raise_on_failure, 

213 ) 

214 

215 @classmethod 

216 def from_dict(cls, data: dict[str, Any]) -> "LLMRanker": 

217 """ 

218 Deserialize this component from a dictionary. 

219 

220 :param data: 

221 The dictionary representation of the component. 

222 :returns: 

223 The deserialized component instance. 

224 """ 

225 init_params = data.get("init_parameters", {}) 

226 if init_params.get("chat_generator"): 

227 deserialize_chatgenerator_inplace(init_params, key="chat_generator") 

228 return default_from_dict(cls, data) 

229 

230 @component.output_types(documents=list[Document]) 

231 def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]: 

232 """ 

233 Rank documents for a query using an LLM. 

234 

235 Before ranking, duplicate documents are removed. 

236 

237 :param query: 

238 The query used for reranking. 

239 :param documents: 

240 Candidate documents to rerank. 

241 :param top_k: 

242 The maximum number of documents to return. Overrides the instance's `top_k` if provided. 

243 :returns: 

244 A dictionary with the ranked documents under the `documents` key. 

245 """ 

246 if top_k is not None and top_k <= 0: 

247 raise ValueError(f"top_k must be > 0, but got {top_k}") 

248 

249 if not documents: 

250 return {"documents": []} 

251 

252 top_k = self.top_k if top_k is None else top_k 

253 deduplicated_documents = _deduplicate_documents(documents) 

254 fallback_documents = deduplicated_documents 

255 

256 if not query.strip(): 

257 logger.warning("Empty query provided to LLMRanker. Returning documents without reranking.") 

258 return {"documents": fallback_documents} 

259 

260 self.warm_up() 

261 

262 prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents) 

263 messages = [ChatMessage.from_user(prompt["prompt"])] 

264 

265 try: 

266 with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span: 

267 result = self._chat_generator.run(messages=messages) 

268 span.set_content_tag("haystack.component.output", result) 

269 except Exception as exc: 

270 if self.raise_on_failure: 

271 raise 

272 logger.warning( 

273 "LLMRanker failed during chat generation. Returning fallback order. Error: {error}", error=exc 

274 ) 

275 return {"documents": fallback_documents} 

276 

277 try: 

278 reply_text = self._get_reply_text(result) 

279 ranked_documents = self._rank_documents_from_reply(reply_text=reply_text, documents=deduplicated_documents) 

280 except (TypeError, ValueError) as exc: 

281 if self.raise_on_failure: 

282 raise 

283 logger.warning( 

284 "LLMRanker failed while processing the chat response. Returning fallback order. Error: {error}", 

285 error=exc, 

286 ) 

287 return {"documents": fallback_documents} 

288 

289 return {"documents": ranked_documents[:top_k]} 

290 

291 @component.output_types(documents=list[Document]) 

292 async def run_async( 

293 self, query: str, documents: list[Document], top_k: int | None = None 

294 ) -> dict[str, list[Document]]: 

295 """ 

296 Asynchronously rank documents for a query using an LLM. 

297 

298 Before ranking, duplicate documents are removed. 

299 

300 This is the asynchronous version of the `run` method. It has the same parameters and return values 

301 but can be used with `await` in an async code. If the chat generator only implements a synchronous 

302 `run` method, it is executed in a thread to avoid blocking the event loop. 

303 

304 :param query: 

305 The query used for reranking. 

306 :param documents: 

307 Candidate documents to rerank. 

308 :param top_k: 

309 The maximum number of documents to return. Overrides the instance's `top_k` if provided. 

310 :returns: 

311 A dictionary with the ranked documents under the `documents` key. 

312 """ 

313 if top_k is not None and top_k <= 0: 

314 raise ValueError(f"top_k must be > 0, but got {top_k}") 

315 

316 if not documents: 

317 return {"documents": []} 

318 

319 top_k = self.top_k if top_k is None else top_k 

320 deduplicated_documents = _deduplicate_documents(documents) 

321 fallback_documents = deduplicated_documents 

322 

323 if not query.strip(): 

324 logger.warning("Empty query provided to LLMRanker. Returning documents without reranking.") 

325 return {"documents": fallback_documents} 

326 

327 await self.warm_up_async() 

328 

329 prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents) 

330 messages = [ChatMessage.from_user(prompt["prompt"])] 

331 

332 try: 

333 with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span: 

334 result = await _execute_component_async(self._chat_generator, messages=messages) 

335 span.set_content_tag("haystack.component.output", result) 

336 except Exception as exc: 

337 if self.raise_on_failure: 

338 raise 

339 logger.warning( 

340 "LLMRanker failed during chat generation. Returning fallback order. Error: {error}", error=exc 

341 ) 

342 return {"documents": fallback_documents} 

343 

344 try: 

345 reply_text = self._get_reply_text(result) 

346 ranked_documents = self._rank_documents_from_reply(reply_text=reply_text, documents=deduplicated_documents) 

347 except (TypeError, ValueError) as exc: 

348 if self.raise_on_failure: 

349 raise 

350 logger.warning( 

351 "LLMRanker failed while processing the chat response. Returning fallback order. Error: {error}", 

352 error=exc, 

353 ) 

354 return {"documents": fallback_documents} 

355 

356 return {"documents": ranked_documents[:top_k]} 

357 

358 @staticmethod 

359 def _get_reply_text(result: dict[str, Any]) -> str: 

360 replies = result.get("replies") or [] 

361 if not replies: 

362 raise ValueError("ChatGenerator returned no replies.") 

363 

364 reply_text = replies[0].text 

365 if reply_text is None: 

366 raise ValueError("ChatGenerator returned a reply without text.") 

367 

368 return reply_text 

369 

370 @staticmethod 

371 def _rank_documents_from_reply(reply_text: str, documents: list[Document]) -> list[Document]: 

372 parsed_response = _parse_dict_from_json(reply_text, expected_keys=["documents"], raise_on_failure=True) 

373 ranked_entries = parsed_response["documents"] 

374 

375 if not isinstance(ranked_entries, list): 

376 raise TypeError("Expected 'documents' in ranking response to be a list.") 

377 

378 if not ranked_entries: 

379 return [] 

380 

381 ranked_documents: list[Document] = [] 

382 

383 for entry in ranked_entries: 

384 if not isinstance(entry, dict): 

385 raise TypeError("Expected each ranked document entry to be a JSON object.") 

386 

387 document_index = entry.get("index") 

388 if document_index is None: 

389 continue 

390 

391 try: 

392 # LLMs can return numeric indices as strings even when asked for integers. 

393 document_index = int(document_index) 

394 except (TypeError, ValueError): 

395 continue 

396 

397 # Jinja's `loop.index` is 1-based: 

398 # https://jinja.palletsprojects.com/en/stable/templates/#for 

399 if document_index < 1 or document_index > len(documents): 

400 continue 

401 

402 document = documents[document_index - 1] 

403 ranked_documents.append(document) 

404 

405 if not ranked_documents: 

406 raise ValueError("Ranking response did not contain any valid document indices.") 

407 

408 return ranked_documents