Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-memory/src/lexigram/ai/memory/pruning/pruner.py: 26%

1"""Dynamic context pruner for token-aware memory management.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any 

6 

7from lexigram.ai.memory.pruning.scorer import HybridScorerImpl, RecencyScorerImpl 

8from lexigram.ai.memory.pruning.types import PruningResult, PruningStrategy 

9from lexigram.logging import ( 

10 get_logger, 

11) 

12 

13if TYPE_CHECKING: 

14 from lexigram.contracts.ai.llm import TokenCounterProtocol 

15 from lexigram.contracts.ai.memory import MemoryEntry 

16 

17logger = get_logger(__name__) 

18 

19 

20class DynamicContextPruner: 

21 """Prunes MemoryEntry lists to fit within a token budget. 

22 

23 Uses pluggable scoring strategies to rank entries by importance, then 

24 greedily selects entries until the token budget is exhausted. 

25 

26 Attributes: 

27 token_counter: Protocol for counting tokens in text. 

28 default_strategy: Default pruning strategy when none is specified. 

29 """ 

30 

31 def __init__( 

32 self, 

33 token_counter: TokenCounterProtocol, 

34 default_strategy: PruningStrategy = PruningStrategy.HYBRID, 

35 ) -> None: 

36 """Initialize the pruner. 

37 

38 Args: 

39 token_counter: Implementation of TokenCounterProtocol for token counting. 

40 default_strategy: Default strategy to use if not overridden. Default is HYBRID. 

41 """ 

42 self.token_counter = token_counter 

43 self.default_strategy = default_strategy 

44 

45 async def prune( 

46 self, 

47 entries: list[MemoryEntry], 

48 token_budget: int, 

49 query: str | None = None, 

50 strategy: PruningStrategy | None = None, 

51 **kwargs, 

52 ) -> PruningResult: 

53 """Prune a list of memory entries to fit within a token budget. 

54 

55 Scores all entries using the specified strategy, sorts them by score 

56 (descending), then greedily keeps entries until adding the next entry 

57 would exceed the remaining budget. 

58 

59 Args: 

60 entries: List of MemoryEntry objects to prune. 

61 token_budget: Maximum number of tokens to keep. 

62 query: Optional query context for relevance-based scoring. 

63 strategy: Override the default pruning strategy. If None, uses default_strategy. 

64 **kwargs: Additional keyword arguments (reserved for future use). 

65 

66 Returns: 

67 PruningResult containing kept entries (score-ordered), counts, 

68 and metadata about the pruning operation. 

69 """ 

70 # Use default strategy if none provided 

71 selected_strategy = strategy or self.default_strategy 

72 

73 # Handle empty input 

74 if not entries: 

75 return PruningResult( 

76 kept=[], 

77 pruned_count=0, 

78 original_count=0, 

79 token_budget=token_budget, 

80 strategy=selected_strategy, 

81 metadata={}, 

82 ) 

83 

84 # Select scorer based on strategy 

85 scorer = self._get_scorer(selected_strategy) 

86 

87 # Score all entries 

88 if hasattr(scorer, "score_batch"): 

89 batch_scores = scorer.score_batch(entries, query) 

90 scored_entries = list(zip(batch_scores, entries, strict=True)) 

91 else: 

92 scored_entries = [(scorer.score(entry, query), entry) for entry in entries] 

93 

94 # Sort by score descending (highest scores first) 

95 scored_entries.sort(key=lambda x: x[0], reverse=True) 

96 

97 # Greedily select entries until budget is exhausted 

98 kept: list[MemoryEntry] = [] 

99 remaining_budget = token_budget 

100 

101 for _score, entry in scored_entries: 

102 # Count tokens in this entry's content 

103 entry_tokens = self.token_counter.count(entry.content) 

104 

105 # Check if entry fits in remaining budget 

106 if entry_tokens <= remaining_budget: 

107 kept.append(entry) 

108 remaining_budget -= entry_tokens 

109 # If we have no entries yet and this entry exceeds budget, 

110 # we still add it to avoid returning empty results for single large entries 

111 elif not kept and token_budget > 0: 

112 # Force include the first (highest-scored) entry only when budget > 0 

113 kept.append(entry) 

114 remaining_budget = 0 

115 break 

116 

117 pruned_count = len(entries) - len(kept) 

118 

119 logger.debug( 

120 "context_pruned", 

121 original_count=len(entries), 

122 kept_count=len(kept), 

123 pruned_count=pruned_count, 

124 strategy=selected_strategy, 

125 ) 

126 

127 return PruningResult( 

128 kept=kept, 

129 pruned_count=pruned_count, 

130 original_count=len(entries), 

131 token_budget=token_budget, 

132 strategy=selected_strategy, 

133 metadata={ 

134 "remaining_budget": remaining_budget, 

135 "scorer_type": type(scorer).__name__, 

136 }, 

137 ) 

138 

139 def _get_scorer(self, strategy: PruningStrategy) -> Any: 

140 """Get the scorer implementation for a given strategy. 

141 

142 Args: 

143 strategy: The pruning strategy to use. 

144 

145 Returns: 

146 A scorer instance matching the strategy. 

147 

148 Raises: 

149 ValueError: If the strategy is not recognized. 

150 """ 

151 if strategy == PruningStrategy.RELEVANCE: 

152 logger.warning( 

153 "pruning_strategy_relevance_not_implemented", 

154 fallback="recency", 

155 message="RELEVANCE strategy falls back to RecencyScorerImpl; embedding-based scoring not yet available", 

156 ) 

157 

158 # Use a registry dict instead of if/elif chains 

159 scorer_registry = { 

160 PruningStrategy.RECENCY: RecencyScorerImpl(), 

161 PruningStrategy.RELEVANCE: RecencyScorerImpl(), # Same as recency for now 

162 PruningStrategy.HYBRID: HybridScorerImpl(), 

163 } 

164 

165 if strategy not in scorer_registry: 

166 msg = f"Unknown pruning strategy: {strategy}" 

167 raise ValueError(msg) 

168 

169 return scorer_registry[strategy] 

170 

171 

172__all__ = ["DynamicContextPruner"]