NexusOS/synapse/memory/extractor.py

"""Memory extraction — asks the curator model to read a finished conversation
and decide which new permanent personal facts it contains.

Reads the whole transcript, not a single exchange. A fact is rarely complete in
the turn that introduces it ("I have a Honda" at turn 3 becomes a 2006 Accord
with 312k miles by turn 7), and the store is append-only for the curator, so
extracting per-exchange could only ever accumulate fragments of the same fact.
Waiting until the conversation is idle is also the only reliable signal that a
statement is finished rather than half-said."""

from __future__ import annotations

import json
import logging
import re

from ..chat import _synapse_trace  # append curator reasoning to the same MindTrace log
from ..nexus_config import DEFAULT_MEMORY_MODEL

_log = logging.getLogger(__name__)

_TR = "┅" * 55

_PROMPT = """\
You are a memory curator for a personal AI assistant named Nexus.

Extract EVERY new, permanent personal fact the USER stated in this conversation.
There may be SEVERAL facts across the conversation — output one JSON object for
each.

ONLY the USER lines are a source of facts. The ASSISTANT lines and the existing
memory below are context to help you understand the USER lines — never extract
anything from them. If the assistant said it and the user did not, it is NOT a
fact.

You are reading the WHOLE conversation, so record the FINAL, most complete form
of each fact. If the user gives a detail early and refines it later, output one
object with the refined version — never one per stage. If the user corrects or
retracts something, keep only what they ended up saying. Copy what the user actually said; do not infer, embellish, or add a
judgement the user did not make (never call something their "favorite",
"main", or "best" unless the user used that word).

Never save a statement about what is unknown, unspecified, or absent — no
"X is unknown", "has not said", "has no Y". Silence is not a fact.

SAVE facts that are stable and biographical, such as: identity (name, age,
location), relationships (family, partner, friends), pets, possessions (vehicles,
home, devices), career (job, employer, skills), hobbies and interests,
long-running projects or goals (not today's to-dos), and durable preferences.

DO NOT SAVE — these are ephemeral and would clutter memory:
- Anything time-bound: "today I have...", "I'm working on X today"
- Mood or energy: "I'm tired", "feeling good", "having a rough day"
- Greetings or small talk: "good morning", "how are you"
- Questions the user asked the assistant

The existing memory is below FOR CONTEXT ONLY — it is already saved, so never
output anything from it, and never comment on what it does or does not contain.
If the user ADDS NEW DETAIL to something already known (e.g. a new detail about
a known pet, car, or project), DO save that new detail as its own fact. Skip
anything already covered by a listed fact, even when the wording differs.

Existing memory:
{existing_texts}

For "section", REUSE one of these existing section names whenever it fits:
{existing_sections}
Only invent a new section if none fit, and make it a SHORT single word
(e.g. Hobbies, Pets, Health). Never use a sentence or long phrase as a section.

---
{transcript}
---

Respond with JSON only — no prose, no markdown fences. Output one object PER
new fact (several objects, one per line, if there are several):
{{"save": true, "section": "<short section name>", "text": "<concise fact about the user, third person>"}}
If there is nothing new to save, output exactly:
{{"save": false}}"""


_ABSENCE = re.compile(
    r"\b(?:no longer|does ?n[o']t|do ?n[o']t|did ?n[o']t|has not|hasn't|have not|"
    r"haven't|is not|isn't|are not|aren't|has no|have no|not specified|"
    r"unspecified|unknown|not mentioned|no pets|as per the|according to the "
    r"(?:current )?memory)\b",
    re.IGNORECASE,
)

# The subject's name and the assistant's are in every fact by instruction, so
# they prove nothing about grounding. "The"/"User" are here because the regex
# below treats any capitalised word as distinctive.
_STOP_TOKENS = {"jon", "nexus", "the", "user"}


def _distinctive(text: str) -> set[str]:
    """Proper nouns and numbers in a fact - the parts that can't be invented
    from thin air without showing up in what the user actually typed."""
    tokens = set(re.findall(r"\b[A-Z][A-Za-z0-9.+-]{2,}\b|\b\d[\d.]*\w*\b", text))
    return {t for t in tokens if t.lower() not in _STOP_TOKENS}


def _reject_reason(fact: str, user_message: str) -> str | None:
    """Why this fact must not be saved, or None to keep it.

    Two failure classes the curator model keeps producing no matter how the
    prompt is worded (verified against mistral:7b):

    1. Absence claims. It reads the existing-memory block and writes things like
       "the user does not have any pets" - which contradicted four cats already on
       file. Silence is not a fact.
    2. Assistant-sourced specifics. It lifts names the ASSISTANT said and
       attributes them to the user: a reply that echoed a stale memory row
       produced "the user's main development machine is a MacBook Pro" off the user
       message "What am I developing on?".

    The grounding test only fires when a fact carries distinctive tokens and
    NONE of them appear in the user's own message. A fact with no proper nouns
    or numbers ("prefers casual conversation") is left to the prompt.
    """
    if _ABSENCE.search(fact):
        return "absence claim"
    marks = _distinctive(fact)
    if marks and not any(m.lower() in user_message.lower() for m in marks):
        return f"ungrounded in the user message (invented {sorted(marks)[:3]})"
    return None


