Coverage for src / lexigram / contracts / ai / memory.py: 10%

1"""Memory system contracts: working, episodic, and semantic memory protocols.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6from typing import TYPE_CHECKING, Any 

7 

8from typing_extensions import Protocol, runtime_checkable 

9 

10if TYPE_CHECKING: 

11 from datetime import datetime 

12 

13 from lexigram.contracts.core import HealthCheckResult 

14 

15 

16@dataclass(frozen=True) 

17class MemoryEntry: 

18 """A single memory entry with temporal and importance metadata. 

19 

20 Attributes: 

21 id: Unique identifier for the entry 

22 owner_id: Owner scope for the entry (user, session, or composite) 

23 content: The actual content/text of the memory 

24 role: The role (e.g., "user", "assistant", "system") 

25 timestamp: When this entry was created 

26 importance: Score 0-1 indicating importance (default 0.5) 

27 metadata: Optional metadata dictionary 

28 embedding: Optional vector embedding of content 

29 """ 

30 

31 id: str 

32 owner_id: str 

33 content: str 

34 role: str 

35 timestamp: datetime 

36 importance: float = 0.5 

37 metadata: dict[str, Any] = field(default_factory=dict) 

38 embedding: list[float] | None = None 

39 

40 

41@dataclass(frozen=True) 

42class MemoryQuery: 

43 """Query parameters for searching memory. 

44 

45 Attributes: 

46 owner_id: Owner scope the query is restricted to (user, session, or composite) 

47 query: The search query text 

48 top_k: Number of results to return (default 10) 

49 min_relevance: Minimum relevance score threshold (default 0.0) 

50 recency_weight: Weight for recency in ranking (default 0.3) 

51 importance_weight: Weight for importance in ranking (default 0.3) 

52 relevance_weight: Weight for relevance in ranking (default 0.4) 

53 filters: Optional filters as key-value pairs 

54 time_range: Optional (start_datetime, end_datetime) range 

55 """ 

56 

57 owner_id: str 

58 query: str 

59 top_k: int = 10 

60 min_relevance: float = 0.0 

61 recency_weight: float = 0.3 

62 importance_weight: float = 0.3 

63 relevance_weight: float = 0.4 

64 filters: dict[str, Any] = field(default_factory=dict) 

65 time_range: tuple[datetime, datetime] | None = None 

66 

67 

68@dataclass(frozen=True) 

69class MemorySearchResult: 

70 """Result of a memory search. 

71 

72 Attributes: 

73 entry: The matching memory entry 

74 score: Relevance score (typically 0-1) 

75 source: Where this entry came from (e.g., "episodic", "semantic") 

76 """ 

77 

78 entry: MemoryEntry 

79 score: float 

80 source: str 

81 

82 

83@dataclass(frozen=True) 

84class ConsolidationResult: 

85 """Result of memory consolidation operation. 

86 

87 Attributes: 

88 entries_processed: Total entries examined 

89 entries_consolidated: Entries that were compressed/merged 

90 entries_pruned: Entries that were deleted 

91 entities_extracted: Number of entities extracted 

92 duration_ms: Time taken to consolidate in milliseconds 

93 """ 

94 

95 entries_processed: int 

96 entries_consolidated: int 

97 entries_pruned: int 

98 entities_extracted: int 

99 duration_ms: float 

100 

101 

102from lexigram.contracts.ai.exceptions import AIMemoryError 

103 

104# Memory Errors — AIMemoryError is the canonical base defined in ai/exceptions.py 

105 

106 

107class ConsolidationError(AIMemoryError): 

108 """Error raised during memory consolidation.""" 

109 

110 _code = "LEX_ERR_MEM_002" 

111 

112 

113class StorageError(AIMemoryError): 

114 """Error raised during memory storage operations.""" 

115 

116 _code = "LEX_ERR_MEM_003" 

117 

118 

119@runtime_checkable 

120class MemoryStoreProtocol(Protocol): 

121 """Protocol for storing and retrieving memory entries. 

122 

123 Implementations should provide persistent or semi-persistent storage 

124 for memory entries, supporting both sequential and search-based access. 

125 """ 

126 

127 async def store(self, entry: MemoryEntry) -> None: 

128 """Store a single memory entry. 

129 

130 Args: 

131 entry: The memory entry to store 

132 """ 

133 ... 

134 

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

136 """Search for memory entries based on a query. 

137 

138 Args: 

139 query: Query parameters 

140 

141 Returns: 

142 List of matching entries with scores, ordered by relevance 

143 """ 

144 ... 

145 

146 async def get_recent(self, n: int, owner_id: str) -> list[MemoryEntry]: 

147 """Get the N most recent entries owned by ``owner_id``. 

148 

149 Args: 

150 n: Number of entries to return 

151 owner_id: Owner scope; only this owner's entries are returned 

