Coverage for merco/memory/recall.py: 81%

1"""Memory recall interfaces — ABC + FTS5 + Memory + Hybrid recallers.""" 

2 

3from __future__ import annotations 

4 

5import json 

6import logging 

7from abc import ABC, abstractmethod 

8from dataclasses import dataclass, field 

9from typing import TYPE_CHECKING 

10 

11if TYPE_CHECKING: 

12 from .session_search import SessionSearch 

13 

14from .store import MemoryStore 

15 

16 

17@dataclass 

18class RecallResult: 

19 """A single recall result from any recaller.""" 

20 

21 snippet: str 

22 session_title: str = "" 

23 score: float = 0.0 

24 source: str = "memory" # "fts5" | "memory" 

25 

26 

27class BaseRecaller(ABC): 

28 """Abstract base for all recallers.""" 

29 

30 name: str = "" 

31 

32 @abstractmethod 

33 async def recall(self, query: str, limit: int = 10) -> list[RecallResult]: 

34 """Return recall results for the given query.""" 

35 ... 

36 

37 

38class FTS5Recaller(BaseRecaller): 

39 """Full-text search recaller backed by SessionSearch.""" 

40 

41 name = "fts5" 

42 

43 def __init__(self, session_search: SessionSearch) -> None: 

44 self._search = session_search 

45 

46 async def recall(self, query: str, limit: int = 10) -> list[RecallResult]: 

47 raw = self._search.search(query, limit=limit) 

48 total = len(raw) 

49 results: list[RecallResult] = [] 

50 for i, item in enumerate(raw): 

51 score = (total - i) / max(total, 1) if total > 0 else 0.0 

52 results.append( 

53 RecallResult( 

54 snippet=item["snippet"], 

55 session_title=item.get("session_title", ""), 

56 score=round(score, 4), 

57 source="fts5", 

58 ) 

59 ) 

60 return results 

61 

62 

63class MemoryRecaller(BaseRecaller): 

64 """Key-value memory store recaller backed by MemoryStore.""" 

65 

66 name = "memory" 

67 

68 def __init__(self, memory_store: MemoryStore) -> None: 

69 self._store = memory_store 

70 

71 async def recall(self, query: str, limit: int = 10) -> list[RecallResult]: 

72 raw = self._store.search(query)[:limit] 

73 total = len(raw) 

74 results: list[RecallResult] = [] 

75 for i, item in enumerate(raw): 

76 score = (total - i) / max(total, 1) if total > 0 else 0.0 

77 snippet = item.get("value", "") 

78 if isinstance(snippet, (dict, list)): 

79 snippet = json.dumps(snippet, ensure_ascii=False) 

80 else: 

81 snippet = str(snippet) 

82 results.append( 

83 RecallResult( 

84 snippet=snippet, 

85 session_title=item.get("key", ""), 

86 score=round(score, 4), 

87 source="memory", 

88 ) 

89 ) 

90 return results 

91 

92 

93class HybridRecaller(BaseRecaller): 

94 """Aggregates multiple recallers, deduplicates, and limits by chars. 

95 

96 Caches results by (query, limit) so repeated calls return instantly. 

97 """ 

98 

99 name = "hybrid" 

100 

101 def __init__( 

102 self, 

103 recallers: list[BaseRecaller] | None = None, 

104 limit: int = 3, 

105 max_chars: int = 300, 

106 ) -> None: 

107 self._recallers: list[BaseRecaller] = list(recallers) if recallers else [] 

108 self._limit = limit 

109 self._max_chars = max_chars 

110 self._cache: dict[tuple[str, int], list[RecallResult]] = {} 

111 

112 def add(self, recaller: BaseRecaller) -> HybridRecaller: 

113 """Add a recaller (fluent interface).""" 

114 self._recallers.append(recaller) 

115 return self 

116 

117 async def recall(self, query: str, limit: int = 10) -> list[RecallResult]: 

118 cache_key = (query, limit) 

119 if cache_key in self._cache: 

120 return self._cache[cache_key] 

121 

122 # Gather from all recallers 

123 all_results: list[RecallResult] = [] 

124 for rec in self._recallers: 

125 try: 

126 batch = await rec.recall(query, limit=limit) 

127 all_results.extend(batch) 

128 except Exception: 

129 # Swallow individual recaller errors so one failure 

130 # doesn't break the whole hybrid pipeline. 

131 logger = logging.getLogger(__name__) 

132 logger.warning("Recaller %s failed for query %r", rec.name, query, exc_info=True) 

133 continue 

134 

135 # Sort by score descending 

136 all_results.sort(key=lambda r: r.score, reverse=True) 

137 

138 # Deduplicate by first 80 chars of snippet 

139 seen: set[str] = set() 

140 deduped: list[RecallResult] = [] 

141 for r in all_results: 

142 prefix = r.snippet[:80] 

143 if prefix not in seen: 

144 seen.add(prefix) 

145 deduped.append(r) 

146 

147 # Truncate by max_chars (cumulative snippet length) 

148 total_chars = 0 

149 truncated: list[RecallResult] = [] 

150 for r in deduped: 

151 total_chars += len(r.snippet) 

152 if total_chars > self._max_chars and truncated: 

153 break 

154 truncated.append(r) 

155 

156 # Apply top-level limit (use _limit as a cap, not a floor) 

157 final = truncated[: min(limit, self._limit)] 

158 

159 self._cache[cache_key] = final 

160 return final 

161 

162 

163# --------------------------------------------------------------------------- 

164# Backward-compatible synchronous MemoryRecall (existing API) 

165# --------------------------------------------------------------------------- 

166 

167class MemoryRecall: 

168 """Legacy memory recall — kept for backward compatibility. 

169 

170 Prefer MemoryRecaller (async) for new code. 

171 """ 

172 

173 def __init__(self, store: MemoryStore) -> None: 

174 self.store = store 

175 

176 def recall(self, context: str, max_results: int = 5) -> list[dict]: 

177 """根据上下文召回相关记忆""" 

178 results = self.store.search(context) 

179 return results[:max_results] 

180 

181 def recall_by_tag(self, tag: str) -> list[dict]: 

182 """按标签召回记忆""" 

183 keys = self.store.list_keys(tag=tag) 

184 memories = [] 

185 for key in keys: 

186 memory = self.store.load(key) 

187 if memory: 

188 memories.append(memory) 

189 return memories 

190 

191 def get_relevant_context(self, query: str) -> str: 

192 """获取相关上下文文本""" 

193 memories = self.recall(query) 

194 if not memories: 

195 return "" 

196 

197 context_parts = [] 

198 for m in memories: 

199 context_parts.append(f"[{m.get('key', '?')}]: {m.get('value', '')}") 

200 

201 return "\n".join(context_parts)