Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/synthesis/context/ranker.py: 29%

1"""Context ranker for synthesis. 

2 

3This module implements context ranking to reorder chunks by relevance 

4before synthesis. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING 

10 

11if TYPE_CHECKING: 

12 from lexigram.contracts.ai import EmbeddingClientProtocol 

13 

14from lexigram.ai.rag.synthesis.types import ContextChunk 

15 

16 

17class ContextRanker: 

18 """Rank context chunks by relevance. 

19 

20 This component reorders chunks to prioritize the most relevant content 

21 for synthesis. 

22 

23 Attributes: 

24 embedding_client: Optional embedding client for semantic ranking 

25 use_scores: Whether to use existing chunk scores 

26 use_recency: Whether to consider recency (if available in metadata) 

27 """ 

28 

29 def __init__( 

30 self, 

31 embedding_client: EmbeddingClientProtocol | None = None, 

32 use_scores: bool = True, 

33 use_recency: bool = False, 

34 ): 

35 """Initialize the context ranker. 

36 

37 Args: 

38 embedding_client: Optional embedding client 

39 use_scores: Whether to use chunk scores 

40 use_recency: Whether to consider recency 

41 """ 

42 self.embedding_client = embedding_client 

43 self.use_scores = use_scores 

44 self.use_recency = use_recency 

45 

46 async def rank_chunks( 

47 self, 

48 query: str, 

49 chunks: list[ContextChunk], 

50 ) -> list[ContextChunk]: 

51 """Rank chunks by relevance to query. 

52 

53 Args: 

54 query: The user query 

55 chunks: Chunks to rank 

56 

57 Returns: 

58 Ranked list of chunks 

59 """ 

60 if not chunks: 

61 return [] 

62 

63 # If using existing scores, just sort by score 

64 if self.use_scores and all( 

65 chunk.score is not None and chunk.score > 0 for chunk in chunks 

66 ): 

67 ranked = sorted( 

68 chunks, 

69 key=lambda c: c.score if c.score is not None else 0.0, 

70 reverse=True, 

71 ) 

72 

73 # Update ranks 

74 for i, chunk in enumerate(ranked): 

75 object.__setattr__(chunk, "rank", i) 

76 

77 return ranked 

78 

79 # Otherwise, use simple ranking 

80 # For now, preserve original order and set ranks 

81 for i, chunk in enumerate(chunks): 

82 object.__setattr__(chunk, "rank", i) 

83 

84 return chunks 

85 

86 async def rerank_chunks( 

87 self, 

88 query: str, 

89 chunks: list[ContextChunk], 

90 top_k: int | None = None, 

91 ) -> list[ContextChunk]: 

92 """Rerank chunks and optionally limit to top K. 

93 

94 Args: 

95 query: The user query 

96 chunks: Chunks to rerank 

97 top_k: Number of top chunks to return (None = all) 

98 

99 Returns: 

100 Reranked and filtered chunks 

101 """ 

102 ranked = await self.rank_chunks(query, chunks) 

103 

104 if top_k: 

105 return ranked[:top_k] 

106 

107 return ranked