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

1"""Length optimizer for context. 

2 

3This module implements length optimization to fit context within token limits 

4while preserving the most relevant information. 

5""" 

6 

7from __future__ import annotations 

8 

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

10 

11 

12class LengthOptimizer: 

13 """Optimize context length to fit within limits. 

14 

15 This component truncates or summarizes context to fit within token 

16 limits while prioritizing the most relevant content. 

17 

18 Attributes: 

19 max_tokens: Maximum total tokens allowed 

20 chars_per_token: Approximate characters per token 

21 preserve_order: Whether to preserve chunk order 

22 """ 

23 

24 def __init__( 

25 self, 

26 max_tokens: int = 4000, 

27 chars_per_token: int = 4, 

28 preserve_order: bool = True, 

29 ): 

30 """Initialize the length optimizer. 

31 

32 Args: 

33 max_tokens: Maximum tokens allowed 

34 chars_per_token: Approximate characters per token 

35 preserve_order: Whether to preserve order 

36 """ 

37 self.max_tokens = max_tokens 

38 self.chars_per_token = chars_per_token 

39 self.preserve_order = preserve_order 

40 

41 def _estimate_tokens(self, text: str) -> int: 

42 """Estimate token count from text. 

43 

44 Args: 

45 text: Input text 

46 

47 Returns: 

48 Estimated token count 

49 """ 

50 return len(text) // self.chars_per_token 

51 

52 async def optimize_length( 

53 self, 

54 chunks: list[ContextChunk], 

55 ) -> list[ContextChunk]: 

56 """Optimize chunk list to fit within token limit. 

57 

58 Args: 

59 chunks: Chunks to optimize 

60 

61 Returns: 

62 Optimized list of chunks 

63 """ 

64 if not chunks: 

65 return [] 

66 

67 # Calculate total tokens 

68 total_tokens = sum(self._estimate_tokens(c.text) for c in chunks) 

69 

70 # If already within limit, return as-is 

71 if total_tokens <= self.max_tokens: 

72 return chunks 

73 

74 # Sort by score/rank to prioritize best chunks 

75 sorted_chunks = sorted( 

76 chunks, 

77 key=lambda c: (c.score if c.score else 0, -c.rank), 

78 reverse=True, 

79 ) 

80 

81 # Greedily select chunks until limit 

82 selected_chunks: list[ContextChunk] = [] 

83 current_tokens = 0 

84 

85 for chunk in sorted_chunks: 

86 chunk_tokens = self._estimate_tokens(chunk.text) 

87 

88 if current_tokens + chunk_tokens <= self.max_tokens: 

89 selected_chunks.append(chunk) 

90 current_tokens += chunk_tokens 

91 elif current_tokens < self.max_tokens: 

92 # Try to fit partial chunk 

93 remaining_tokens = self.max_tokens - current_tokens 

94 remaining_chars = remaining_tokens * self.chars_per_token 

95 

96 if remaining_chars >= 100: # Minimum useful chunk 

97 # Truncate chunk 

98 truncated_text = chunk.text[:remaining_chars] + "..." 

99 truncated_chunk = ContextChunk( 

100 text=truncated_text, 

101 source=chunk.source, 

102 score=chunk.score, 

103 metadata={**chunk.metadata, "truncated": True}, 

104 rank=chunk.rank, 

105 ) 

106 selected_chunks.append(truncated_chunk) 

107 break 

108 

109 # Restore original order if requested 

110 if self.preserve_order: 

111 selected_chunks.sort(key=lambda c: c.rank) 

112 

113 return selected_chunks 

114 

115 async def optimize_with_budget( 

116 self, 

117 chunks: list[ContextChunk], 

118 token_budget: int, 

119 ) -> list[ContextChunk]: 

120 """Optimize chunks with a specific token budget. 

121 

122 Args: 

123 chunks: Chunks to optimize 

124 token_budget: Token budget for this optimization 

125 

126 Returns: 

127 Optimized chunks 

128 """ 

129 original_max = self.max_tokens 

130 self.max_tokens = token_budget 

131 

132 result = await self.optimize_length(chunks) 

133 

134 self.max_tokens = original_max 

135 return result