Coverage for agentos/memory/working.py: 47%

1""" 

2工作记忆 — 当前会话上下文。 

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass, field 

8 

9 

10@dataclass 

11class MemoryItem: 

12 """工作记忆项。""" 

13 key: str 

14 content: str 

15 ttl: str = "session" # session | days | permanent 

16 metadata: dict = field(default_factory=dict) 

17 

18 

19class WorkingMemory: 

20 """工作记忆 — 当前会话内有效，会话结束即销毁。""" 

21 

22 def __init__(self, max_items: int = 100): 

23 self.max_items = max_items 

24 self._items: dict[str, MemoryItem] = {} 

25 

26 def add(self, item: MemoryItem): 

27 self._items[item.key] = item 

28 if len(self._items) > self.max_items: 

29 oldest = next(iter(self._items)) 

30 del self._items[oldest] 

31 

32 def get(self, key: str) -> MemoryItem | None: 

33 return self._items.get(key) 

34 

35 def search(self, query: str, limit: int = 5) -> list[MemoryItem]: 

36 """简单关键词匹配。""" 

37 results = [] 

38 for item in self._items.values(): 

39 if query.lower() in item.content.lower(): 

40 results.append(item) 

41 return results[:limit] 

42 

43 def clear(self): 

44 self._items.clear() 

45 

46 # ── Persistence (v1.14.9) ──────────────── 

47 

48 def get_state(self) -> dict[str, Any]: 

49 """Export working memory state for persistence.""" 

50 return { 

51 "max_items": self.max_items, 

52 "items": { 

53 key: { 

54 "key": item.key, 

55 "content": item.content, 

56 "ttl": item.ttl, 

57 "metadata": item.metadata, 

58 } 

59 for key, item in self._items.items() 

60 }, 

61 } 

62 

63 def restore_state(self, state: dict[str, Any]) -> None: 

64 """Restore working memory from a persisted snapshot.""" 

65 self.max_items = state.get("max_items", self.max_items) 

66 self._items.clear() 

67 for key, item_data in state.get("items", {}).items(): 

68 self._items[key] = MemoryItem( 

69 key=item_data.get("key", key), 

70 content=item_data.get("content", ""), 

71 ttl=item_data.get("ttl", "session"), 

72 metadata=item_data.get("metadata", {}), 

73 )