152 

153 Returns: 

154 List of recent entries in descending temporal order 

155 """ 

156 ... 

157 

158 async def delete(self, entry_id: str, owner_id: str) -> None: 

159 """Delete an entry by ID, scoped to ``owner_id``. 

160 

161 Args: 

162 entry_id: ID of the entry to delete 

163 owner_id: Owner scope; only this owner's entry may be deleted 

164 """ 

165 ... 

166 

167 async def clear(self, owner_id: str) -> None: 

168 """Clear all entries owned by ``owner_id`` from the store. 

169 

170 Args: 

171 owner_id: Owner scope; only this owner's entries are cleared 

172 """ 

173 ... 

174 

175 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

176 """Perform a lightweight connectivity check.""" 

177 ... 

178 

179 

180@runtime_checkable 

181class WorkingMemoryProtocol(Protocol): 

182 """Protocol for working memory that assembles context for the current request. 

183 

184 Working memory is the in-context window that gets passed to the LLM, 

185 assembled from episodic and semantic memory based on token budgets. 

186 """ 

187 

188 async def assemble( 

189 self, 

190 query: str, 

191 token_budget: int, 

192 *, 

193 owner_id: str, 

194 session_id: str | None = None, 

195 ) -> list[MemoryEntry]: 

196 """Assemble context window from available memory tiers. 

197 

198 Retrieves relevant entries from episodic and semantic memory 

199 and fits them within the token budget. 

200 

201 Args: 

202 query: The current query/prompt 

203 token_budget: Maximum tokens available for context 

204 owner_id: Owner scope; only this owner's memory is assembled 

205 session_id: Optional session scope for session-specific memory 

206 

207 Returns: 

208 Ordered list of memory entries for context 

209 """ 

210 ... 

211 

212 async def add(self, entry: MemoryEntry) -> None: 

213 """Add a new entry to the working memory stream. 

214 

215 Args: 

216 entry: The entry to add 

217 """ 

218 ... 

219 

220 async def get_context_entries(self) -> list[MemoryEntry]: 

221 """Get current assembled context entries. 

222 

223 Returns: 

224 List of entries currently in context window 

225 """ 

226 ... 

227 

228 async def flush(self) -> None: 

229 """Clear the current context assembly.""" 

230 ... 

231 

232 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

233 """Check if the working memory backend is reachable and healthy. 

234 

235 Args: 

236 timeout: Maximum seconds to wait for a response. 

237 

238 Returns: 

239 Health check result with status and details. 

240 """ 

241 ... 

242 

243 

244@runtime_checkable 

245class EpisodicMemoryProtocol(Protocol): 

246 """Protocol for episodic memory (conversation history with temporal grounding). 

247 

248 Episodic memory stores timestamped events and conversations, 

249 supporting hybrid retrieval based on recency and relevance. 

250 """ 

251 

252 async def record(self, entry: MemoryEntry) -> None: 

253 """Record a new episode/conversation turn. 

254 

255 Args: 

256 entry: The episode to record 

257 """ 

258 ... 

259 

260 async def recall(self, query: MemoryQuery) -> list[MemorySearchResult]: 

261 """Recall episodes based on relevance and recency. 

262 

263 Args: 

264 query: Query parameters with weighting preferences 

265 

266 Returns: 

267 Ranked list of matching episodes 

268 """ 

269 ... 

270 

271 async def forget(self, entry_id: str, owner_id: str) -> None: 

272 """Forget/delete a specific episode, scoped to ``owner_id``. 

273 

274 Args: 

275 entry_id: ID of the episode to forget 

276 owner_id: Owner scope; only this owner's episode may be forgotten 

277 """ 

278 ... 

279 

280 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

281 """Check if the episodic memory backend is reachable and healthy. 

282 

283 Args: 

284 timeout: Maximum seconds to wait for a response. 

285 

286 Returns: 

287 Health check result with status and details. 

288 """ 

289 ... 

290 

291 

292@runtime_checkable 

293class SemanticMemoryProtocol(Protocol): 

294 """Protocol for semantic memory (extracted facts and entities). 

295 

296 Semantic memory stores generalized knowledge in the form of facts, 

297 entities, and their relationships for long-term retention. 

298 """ 

299 

300 async def store_fact( 

301 self, subject: str, predicate: str, object_: str, confidence: float 

302 ) -> None: 

303 """Store a fact as a subject-predicate-object triple. 

304 

305 Args: 

306 subject: The subject entity 

307 predicate: The relationship type 

308 object_: The object entity 

309 confidence: Confidence score (0-1) for this fact 

310 """ 

311 ... 

312 

313 async def query_facts(self, subject: str) -> list[dict[str, Any]]: 

314 """Query facts by subject. 

315 

316 Args: 

317 subject: The subject entity to query 

318 

319 Returns: 

320 List of facts with this subject 

