Coverage for src / qdrant_loader_mcp_server / search / hybrid / components / reranking.py: 78%

1from __future__ import annotations 

2 

3import logging 

4from typing import Any 

5 

6from .cross_encoder_reranker import CrossEncoderReranker 

7 

8 

9class HybridReranker: 

10 """Hybrid reranker that optionally applies cross-encoder re-ranking.""" 

11 

12 def __init__( 

13 self, 

14 enabled: bool = False, 

15 model: str = "cross-encoder/ms-marco-MiniLM-L-12-v2", 

16 device: str = "cpu", 

17 batch_size: int = 32, 

18 ): 

19 self.logger = logging.getLogger(__name__) 

20 self.cross_encoder = None 

21 

22 if enabled: 

23 try: 

24 self.cross_encoder = CrossEncoderReranker( 

25 model_name=model, 

26 device=device, 

27 batch_size=batch_size, 

28 enabled=True, 

29 ) 

30 self.logger.info("Cross-encoder reranker enabled") 

31 except Exception: 

32 self.logger.exception( 

33 "Failed to initialize cross-encoder reranker; continuing without reranking" 

34 ) 

35 else: 

36 self.logger.info("Cross-encoder reranker disabled") 

37 

38 def rerank( 

39 self, 

40 query: str, 

41 results: list[Any], 

42 top_k: int | None = None, 

43 text_key: str = "text", 

44 ) -> list[Any]: 

45 

46 if not results or self.cross_encoder is None: 

47 return results 

48 

49 try: 

50 return self.cross_encoder.rerank( 

51 query=query, 

52 results=results, 

53 top_k=top_k, 

54 text_key=text_key, 

55 ) 

56 except Exception: 

57 self.logger.exception( 

58 "Cross-encoder reranking failed; returning original ranking" 

59 ) 

60 return results