Coverage for src/lilbee/retrieval/query/memory

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

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"""

10from __future__ import annotations

12import logging

13from collections.abc import Callable

14from dataclasses import dataclass

16from lilbee.core.llm_json import first_json_array

17from lilbee.data.store import MemoryKind

19log = logging.getLogger(__name__)

21ChatFn = Callable[..., str]

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.

25MEMORY_EXTRACT_MAX_TOKENS = 800

26_MIN_MEMORY_CHARS = 10

27_MAX_MEMORY_CHARS = 500

29_EXTRACT_SYSTEM_PROMPT = (

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

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

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

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

34 "just this question. "

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

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

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

38)

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

43@dataclass(frozen=True)

44class ExtractedMemory:

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

47 text: str

48 kind: MemoryKind

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

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

53 return [

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

55 {

56 "role": "user",

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

58 },

59 ]

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

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

64 if not isinstance(value, str):

65 return MemoryKind.FACT

66 try:

67 return MemoryKind(value)

68 except ValueError:

69 return MemoryKind.FACT

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

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

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

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

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

78 """

79 items = first_json_array(raw)

80 if items is None:

81 return []

83 memories: list[ExtractedMemory] = []

84 for item in items:

85 if not isinstance(item, dict):

86 continue

87 text = item.get("text")

88 if not isinstance(text, str):

89 continue

90 text = text.strip()

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

92 continue

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

94 return memories

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

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

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

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

102 chat session.

103 """

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

105 return []

106 try:

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

108 except Exception:

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

110 return []

111 return parse_extraction(raw)

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

51 statements