Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-memory/src/lexigram/ai/memory/retrieval/prune.py: 37%

1"""Memory pruner — removes stale or low-relevance entries from memory stores.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from datetime import UTC, datetime 

7from typing import TYPE_CHECKING, Any 

8 

9from lexigram.logging import ( 

10 get_logger, 

11) 

12 

13if TYPE_CHECKING: 

14 from lexigram.contracts.ai.memory import MemoryStoreProtocol 

15 

16logger = get_logger(__name__) 

17 

18 

19@dataclass 

20class PruneResult: 

21 """Outcome of a prune operation. 

22 

23 Attributes: 

24 pruned_count: Number of entries removed. 

25 remaining_count: Number of entries kept. 

26 metadata: Extra details about the prune pass. 

27 """ 

28 

29 pruned_count: int 

30 remaining_count: int 

31 metadata: dict[str, Any] 

32 

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

34 """Convert to dictionary for serialization.""" 

35 return { 

36 "pruned_count": self.pruned_count, 

37 "remaining_count": self.remaining_count, 

38 "metadata": self.metadata, 

39 } 

40 

41 

42class MemoryPruner: 

43 """Removes stale or low-relevance entries from a memory store. 

44 

45 Supports two complementary criteria: 

46 

47 * **Age-based**: entries older than ``max_age_hours`` are pruned. 

48 * **Threshold-based**: entries whose ``importance`` score falls below 

49 ``importance_threshold`` are pruned. 

50 

51 Both criteria can be combined — an entry is pruned if it matches 

52 *either* condition. 

53 

54 Example:: 

55 

56 pruner = MemoryPruner(store) 

57 result = await pruner.prune( 

58 importance_threshold=0.2, 

59 max_age_hours=72, 

60 ) 

61 print(f"Pruned {result.pruned_count} entries") 

62 

63 Args: 

64 store: Memory store to prune entries from. 

65 """ 

66 

67 def __init__(self, store: MemoryStoreProtocol) -> None: 

68 """Initialise the pruner. 

69 

70 Args: 

71 store: The target memory store backend. 

72 """ 

73 self._store = store 

74 

75 async def prune( 

76 self, 

77 owner_id: str, 

78 importance_threshold: float = 0.1, 

79 max_age_hours: float = 0.0, 

80 dry_run: bool = False, 

81 ) -> PruneResult: 

82 """Prune entries that fall below the given thresholds. 

83 

84 An entry is eligible for pruning if it matches *either* condition: 

85 - ``importance < importance_threshold`` 

86 - ``age > max_age_hours`` (if ``max_age_hours > 0``) 

87 

88 Args: 

89 owner_id: Owner scope restricting pruning to one owner's entries. 

90 importance_threshold: Prune entries below this importance score. 

91 max_age_hours: Prune entries older than this (0 = disabled). 

92 dry_run: If ``True``, report what would be pruned without deleting. 

93 

94 Returns: 

95 ``PruneResult`` with counts and metadata. 

96 """ 

97 from lexigram.contracts.ai.memory import MemoryQuery 

98 

99 # Retrieve all entries with a broad query 

100 all_results = await self._store.retrieve( 

101 MemoryQuery( 

102 owner_id=owner_id, 

103 query="", 

104 top_k=10_000, 

105 recency_weight=0.0, 

106 importance_weight=0.0, 

107 relevance_weight=1.0, 

108 ) 

109 ) 

110 

111 now = datetime.now(UTC) 

112 to_prune: list[str] = [] 

113 age_pruned = 0 

114 importance_pruned = 0 

115 

116 for result in all_results: 

117 entry = result.entry 

118 prune_this = False 

119 

120 # Age check 

121 if max_age_hours > 0: 

122 age_hours = (now - entry.timestamp).total_seconds() / 3600 

123 if age_hours > max_age_hours: 

124 prune_this = True 

125 age_pruned += 1 

126 

127 # Importance check 

128 if entry.importance < importance_threshold: 

129 prune_this = True 

130 importance_pruned += 1 

131 

132 if prune_this: 

133 to_prune.append(entry.id) 

134 

135 if not dry_run and to_prune: 

136 for entry_id in to_prune: 

137 await self._store.delete(entry_id, owner_id) 

138 

139 remaining = len(all_results) - len(to_prune) 

140 

141 logger.info( 

142 "memory_pruned", 

143 pruned=len(to_prune), 

144 remaining=remaining, 

145 dry_run=dry_run, 

146 age_pruned=age_pruned, 

147 importance_pruned=importance_pruned, 

148 ) 

149 

150 return PruneResult( 

151 pruned_count=len(to_prune), 

152 remaining_count=remaining, 

153 metadata={ 

154 "importance_threshold": importance_threshold, 

155 "max_age_hours": max_age_hours, 

156 "dry_run": dry_run, 

157 "age_pruned": age_pruned, 

158 "importance_pruned": importance_pruned, 

159 }, 

160 ) 

161 

162 

163__all__ = ["MemoryPruner", "PruneResult"]