def render_transcript(messages: list[dict]) -> str:
    """The conversation as the curator sees it. Assistant turns are truncated
    hard: they are context for reading the user's lines, never a fact source,
    and a long reply would otherwise crowd out the lines that matter."""
    lines = []
    for m in messages:
        role = (m.get("role") or "").upper()
        if role not in ("USER", "ASSISTANT"):
            continue
        text = (m.get("content") or "").strip()
        if not text:
            continue
        lines.append(f"{role}: {text[:3000] if role == 'USER' else text[:600]}")
    return "\n".join(lines)


async def extract_memory(
    messages: list[dict],
    existing_sections: list[str],
    existing_texts: list[str],
    ollama_manager,
    model: str = DEFAULT_MEMORY_MODEL,
    num_gpu: int | None = 0,
) -> list[dict]:
    """Extract saveable memory facts from a whole conversation.

    `messages` is the transcript as [{role, content}]. Returns a list of
    {"section": ..., "text": ...} — possibly empty. One conversation can hold
    several facts, and the model emits one JSON object per fact.
    """
    transcript = render_transcript(messages)
    if not transcript:
        return []
    # Grounding is checked against everything the user actually typed, so a fact
    # assembled from details spread across several of their turns still passes.
    user_text = "\n".join(
        (m.get("content") or "") for m in messages if (m.get("role") or "") == "user"
    )
    if existing_texts:
        # ponytail: only the 12 most-recent facts go in the dedup context, not all
        # ~40. On a CPU-bound curator (num_gpu=0) prompt-eval dominates, and 40
        # facts made a ~1200-token prompt that took ~50s+ to process. If dedup
        # starts re-saving older facts, move dedup to a difflib check in the
        # service instead of stuffing every fact into the prompt.
        texts_block = "\n".join(f"- {t}" for t in existing_texts[-12:])
    else:
        texts_block = "(none yet)"
    # Give Mistral the real section names to reuse, so it stops inventing
    # sentence-long sections out of the category descriptions in the prompt.
    # Only offer SHORT, clean names — never feed a junk sentence-section (e.g. a
    # past bad "Long-running projects or goals") back as a valid choice.
    clean = sorted(s for s in existing_sections if s and len(s.split()) <= 2 and len(s) <= 24)
    sections_line = ", ".join(clean) if clean else (
        "Identity, Relationships, Pets, Possessions, Career, Hobbies, Projects, Preferences"
    )
    prompt = _PROMPT.format(
        existing_texts=texts_block,
        existing_sections=sections_line,
        transcript=transcript,
    )
    # MindTrace: curator pre-flight (full prompt) so its reasoning is visible in
    # the same console as the frontline model, not just Python warnings on failure.
    _synapse_trace(
        f"\n{_TR}\n◆ CURATOR: {model} (num_gpu={num_gpu})\n"
        f"  PROMPT ({len(prompt)} chars):\n{prompt}\n{_TR}\n"
    )
    try:
        # num_gpu=0 (the default) pins the curator fully in system RAM instead of
        # the GPU, so it coexists with the GPU-resident chat model instead of
        # evicting it. Without this, on a small GPU the two thrash: every
        # exchange cold-loads the curator (~45s) and extraction times out,
        # silently saving nothing. Boxes with spare VRAM override via settings.
        response = await ollama_manager.chat(
            messages=[{"role": "user", "content": prompt}],
            model=model,
            stream=False,
            temperature=0.0,
            num_gpu=num_gpu,
        )
        if not response:
            _synapse_trace("◆ CURATOR RAW: (empty response)\n")
            return []

        text = response.strip()
        _synapse_trace(f"◆ CURATOR RAW:\n{text}\n")

        # Strip markdown code fences if the model added them
        if "```" in text:
            m = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
            if m:
                text = m.group(1).strip()

        # For several facts Mistral is inconsistent: sometimes ONE JSON object
        # per fact newline-separated, sometimes a single JSON ARRAY of objects.
        # raw_decode pulls each top-level value (handles the newline case and
        # plain "Extra data"); we then flatten any array so both shapes save all
        # facts. Plain json.loads() would die on the newline case and skip the
        # array (a list isn't a dict), losing every fact either way.
        results: list[dict] = []

        def _keep(o):
            if isinstance(o, dict) and o.get("save") and o.get("section") and o.get("text"):
                fact = str(o["text"]).strip()
                reason = _reject_reason(fact, user_text)
                if reason:
                    _synapse_trace(f"◆ CURATOR DROPPED ({reason}): {fact}\n")
                    return
                results.append({"section": str(o["section"]).strip(), "text": fact})

        dec = json.JSONDecoder()
        idx = 0
        while idx < len(text):
            while idx < len(text) and text[idx] in " \t\r\n,":
                idx += 1
            if idx >= len(text):
                break
            try:
                obj, idx = dec.raw_decode(text, idx)
            except json.JSONDecodeError:
                break
            if isinstance(obj, list):
                for o in obj:
                    _keep(o)
            else:
                _keep(obj)
        if not results:
            _log.warning("memory: nothing saved. curator said: %.300r", text)
            _synapse_trace(f"◆ CURATOR VERDICT: nothing to save\n{_TR}\n\n")
        else:
            _facts = "; ".join(f"[{r['section']}] {r['text']}" for r in results)
            _synapse_trace(f"◆ CURATOR VERDICT: {len(results)} fact(s) — {_facts}\n{_TR}\n\n")
        return results
    except Exception as e:
        _log.warning("memory extraction failed: %s", e)
        _synapse_trace(f"◆ CURATOR ERROR: {e}\n{_TR}\n\n")
    return []