Coverage for src / lilbee / app / memory.py: 100%

1"""Use-case orchestration for long-term chat memory, shared by every surface. 

2 

3Surfaces (TUI, CLI, MCP, REST, Python API) call these functions rather than 

4constructing ``MemoryRow`` objects or building owner predicates themselves, so 

5embedding, id/timestamp assignment, and scoping live in one place. 

6""" 

7 

8from __future__ import annotations 

9 

10import uuid 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13from datetime import UTC, datetime 

14 

15from lilbee.app.services import get_services 

16from lilbee.core.config import cfg 

17from lilbee.data.store import ( 

18 LOCAL_OWNER, 

19 MemoryKind, 

20 MemoryRow, 

21 MemorySource, 

22 agent_recall_predicate, 

23 escape_sql_string, 

24 local_owner_predicate, 

25) 

26 

27 

28def make_memory_row( 

29 text: str, 

30 embed: Callable[[str], list[float]], 

31 *, 

32 owner: str = LOCAL_OWNER, 

33 kind: MemoryKind = MemoryKind.FACT, 

34 source: MemorySource = MemorySource.MANUAL, 

35 shared: bool = False, 

36) -> MemoryRow: 

37 """Build a fully populated ``MemoryRow`` with a fresh id, timestamps, and 

38 an embedded vector. The single id/timestamp/embedding assignment point, so 

39 callers supply only their own ``embed`` and store the result. 

40 """ 

41 now = datetime.now(UTC).isoformat() 

42 return MemoryRow( 

43 id=uuid.uuid4().hex, 

44 owner=owner, 

45 shared=shared, 

46 kind=kind, 

47 source=source, 

48 text=text, 

49 vector=embed(text), 

50 created_at=now, 

51 updated_at=now, 

52 ) 

53 

54 

55MEMORY_DISABLED_HINT = ( 

56 "Memory is off. Turn it on with /set memory_enabled true in the TUI, " 

57 "settings_set via MCP, or memory_enabled = true in config.toml." 

58) 

59 

60 

61def memory_enabled() -> bool: 

62 """True when the memory subsystem is switched on (off by default).""" 

63 return cfg.memory_enabled 

64 

65 

66def remember( 

67 text: str, 

68 *, 

69 owner: str = LOCAL_OWNER, 

70 kind: MemoryKind = MemoryKind.FACT, 

71 source: MemorySource = MemorySource.MANUAL, 

72 shared: bool = False, 

73) -> str: 

74 """Embed *text* and store it as a memory; returns the stored id.""" 

75 services = get_services() 

76 record = make_memory_row( 

77 text, 

78 services.embedder.embed, 

79 owner=owner, 

80 kind=kind, 

81 source=source, 

82 shared=shared, 

83 ) 

84 return services.store.add_memory(record) 

85 

86 

87def recall(query: str, owner: str = LOCAL_OWNER, *, top_k: int | None = None) -> list[MemoryRow]: 

88 """Recall *owner*'s facts (plus human-shared facts for agents).""" 

89 services = get_services() 

90 predicate = local_owner_predicate() if owner == LOCAL_OWNER else agent_recall_predicate(owner) 

91 return services.store.search_memories( 

92 services.embedder.embed(query), 

93 owner_predicate=predicate, 

94 top_k=cfg.memory_top_k if top_k is None else top_k, 

95 max_distance=cfg.memory_max_distance, 

96 ) 

97 

98 

99def list_memories(owner: str = LOCAL_OWNER) -> list[MemoryRow]: 

100 """List all of *owner*'s memories (any kind), newest first.""" 

101 predicate = ( 

102 local_owner_predicate() if owner == LOCAL_OWNER else f"owner = '{escape_sql_string(owner)}'" 

103 ) 

104 return get_services().store.get_memories(owner_predicate=predicate) 

105 

106 

107def forget(memory_id: str) -> None: 

108 """Delete a memory by id.""" 

109 get_services().store.delete_memory(memory_id) 

110 

111 

112def set_memory_shared(memory_id: str, *, shared: bool) -> bool: 

113 """Set a memory's shared-with-agents flag; returns True when the id exists.""" 

114 return get_services().store.update_memory(memory_id, shared=shared) 

115 

116 

117def auto_extract_enabled() -> bool: 

118 """True when auto-extraction is on (requires the master gate too).""" 

119 return cfg.memory_enabled and cfg.memory_auto_extract 

120 

121 

122@dataclass(frozen=True, slots=True) 

123class SavedMemory: 

124 """A memory created by auto-extraction: its stored id, kind, and text.""" 

125 

126 id: str 

127 kind: MemoryKind 

128 text: str 

129 

130 

131def auto_extract(question: str, answer: str) -> list[SavedMemory]: 

132 """Extract durable memories from a chat turn and store them. 

133 

134 Returns one :class:`SavedMemory` per stored memory. Stored memories are 

135 ``source=EXTRACTED`` and are recalled like any other; the user manages them 

136 in ``/memories``. A no-op (returns ``[]``) unless both the master gate and 

137 ``memory_auto_extract`` are on. 

138 """ 

139 from lilbee.retrieval.query.memory_extract import extract_memories 

140 

141 if not auto_extract_enabled(): 

142 return [] 

143 services = get_services() 

144 extracted = extract_memories(question, answer, services.provider.chat) 

145 stored: list[SavedMemory] = [] 

146 for memory in extracted: 

147 memory_id = remember(memory.text, kind=memory.kind, source=MemorySource.EXTRACTED) 

148 stored.append(SavedMemory(id=memory_id, kind=memory.kind, text=memory.text)) 

149 return stored