321 """ 

322 ... 

323 

324 async def get_entity_facts(self, entity: str) -> list[dict[str, Any]]: 

325 """Get all facts mentioning an entity (subject or object). 

326 

327 Args: 

328 entity: The entity name 

329 

330 Returns: 

331 List of facts involving this entity 

332 """ 

333 ... 

334 

335 async def update_fact(self, fact_id: str, confidence: float) -> None: 

336 """Update the confidence score of a fact. 

337 

338 Args: 

339 fact_id: ID of the fact to update 

340 confidence: New confidence score (0-1) 

341 """ 

342 ... 

343 

344 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

345 """Check if the semantic memory backend is reachable and healthy. 

346 

347 Args: 

348 timeout: Maximum seconds to wait for a response. 

349 

350 Returns: 

351 Health check result with status and details. 

352 """ 

353 ... 

354 

355 

356@runtime_checkable 

357class MemoryConsolidatorProtocol(Protocol): 

358 """Protocol for consolidating and optimizing memory. 

359 

360 Consolidation compresses old entries, deduplicates, extracts entities, 

361 and manages memory size to prevent unbounded growth. 

362 """ 

363 

364 async def consolidate(self, entries: list[MemoryEntry]) -> ConsolidationResult: 

365 """Consolidate a batch of memory entries. 

366 

367 Applies compression, deduplication, entity extraction, and pruning 

368 strategies to optimize memory storage and retrieval. 

369 

370 Args: 

371 entries: Entries to consolidate 

372 

373 Returns: 

374 Consolidation result with statistics 

375 """ 

376 ... 

377 

378 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

379 """Check if the memory consolidator backend is reachable and healthy. 

380 

381 Args: 

382 timeout: Maximum seconds to wait for a response. 

383 

384 Returns: 

385 Health check result with status and details. 

386 """ 

387 ... 

388 

389 

390from typing import Any, Protocol, runtime_checkable 

391 

392 

393@runtime_checkable 

394class MemoryProtocol(Protocol): 

395 """Protocol for conversation memory (G-02 parity).""" 

396 

397 async def add_message(self, message: Any) -> None: 

398 """Add a message to memory.""" 

399 ... 

400 

401 async def get_messages(self) -> list[Any]: 

402 """Get all messages from memory.""" 

403 ... 

404 

405 async def clear(self) -> None: 

406 """Clear all messages from memory.""" 

407 ... 

408 

409 

410class ConversationMemory: 

411 """Simple conversation memory that keeps all messages. 

412 

413 Like LangChain's ConversationBufferMemory. 

414 """ 

415 

416 def __init__(self, max_messages: int | None = None) -> None: 

417 self._messages: list[Any] = [] 

418 self.max_messages = max_messages 

419 

420 def add_message(self, message: Any) -> None: 

421 self._messages.append(message) 

422 if self.max_messages and len(self._messages) > self.max_messages: 

423 self._messages.pop(0) 

424 

425 def get_messages(self) -> list[Any]: 

426 return list(self._messages) 

427 

428 def clear(self) -> None: 

429 self._messages.clear() 

430 

431 async def aadd_message(self, message: Any) -> None: 

432 self.add_message(message) 

433 

434 async def aget_messages(self) -> list[Any]: 

435 return self.get_messages() 

436 

437 async def aclear(self) -> None: 

438 self.clear() 

439 

440 

441class WindowMemory: 

442 """Memory that keeps only the last N messages. 

443 

444 Like LangChain's ConversationBufferWindowMemory. 

445 """ 

446 

447 def __init__(self, window_size: int = 5) -> None: 

448 self.window_size = window_size 

449 self._messages: list[Any] = [] 

450 

451 def add_message(self, message: Any) -> None: 

452 self._messages.append(message) 

453 if len(self._messages) > self.window_size: 

454 self._messages.pop(0) 

455 

456 def get_messages(self) -> list[Any]: 

457 return list(self._messages) 

458 

459 def clear(self) -> None: 

460 self._messages.clear() 

461 

462 async def aadd_message(self, message: Any) -> None: 

463 self.add_message(message) 

464 

465 async def aget_messages(self) -> list[Any]: 

466 return self.get_messages() 

467 

468 async def aclear(self) -> None: 

469 self.clear() 

470 

471 

472__all__ = [ 

473 "ConsolidationError", 

474 "ConsolidationResult", 

475 "ConversationMemory", 

476 "EpisodicMemoryProtocol", 

477 "MemoryConsolidatorProtocol", 

478 "MemoryEntry", 

479 "MemoryProtocol", 

480 "MemoryQuery", 

481 "MemorySearchResult", 

482 "MemoryStoreProtocol", 

483 "SemanticMemoryProtocol", 

484 "StorageError", 

485 "WindowMemory", 

486 "WorkingMemoryProtocol", 

487]