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 

14 key: str 

15 content: str 

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

17 metadata: dict = field(default_factory=dict) 

18 

19 

20class WorkingMemory: 

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

22 

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

24 self.max_items = max_items 

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

26 

27 def add(self, item: MemoryItem): 

28 self._items[item.key] = item 

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

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

31 del self._items[oldest] 

32 

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

34 return self._items.get(key) 

35 

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

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

38 results = [] 

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

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

41 results.append(item) 

42 return results[:limit] 

43 

44 def clear(self): 

45 self._items.clear() 

46 

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

48 

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

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

51 return { 

52 "max_items": self.max_items, 

53 "items": { 

54 key: { 

55 "key": item.key, 

56 "content": item.content, 

57 "ttl": item.ttl, 

58 "metadata": item.metadata, 

59 } 

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

61 }, 

62 } 

63 

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

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

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

67 self._items.clear() 

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

69 self._items[key] = MemoryItem( 

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

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

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

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

74 )