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

1"""Memory retriever — unified retrieval across multiple MemoryStoreProtocol sources.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any 

6 

7from lexigram.ai.memory.retrieval.ranking import RelevanceRanker 

8from lexigram.logging import ( 

9 get_logger, 

10) 

11 

12if TYPE_CHECKING: 

13 from lexigram.contracts.ai.memory import ( 

14 MemoryQuery, 

15 MemorySearchResult, 

16 MemoryStoreProtocol, 

17 ) 

18 

19logger = get_logger(__name__) 

20 

21 

22class MemoryRetriever: 

23 """Queries one or more MemoryStoreProtocol backends and merges results. 

24 

25 Results from each backend are pooled, deduplicated by entry ID, and 

26 re-ranked by the configured RelevanceRanker. Retrieval counts are 

27 tracked per entry to enable downstream analytics and relevance decay. 

28 """ 

29 

30 def __init__( 

31 self, 

32 sources: list[MemoryStoreProtocol], 

33 ranker: RelevanceRanker | None = None, 

34 ) -> None: 

35 """Initialise the retriever. 

36 

37 Args: 

38 sources: Memory store backends to query in parallel. 

39 ranker: Optional ranker. Defaults to a new ``RelevanceRanker``. 

40 """ 

41 self._sources = sources 

42 self._ranker = ranker or RelevanceRanker() 

43 self._retrieval_counts: dict[str, int] = {} 

44 

45 async def retrieve(self, query: MemoryQuery) -> list[MemorySearchResult]: 

46 """Retrieve and merge results from all configured sources. 

47 

48 Args: 

49 query: Search query with weights and filters. 

50 

51 Returns: 

52 De-duplicated, re-ranked results capped at ``query.top_k``. 

53 """ 

54 import asyncio 

55 

56 tasks = [source.retrieve(query) for source in self._sources] 

57 all_results_nested = await asyncio.gather(*tasks, return_exceptions=False) 

58 

59 merged: dict[str, MemorySearchResult] = {} 

60 for batch in all_results_nested: 

61 for result in batch: 

62 # Keep highest raw score if the same entry appears in multiple sources 

63 existing = merged.get(result.entry.id) 

64 if existing is None or result.score > existing.score: 

65 merged[result.entry.id] = result 

66 

67 pooled = list(merged.values()) 

68 ranked = self._ranker.top_k(pooled, query) 

69 

70 # Track retrieval counts for each returned entry 

71 for result in ranked: 

72 entry_id = result.entry.id 

73 self._retrieval_counts[entry_id] = ( 

74 self._retrieval_counts.get(entry_id, 0) + 1 

75 ) 

76 

77 return ranked 

78 

79 def add_source(self, source: MemoryStoreProtocol) -> None: 

80 """Register a new backend source. 

81 

82 Args: 

83 source: Additional MemoryStoreProtocol to query. 

84 """ 

85 self._sources.append(source) 

86 

87 def get_retrieval_count(self, entry_id: str) -> int: 

88 """Return how many times a specific entry has been retrieved. 

89 

90 Args: 

91 entry_id: The memory entry ID. 

92 

93 Returns: 

94 Number of times the entry appeared in retrieval results. 

95 """ 

96 return self._retrieval_counts.get(entry_id, 0) 

97 

98 def get_retrieval_stats(self) -> dict[str, Any]: 

99 """Return aggregate retrieval statistics. 

100 

101 Returns: 

102 Dictionary with total retrievals, unique entries, and top-10 

103 most-retrieved entries. 

104 """ 

105 total = sum(self._retrieval_counts.values()) 

106 unique = len(self._retrieval_counts) 

107 sorted_entries = sorted( 

108 self._retrieval_counts.items(), key=lambda x: x[1], reverse=True 

109 ) 

110 return { 

111 "total_retrievals": total, 

112 "unique_entries": unique, 

113 "top_entries": sorted_entries[:10], 

114 } 

115 

116 def reset_stats(self) -> None: 

117 """Clear all retrieval tracking data.""" 

118 self._retrieval_counts.clear() 

119 

120 

121__all__ = ["MemoryRetriever"]