Coverage for src/lilbee/retrieval/query/memory_extract.py: 100%

1"""Auto-extraction of durable memories from a chat turn. 

2 

3A small LLM pass over the user's message and the assistant's answer that 

4proposes durable facts/preferences worth remembering. The model is asked for 

5a strict JSON array; parsing is defensive (a non-conforming reply yields no 

6memories rather than an error). Callers store the results, which the user can 

7review and remove in ``/memories``. 

8""" 

9 

10from __future__ import annotations 

11 

12import logging 

13from collections.abc import Callable 

14from dataclasses import dataclass 

15 

16from lilbee.core.llm_json import first_json_array 

17from lilbee.data.store import MemoryKind 

18 

19log = logging.getLogger(__name__) 

20 

21ChatFn = Callable[..., str] 

22 

23# Bounds mirror the prior-art auto-capture filter: too-short strings carry no 

24# durable signal, too-long ones are usually the model restating the answer. 

25_MIN_MEMORY_CHARS = 10 

26_MAX_MEMORY_CHARS = 500 

27 

28_EXTRACT_SYSTEM_PROMPT = ( 

29 "You extract durable, long-term memories about the user from a single chat turn. " 

30 "A memory is a stable fact about the user or their project (not a one-off question) " 

31 "or a standing preference for how they want help. " 

32 "Ignore transient details, the assistant's own content, and anything specific to " 

33 "just this question. " 

34 "Respond with ONLY a JSON array (no prose). Each element is an object " 

35 '{"text": "<the memory, third person>", "kind": "fact" | "preference"}. ' 

36 "Return [] when nothing is worth remembering." 

37) 

38 

39_EXTRACT_USER_TEMPLATE = "User said:\n{question}\n\nAssistant replied:\n{answer}" 

40 

41 

42@dataclass(frozen=True) 

43class ExtractedMemory: 

44 """A single memory proposed by the extraction pass.""" 

45 

46 text: str 

47 kind: MemoryKind 

48 

49 

50def build_extract_messages(question: str, answer: str) -> list[dict[str, str]]: 

51 """Build the system+user message pair for the extraction prompt.""" 

52 return [ 

53 {"role": "system", "content": _EXTRACT_SYSTEM_PROMPT}, 

54 { 

55 "role": "user", 

56 "content": _EXTRACT_USER_TEMPLATE.format(question=question, answer=answer), 

57 }, 

58 ] 

59 

60 

61def _coerce_kind(value: object) -> MemoryKind: 

62 """Decode a kind string, defaulting to FACT for anything unrecognized.""" 

63 if not isinstance(value, str): 

64 return MemoryKind.FACT 

65 try: 

66 return MemoryKind(value) 

67 except ValueError: 

68 return MemoryKind.FACT 

69 

70 

71def parse_extraction(raw: str) -> list[ExtractedMemory]: 

72 """Parse the model's reply into memories; tolerate a non-conforming reply. 

73 

74 Extracts the first JSON array in *raw* (models often wrap it in prose or a 

75 code fence), keeps only objects with a usable-length ``text``, and decodes 

76 the kind. Any parse failure yields an empty list. 

77 """ 

78 items = first_json_array(raw) 

79 if items is None: 

80 return [] 

81 

82 memories: list[ExtractedMemory] = [] 

83 for item in items: 

84 if not isinstance(item, dict): 

85 continue 

86 text = item.get("text") 

87 if not isinstance(text, str): 

88 continue 

89 text = text.strip() 

90 if not _MIN_MEMORY_CHARS <= len(text) <= _MAX_MEMORY_CHARS: 

91 continue 

92 memories.append(ExtractedMemory(text=text, kind=_coerce_kind(item.get("kind")))) 

93 return memories 

94 

95 

96def extract_memories(question: str, answer: str, chat: ChatFn) -> list[ExtractedMemory]: 

97 """Run the extraction pass for one turn; never raises. 

98 

99 *chat* is the provider's non-streaming chat callable. A model or transport 

100 failure logs and yields no memories so a bad extraction never disrupts the 

101 chat session. 

102 """ 

103 if not question.strip() or not answer.strip(): 

104 return [] 

105 try: 

106 raw = chat(build_extract_messages(question, answer), stream=False) 

107 except Exception: 

108 log.debug("Memory extraction call failed", exc_info=True) 

109 return [] 

110 return parse_extraction(raw)