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

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

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. 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. Only skip a fact that is an EXACT
restatement of one already listed.

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 Jon, third person>"}}
If there is nothing new to save, output exactly:
{{"save": false}}"""


async def extract_memory(
    user_message: str,
    assistant_response: str,
    existing_sections: list[str],
    existing_texts: list[str],
    ollama_manager,
    model: str = "mistral:latest",
    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"):
                results.append({
                    "section": str(o["section"]).strip(),
                    "text": str(o["text"]).strip(),
                })

        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 []