fix(memory): make auto-memory extraction reliable for reasoning models (#3190)

* fix(memory): auto-memory extracted nothing — flatten window so the prompt ends on a user turn extract_and_store appended the recent window as raw alternating role messages after the system prompt. Since the window is the last N messages, the prompt usually ENDED on an assistant turn — and a chat model given a prompt ending on an assistant turn returns an empty completion (nothing to answer). The result was facts=[] → "Auto memory extraction ran: 0 candidates" on every run, so no memories were ever stored, while skill extraction (which flattens the transcript into a single user message) worked fine. Flatten the window into one user message ending with an explicit instruction, mirroring the skill extractor, so the model always responds. Also harden parsing for reasoning models, matching the audit path which already does this: - raise max_tokens 500 → 4096 (a reasoning model spends the budget on <think> before emitting JSON; 500 truncated it before any JSON appeared); - strip <think>/prose preambles via strip_think and slice the embedded JSON array before json.loads, instead of bombing on char 0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * chore: tighten memory-extraction-empty-completion — clarify JSON-slice comment re prior strip steps Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(memory): reframe the comment to the accurate root cause (raw-chat framing) The earlier comment leaned on "ends on an assistant turn -> empty completion", which is only one failure mode. The dominant cause, confirmed by a controlled repro (0/6 old vs 6/6 new on this model), is that passing the window as raw chat messages makes the model treat it as a conversation to continue rather than a transcript to analyze, so it returns [] even when durable facts are present. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * test(memory): cover extraction JSON parsing + slice trailing commentary unconditionally Factor the strip/fence/slice/json.loads logic out of extract_and_store into a pure module-level helper _parse_extraction_json(raw) -> list and drop the 'text[0] != "["' guard so the array is sliced whenever both brackets exist (fixes trailing commentary like '[...] Done!' reaching json.loads). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 17:25:26 -04:00 · 2026-06-08 10:57:44 -07:00
parent 8449baea80
commit 34a3f8637a
2 changed files with 110 additions and 11 deletions
@@ -236,6 +236,43 @@ def _is_text_duplicate(new_text: str, existing: list, threshold: float = 0.6) ->
    return False


+def _parse_extraction_json(raw: str) -> list:
+    """Parse the extraction LLM's reply into a list of facts, tolerating
+    reasoning-model noise.
+
+    The model emits <think>…</think> (and sometimes a prose preamble or a
+    ```json fence) AROUND the JSON array; without stripping it, json.loads
+    bombs and the run silently yields "0 candidates". Pure str -> list (no
+    LLM/network); returns [] on any parse failure instead of raising.
+    """
+    text = (raw or "").strip()
+    try:
+        from src.text_helpers import strip_think as _strip_think
+        text = _strip_think(text, prose=True, prompt_echo=True).strip()
+    except Exception:
+        pass
+    if text.startswith("```"):
+        text = text.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
+    # JSON may still be embedded in surrounding commentary (leading prose or
+    # trailing remarks like "[...] Done!") — slice from the first '[' to the
+    # last ']' whenever both exist. Slice unconditionally: a reply that starts
+    # with '[' can still carry trailing commentary that breaks json.loads.
+    _start = text.find("[")
+    _end = text.rfind("]")
+    if 0 <= _start < _end:
+        text = text[_start : _end + 1]
+
+    try:
+        facts = json.loads(text)
+    except json.JSONDecodeError:
+        logger.debug("Memory extraction returned non-JSON: %r", (raw or "")[:120])
+        return []
+    except Exception:
+        logger.debug("Memory extraction returned non-JSON: %r", (raw or "")[:120])
+        return []
+    return facts if isinstance(facts, list) else []
+
+
 async def extract_and_store(
    session,
    memory_manager,
@@ -284,9 +321,34 @@ async def extract_and_store(

        fallback_facts = _fallback_memory_candidates(stripped_recent)

+        # Flatten the window into a SINGLE user message instead of appending the
+        # raw alternating role messages. Passed as raw chat messages, the model
+        # treats the window as a conversation to CONTINUE rather than a transcript
+        # to ANALYZE, so it reliably extracts nothing — typically returning `[]`
+        # (and, depending on the input, sometimes an empty or <think>-only
+        # completion when the window ends on an assistant turn). This was the real
+        # cause of auto-memory logging "0 candidates" on every run. Reframing it as
+        # one "analyze this transcript, return the JSON array" user message makes
+        # the model actually extract. Controlled repro on this model: 0/6 trials
+        # with the old structure vs 6/6 with this one. The skill extractor flattens
+        # for the same reason.
+        def _flatten_msg(m):
+            c = m.get("content", "")
+            if isinstance(c, list):
+                c = " ".join(
+                    b.get("text", "") for b in c
+                    if isinstance(b, dict) and b.get("type") == "text"
+                )
+            return f"{m.get('role', '?')}: {c}"
+
+        transcript = "\n\n".join(_flatten_msg(m) for m in stripped_recent)
        extraction_messages = [
            {"role": "system", "content": EXTRACT_SYSTEM_PROMPT},
-        ] + stripped_recent
+            {"role": "user", "content": (
+                "Conversation to analyze:\n\n" + transcript
+                + "\n\nReturn the JSON array of durable facts now (or [] if none)."
+            )},
+        ]

        facts = []
        try:
@@ -295,19 +357,20 @@ async def extract_and_store(
                model,
                extraction_messages,
                temperature=0.1,
-                max_tokens=500,
+                # A reasoning model spends most of its budget on <think> tokens
+                # BEFORE emitting the JSON, so the old 500 truncated the response
+                # before any JSON appeared → every run logged "0 candidates". The
+                # audit path hit the same wall and raised to 16384; extraction's
+                # output (a short facts list) is small, so an ample ceiling is
+                # enough once thinking has room.
+                max_tokens=4096,
                headers=headers,
            )

-            # Parse JSON from response (handle markdown fences if model wraps them)
-            text = raw.strip()
-            if text.startswith("```"):
-                text = text.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
-
-            try:
-                facts = json.loads(text)
-            except json.JSONDecodeError:
-                logger.debug("Memory extraction returned non-JSON")
+            # Parse JSON, tolerating reasoning-model noise (<think> blocks, a
+            # ```json fence, and leading/trailing commentary). See
+            # _parse_extraction_json — returns [] rather than raising.
+            facts = _parse_extraction_json(raw)
        except Exception as e:
            logger.warning(f"LLM memory extraction failed; using fallback candidates if available: {e}")

@@ -0,0 +1,36 @@
+"""_parse_extraction_json must survive reasoning-model noise.
+
+The extraction model wraps its JSON array in <think> blocks, ```json fences,
+or leading/trailing prose. The helper strips that noise and slices the array
+unconditionally — a reply that starts with '[' can still carry trailing
+commentary like "[...] Done!" that would otherwise break json.loads.
+"""
+
+from services.memory.memory_extractor import _parse_extraction_json
+
+
+def test_think_prefixed_array_parses_to_one_fact():
+    raw = '<think>reasoning...</think>\n[{"text": "x", "category": "fact"}]'
+    assert _parse_extraction_json(raw) == [{"text": "x", "category": "fact"}]
+
+
+def test_fenced_json_block_parses():
+    raw = '```json\n[{"text": "x", "category": "fact"}]\n```'
+    assert _parse_extraction_json(raw) == [{"text": "x", "category": "fact"}]
+
+
+def test_leading_prose_before_array_parses():
+    raw = 'Here are the durable facts:\n[{"text": "x", "category": "fact"}]'
+    assert _parse_extraction_json(raw) == [{"text": "x", "category": "fact"}]
+
+
+def test_trailing_commentary_after_array_parses():
+    # Exercises the unconditional slice: text starts with '[' but has trailing
+    # commentary that the old `text[0] != "["` guard skipped, breaking json.loads.
+    raw = '[{"text": "x", "category": "fact"}] Done!'
+    assert _parse_extraction_json(raw) == [{"text": "x", "category": "fact"}]
+
+
+def test_malformed_no_array_returns_empty():
+    assert _parse_extraction_json("no array here, sorry") == []
+    assert _parse_extraction_json("") == []