Coverage for src / agent / memory / persistence.py: 100%

1# Copyright 2025-2026 Microsoft Corporation 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Memory persistence utilities. 

16 

17This module provides serialization and persistence for memory state. 

18""" 

19 

20import json 

21import logging 

22from datetime import datetime 

23from pathlib import Path 

24 

25logger = logging.getLogger(__name__) 

26 

27 

28class MemoryPersistence: 

29 """Helper class for memory serialization and persistence. 

30 

31 Provides JSON-based serialization for memory state with version 

32 compatibility handling. Follows ThreadPersistence patterns. 

33 

34 Example: 

35 >>> persistence = MemoryPersistence() 

36 >>> await persistence.save(memory_data, Path("memory.json")) 

37 """ 

38 

39 VERSION = "1.0" 

40 

41 def __init__(self, storage_dir: Path | None = None): 

42 """Initialize memory persistence helper. 

43 

44 Args: 

45 storage_dir: Directory for memory storage (default: ~/.osdu-agent/memory) 

46 """ 

47 if storage_dir is None: 

48 storage_dir = Path.home() / ".osdu-agent" / "memory" 

49 

50 self.storage_dir = Path(storage_dir) 

51 self.storage_dir.mkdir(parents=True, exist_ok=True) 

52 

53 logger.debug(f"Memory persistence initialized: {self.storage_dir}") 

54 

55 async def save(self, memory_data: list[dict], file_path: Path) -> None: 

56 """Save memory state to file. 

57 

58 Args: 

59 memory_data: List of memory entries to serialize 

60 file_path: Path to save file 

61 

62 Raises: 

63 Exception: If serialization or save fails 

64 

65 Example: 

66 >>> memories = [{"role": "user", "content": "Hello"}] 

67 >>> await persistence.save(memories, Path("session-1-memory.json")) 

68 """ 

69 try: 

70 # Build memory state with metadata 

71 state = { 

72 "version": self.VERSION, 

73 "saved_at": datetime.now().isoformat(), 

74 "memory_count": len(memory_data), 

75 "memories": memory_data, 

76 } 

77 

78 # Ensure parent directory exists 

79 file_path.parent.mkdir(parents=True, exist_ok=True) 

80 

81 # Save to file 

82 with open(file_path, "w") as f: 

83 json.dump(state, f, indent=2) 

84 

85 logger.info(f"Saved {len(memory_data)} memories to {file_path}") 

86 

87 except Exception as e: 

88 logger.error(f"Failed to save memory state: {e}") 

89 raise 

90 

91 async def load(self, file_path: Path) -> list[dict] | None: 

92 """Load memory state from file. 

93 

94 Args: 

95 file_path: Path to memory file 

96 

97 Returns: 

98 List of memory entries or None if file doesn't exist 

99 

100 Raises: 

101 Exception: If deserialization fails 

102 

103 Example: 

104 >>> memories = await persistence.load(Path("session-1-memory.json")) 

105 """ 

106 if not file_path.exists(): 

107 logger.debug(f"Memory file not found: {file_path}") 

108 return None 

109 

110 try: 

111 with open(file_path) as f: 

112 state = json.load(f) 

113 

114 # Version compatibility check 

115 version = state.get("version", "unknown") 

116 if version != self.VERSION: 

117 logger.warning(f"Memory version mismatch: {version} != {self.VERSION}") 

118 # Future: Handle version migrations here 

119 

120 memories: list[dict] = state.get("memories", []) 

121 logger.info(f"Loaded {len(memories)} memories from {file_path}") 

122 

123 return memories 

124 

125 except Exception as e: 

126 logger.error(f"Failed to load memory state: {e}") 

127 raise 

128 

129 def get_memory_path(self, session_name: str) -> Path: 

130 """Get memory file path for a session. 

131 

132 Args: 

133 session_name: Name of the session 

134 

135 Returns: 

136 Path to memory file for the session 

137 

138 Example: 

139 >>> path = persistence.get_memory_path("session-1") 

140 >>> str(path) 

141 '~/.osdu-agent/memory/session-1-memory.json' 

142 """ 

143 return self.storage_dir / f"{session_name}-memory.json"