Coverage for src\core\memory.py: 0%
215 statements
« prev ^ index » next coverage.py v7.3.4, created at 2026-04-21 14:54 +0800
1"""
2本地记忆系统
3实现 AI 记忆的存储、检索和管理
4"""
6import json
7import time
8import uuid
9from datetime import datetime, timezone
10from enum import Enum
11from typing import Dict, List, Optional, Any, Union
12from dataclasses import dataclass, asdict, field
13from pathlib import Path
15import aiofiles
16from pydantic import BaseModel, Field
19class MemoryType(str, Enum):
20 """记忆类型"""
21 CONVERSATION = "conversation" # 对话记忆
22 KNOWLEDGE = "knowledge" # 知识记忆
23 TASK = "task" # 任务记忆
24 PREFERENCE = "preference" # 偏好记忆
25 CONTEXT = "context" # 上下文记忆
26 SYSTEM = "system" # 系统记忆
29class MemoryPriority(int, Enum):
30 """记忆优先级"""
31 LOW = 1
32 MEDIUM = 2
33 HIGH = 3
34 CRITICAL = 4
37@dataclass
38class MemoryMetadata:
39 """记忆元数据"""
40 source: str = "" # 来源（设备ID或用户ID）
41 confidence: float = 1.0 # 置信度（0.0-1.0）
42 access_count: int = 0 # 访问次数
43 last_accessed: float = 0.0 # 最后访问时间戳
44 created_at: float = field(default_factory=time.time) # 创建时间戳
45 tags: List[str] = field(default_factory=list) # 标签
46 expires_at: Optional[float] = None # 过期时间戳（None表示永不过期）
49class MemoryItem(BaseModel):
50 """记忆项"""
52 id: str = Field(default_factory=lambda: str(uuid.uuid4()))
53 type: MemoryType
54 content: Dict[str, Any]
55 priority: MemoryPriority = MemoryPriority.MEDIUM
56 metadata: MemoryMetadata = Field(default_factory=MemoryMetadata)
58 # 时间戳
59 created_at: float = Field(default_factory=time.time)
60 updated_at: float = Field(default_factory=time.time)
62 # 关联信息
63 related_memories: List[str] = Field(default_factory=list) # 关联记忆ID
64 device_id: Optional[str] = None # 创建此记忆的设备ID
66 class Config:
67 use_enum_values = True
70class MemoryQuery(BaseModel):
71 """记忆查询条件"""
73 type: Optional[MemoryType] = None
74 content_key: Optional[str] = None
75 content_value: Optional[Any] = None
76 min_priority: Optional[MemoryPriority] = None
77 tags: Optional[List[str]] = None
78 source: Optional[str] = None
79 device_id: Optional[str] = None
80 created_after: Optional[float] = None
81 created_before: Optional[float] = None
82 limit: int = 100
83 offset: int = 0
86class MemorySystem:
87 """记忆系统"""
89 def __init__(self, storage_path: Optional[Path] = None):
90 """
91 初始化记忆系统
93 Args:
94 storage_path: 存储路径，如果为None则使用默认路径
95 """
96 if storage_path is None:
97 storage_path = Path.home() / ".pao" / "memories"
99 self.storage_path = Path(storage_path)
100 self.memories: Dict[str, MemoryItem] = {}
102 # 创建存储目录
103 self.storage_path.mkdir(parents=True, exist_ok=True)
105 # 索引
106 self.type_index: Dict[MemoryType, List[str]] = {}
107 self.tag_index: Dict[str, List[str]] = {}
108 self.device_index: Dict[str, List[str]] = {}
110 async def load(self) -> None:
111 """从存储加载所有记忆"""
112 try:
113 memory_files = list(self.storage_path.glob("*.json"))
115 for file_path in memory_files:
116 try:
117 async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
118 content = await f.read()
120 data = json.loads(content)
121 memory = MemoryItem(**data)
123 # 添加到内存
124 self.memories[memory.id] = memory
126 # 更新索引
127 self._add_to_indices(memory)
129 except Exception as e:
130 print(f"加载记忆文件 {file_path} 失败: {e}")
132 except Exception as e:
133 print(f"加载记忆系统失败: {e}")
135 async def save(self, memory: MemoryItem) -> None:
136 """
137 保存记忆
139 Args:
140 memory: 要保存的记忆
141 """
142 try:
143 # 更新时间戳
144 memory.updated_at = time.time()
146 # 保存到内存
147 self.memories[memory.id] = memory
149 # 更新索引
150 self._add_to_indices(memory)
152 # 保存到文件
153 file_path = self.storage_path / f"{memory.id}.json"
155 async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
156 await f.write(memory.model_dump_json(indent=2))
158 # 更新访问统计
159 memory.metadata.access_count += 1
160 memory.metadata.last_accessed = time.time()
162 except Exception as e:
163 print(f"保存记忆失败: {e}")
164 raise
166 async def get(self, memory_id: str) -> Optional[MemoryItem]:
167 """
168 获取记忆
170 Args:
171 memory_id: 记忆ID
173 Returns:
174 记忆项，如果不存在则返回None
175 """
176 memory = self.memories.get(memory_id)
178 if memory:
179 # 更新访问统计
180 memory.metadata.access_count += 1
181 memory.metadata.last_accessed = time.time()
183 # 异步保存更新
184 asyncio.create_task(self._update_memory_file(memory))
186 return memory
188 async def query(self, query: MemoryQuery) -> List[MemoryItem]:
189 """
190 查询记忆
192 Args:
193 query: 查询条件
195 Returns:
196 匹配的记忆列表
197 """
198 results = []
200 for memory in self.memories.values():
201 if self._matches_query(memory, query):
202 results.append(memory)
204 # 按创建时间排序（最新的在前面）
205 results.sort(key=lambda m: m.created_at, reverse=True)
207 # 应用分页
208 start = query.offset
209 end = start + query.limit
210 return results[start:end]
212 async def delete(self, memory_id: str) -> bool:
213 """
214 删除记忆
216 Args:
217 memory_id: 记忆ID
219 Returns:
220 是否成功删除
221 """
222 try:
223 memory = self.memories.get(memory_id)
224 if not memory:
225 return False
227 # 从内存中删除
228 del self.memories[memory_id]
230 # 从索引中删除
231 self._remove_from_indices(memory)
233 # 删除文件
234 file_path = self.storage_path / f"{memory_id}.json"
235 if file_path.exists():
236 file_path.unlink()
238 return True
240 except Exception as e:
241 print(f"删除记忆失败: {e}")
242 return False
244 async def update(self, memory_id: str, updates: Dict[str, Any]) -> Optional[MemoryItem]:
245 """
246 更新记忆
248 Args:
249 memory_id: 记忆ID
250 updates: 更新字段
252 Returns:
253 更新后的记忆，如果不存在则返回None
254 """
255 memory = await self.get(memory_id)
256 if not memory:
257 return None
259 # 应用更新
260 for key, value in updates.items():
261 if hasattr(memory, key):
262 setattr(memory, key, value)
264 memory.updated_at = time.time()
266 # 保存更新
267 await self.save(memory)
269 return memory
271 async def cleanup_expired(self) -> int:
272 """
273 清理过期记忆
275 Returns:
276 清理的记忆数量
277 """
278 current_time = time.time()
279 expired_ids = []
281 for memory_id, memory in self.memories.items():
282 if memory.metadata.expires_at and memory.metadata.expires_at < current_time:
283 expired_ids.append(memory_id)
285 # 删除过期记忆
286 deleted_count = 0
287 for memory_id in expired_ids:
288 if await self.delete(memory_id):
289 deleted_count += 1
291 return deleted_count
293 def get_stats(self) -> Dict[str, Any]:
294 """
295 获取统计信息
297 Returns:
298 统计信息字典
299 """
300 return {
301 "total_memories": len(self.memories),
302 "by_type": {
303 type_.value: len(memories)
304 for type_, memories in self.type_index.items()
305 },
306 "by_device": {
307 device_id: len(memories)
308 for device_id, memories in self.device_index.items()
309 },
310 "oldest_memory": min(
311 [m.created_at for m in self.memories.values()] or [0]
312 ),
313 "newest_memory": max(
314 [m.created_at for m in self.memories.values()] or [0]
315 )
316 }
318 async def add_memory(self, memory_type: str, content: Any, priority: int = 2) -> str:
319 """添加记忆"""
320 mem_type = MemoryType.CONVERSATION
321 for mt in MemoryType:
322 if mt.value == memory_type:
323 mem_type = mt
324 break
325 # 确保priority是有效的枚举值
326 valid_priority = max(1, min(4, priority))
327 memory = MemoryItem(
328 id=str(uuid.uuid4()),
329 type=mem_type,
330 content=content,
331 priority=valid_priority,
332 metadata=MemoryMetadata()
333 )
334 await self.save(memory)
335 return memory.id
337 async def search_memories(self, query: str, limit: int = 10) -> list:
338 """搜索记忆"""
339 memory_query = MemoryQuery(keyword=query, limit=limit)
340 results = await self.query(memory_query)
341 return [m.content for m in results]
343 def _add_to_indices(self, memory: MemoryItem) -> None:
344 """添加到索引"""
345 # 类型索引
346 if memory.type not in self.type_index:
347 self.type_index[memory.type] = []
348 if memory.id not in self.type_index[memory.type]:
349 self.type_index[memory.type].append(memory.id)
351 # 标签索引
352 for tag in memory.metadata.tags:
353 if tag not in self.tag_index:
354 self.tag_index[tag] = []
355 if memory.id not in self.tag_index[tag]:
356 self.tag_index[tag].append(memory.id)
358 # 设备索引
359 if memory.device_id:
360 if memory.device_id not in self.device_index:
361 self.device_index[memory.device_id] = []
362 if memory.id not in self.device_index[memory.device_id]:
363 self.device_index[memory.device_id].append(memory.id)
365 def _remove_from_indices(self, memory: MemoryItem) -> None:
366 """从索引中删除"""
367 # 类型索引
368 if memory.type in self.type_index:
369 if memory.id in self.type_index[memory.type]:
370 self.type_index[memory.type].remove(memory.id)
372 # 标签索引
373 for tag in memory.metadata.tags:
374 if tag in self.tag_index and memory.id in self.tag_index[tag]:
375 self.tag_index[tag].remove(memory.id)
377 # 设备索引
378 if memory.device_id and memory.device_id in self.device_index:
379 if memory.id in self.device_index[memory.device_id]:
380 self.device_index[memory.device_id].remove(memory.id)
382 def _matches_query(self, memory: MemoryItem, query: MemoryQuery) -> bool:
383 """检查记忆是否匹配查询条件"""
384 if query.type and memory.type != query.type:
385 return False
387 if query.content_key and query.content_value:
388 if query.content_key not in memory.content:
389 return False
390 if memory.content[query.content_key] != query.content_value:
391 return False
393 if query.min_priority and memory.priority.value < query.min_priority.value:
394 return False
396 if query.tags:
397 if not all(tag in memory.metadata.tags for tag in query.tags):
398 return False
400 if query.source and memory.metadata.source != query.source:
401 return False
403 if query.device_id and memory.device_id != query.device_id:
404 return False
406 if query.created_after and memory.created_at < query.created_after:
407 return False
409 if query.created_before and memory.created_at > query.created_before:
410 return False
412 return True
414 async def _update_memory_file(self, memory: MemoryItem) -> None:
415 """异步更新记忆文件"""
416 try:
417 file_path = self.storage_path / f"{memory.id}.json"
418 async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
419 await f.write(memory.model_dump_json(indent=2))
420 except Exception as e:
421 print(f"异步更新记忆文件失败: {e}")
424import asyncio