NexusOS/synapse/memory/extractor.py

"""Memory extraction — asks Mistral to evaluate a conversation exchange and
decide if it contains a new permanent personal fact worth saving."""

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 exchange.
There may be SEVERAL facts in one message — output one JSON object for each.

ONLY the USER line is a source of facts. The ASSISTANT line and the existing
memory below are context to help you understand the USER line — never extract
anything from them. If the assistant said it and the user did not, it is NOT a
fact. 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.

---
USER: {user_message}
ASSISTANT: {assistant_response}
---

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


async def extract_memory(
    user_message: str,
    assistant_response: str,
    existing_sections: list[str],
    existing_texts: list[str],
    ollama_manager,
    model: str = DEFAULT_MEMORY_MODEL,
    num_gpu: int | None = 0,
) -> list[dict]:
    """Ask Mistral to extract saveable memory facts from a conversation exchange.

    Returns a list of {"section": ..., "text": ...} — possibly empty. A single
    exchange can hold several facts, and Mistral emits one JSON object per fact.
    """
    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,
        user_message=user_message[:3000],
        assistant_response=assistant_response[:800],
    )
    # 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_message)
                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. mistral 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 []