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

1""" 

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

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass, field 

8from typing import Any 

9 

10 

11@dataclass 

12class MemoryItem: 

13 """工作记忆项。""" 

14 

15 key: str 

16 content: str 

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

18 metadata: dict = field(default_factory=dict) 

19 

20 

21class WorkingMemory: 

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

23 

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

25 self.max_items = max_items 

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

27 

28 def add(self, item: MemoryItem): 

29 self._items[item.key] = item 

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

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

32 del self._items[oldest] 

33 

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

35 return self._items.get(key) 

36 

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

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

39 results = [] 

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

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

42 results.append(item) 

43 return results[:limit] 

44 

45 def clear(self): 

46 self._items.clear() 

47 

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

49 

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

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

52 return { 

53 "max_items": self.max_items, 

54 "items": { 

55 key: { 

56 "key": item.key, 

57 "content": item.content, 

58 "ttl": item.ttl, 

59 "metadata": item.metadata, 

60 } 

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

62 }, 

63 } 

64 

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

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

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

68 self._items.clear() 

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

70 self._items[key] = MemoryItem( 

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

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

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

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

75 )