Tolan/lumi_memory.py

# lumi_memory.py
# Session-level long-term memory extraction + persistence for Lumi

import json
import os
from typing import Dict, Any, List, Optional

DEFAULT_MEMORY_DIR = "memories"

KEY_MEMORY_EXTRACTION_PROMPT = r"""
You are a memory extraction engine for Lumi.
Your job: read the full session conversation and decide what (if anything) should be saved as long-term memory for personalization in future chats.

Save only info that is:
1) Stable for weeks/months/years, AND
2) Useful for future personalization (tone, interests, preferences, identity, goals, recurring constraints), AND
3) Not overly private or risky.

Good examples:
- Name / preferred nickname (if explicitly stated)
- Preferences (tone, topic likes/dislikes, boundaries like “don’t ask X”)
- Long-term goals (e.g., “practice speaking confidence”)
- Persistent interests/hobbies (if clearly stable)

Do NOT save:
- One-off details (today’s breakfast, today’s mood)
- Sensitive personal data (exact address, ID numbers, medical diagnoses, political opinions, etc.)
- Private secrets the user didn’t ask to store

If user explicitly says “remember this / save this”, you MAY store it unless it is sensitive.
If sensitive, store a safer abstract version.

OUTPUT FORMAT (JSON ONLY):
{
  "should_write_memory": true/false,
  "memories_to_add": [
    {"key":"...", "value":"...", "confidence":0.0-1.0, "evidence":"..."}
  ],
  "memories_to_update": [
    {"key":"...", "old_value":"...", "new_value":"...", "confidence":0.0-1.0, "evidence":"..."}
  ],
  "memories_to_forget": [
    {"key":"...", "reason":"..."}
  ],
  "do_not_store": [
    {"candidate":"...", "reason":"one-off | sensitive | not useful"}
  ],
  "notes_for_personalization": [
    "..."
  ]
}

Rules:
- Return JSON only. No extra text.
- If nothing useful, should_write_memory=false and arrays empty.
- Prefer fewer, higher-quality memories.
- Keys should be namespaced, e.g. identity.name, preference.tone, preference.avoid, goal.long_term.
""".strip()


INCREMENTAL_MEMORY_SUMMARY_PROMPT = r"""
You are Lumi's memory-summary updater.

You will receive:
1) The previous memory summary (may be empty)
2) The latest long-term memory JSON (key-value pairs)

Task:
Update the previous summary into a single, smooth, assistant-facing paragraph
that captures the user's stable preferences, interaction style, and recurring tendencies.

Editing principles (important):
- Prefer minimal edits to the existing summary when still valid.
- Integrate new information naturally into the paragraph only if it provides clear long-term personalization value.
- Remove or soften details that are no longer supported by memory.
- Generalize or omit anything that could become sensitive.
- Avoid rigid traits or absolute statements that would overly constrain future conversations.

Output constraints:
- ONE short paragraph only (no bullet points).
- Around 60–100 words.
- Neutral, warm, assistant-facing tone.
- Describe tendencies and preferences, not fixed identities.
- Avoid explicit instructions, rules, or triggers.
- Do NOT mention internal memory keys or system concepts.
- Plain text only. No headings. No JSON.

The goal is to provide gentle background context that supports natural topic flow,
not to force alignment or explicit recall.
""".strip()


def _ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def get_memory_path(user_id: str, memory_dir: str = DEFAULT_MEMORY_DIR) -> str:
    _ensure_dir(memory_dir)
    safe = "".join(c for c in (user_id or "default") if c.isalnum() or c in ("_", "-", ".")).strip()
    if not safe:
        safe = "default"
    return os.path.join(memory_dir, f"memory_{safe}.json")


def load_memory(user_id: str, memory_dir: str = DEFAULT_MEMORY_DIR) -> Dict[str, Any]:
    path = get_memory_path(user_id, memory_dir)
    if os.path.exists(path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            return data if isinstance(data, dict) else {}
        except Exception:
            return {}
    return {}


def save_memory(user_id: str, memory: Dict[str, Any], memory_dir: str = DEFAULT_MEMORY_DIR) -> None:
    path = get_memory_path(user_id, memory_dir)
    _ensure_dir(memory_dir)

    mem = dict(memory or {})

    # ✅ Ensure meta.memory_summary is always the first key in the saved JSON
    if "meta.memory_summary" in mem:
        first = {"meta.memory_summary": mem.get("meta.memory_summary")}
        rest = {k: v for k, v in mem.items() if k != "meta.memory_summary"}
        mem = {**first, **rest}

    with open(path, "w", encoding="utf-8") as f:
        json.dump(mem, f, ensure_ascii=False, indent=2)


def format_memory_for_runtime(memory: Dict[str, Any], max_items: int = 8) -> str:
    """
    Simple formatter for injecting memory into runtime prompts.
    """
    if not memory:
        return "[LONG-TERM MEMORY]\n(none)"

    items = []
    for k, v in memory.items():
        if v is None:
            continue
        s = str(v).strip()
        if not s:
            continue
        items.append(f"- {k}: {s}")

    return "[LONG-TERM MEMORY]\n" + "\n".join(items[:max_items])


def build_session_snippet(messages: List[dict], max_turns: int = 80) -> str:
    """
    Build a session transcript from messages.
    max_turns limits how many user/assistant messages we include (to avoid huge prompt).
    """
    turns = [m for m in messages if m.get("role") in ("user", "assistant")]
    turns = turns[-max_turns:]  # keep tail if too long

    lines = []
    for m in turns:
        role = "User" if m["role"] == "user" else "Lumi"
        lines.append(f"{role}: {m.get('content', '')}")
    return "\n".join(lines).strip()


def safe_json_load(s: str) -> Dict[str, Any]:
    """
    Robust JSON loader for imperfect model outputs.
    """
    s = (s or "").strip()
    if not s:
        return {}
    try:
        return json.loads(s)
    except Exception:
        start = s.find("{")
        end = s.rfind("}")
        if start != -1 and end != -1 and end > start:
            try:
                return json.loads(s[start: end + 1])
            except Exception:
                return {}
        return {}


def _to_clean_text(v: Any) -> str:
    if v is None:
        return ""
    if isinstance(v, str):
        return v.strip()
    if isinstance(v, list):
        parts = []
        for x in v:
            if x is None:
                continue
            s = x.strip() if isinstance(x, str) else str(x).strip()
            if s:
                parts.append(s)
        return "; ".join(parts).strip()
    if isinstance(v, dict):
        return json.dumps(v, ensure_ascii=False).strip()
    return str(v).strip()


def apply_memory_patch(
        memory: Dict[str, Any],
        patch: Dict[str, Any],
        min_confidence: float = 0.55,
) -> Dict[str, Any]:
    """
    Apply add/update/forget operations returned by the extractor model.
    """
    mem = dict(memory or {})

    for item in (patch.get("memories_to_add") or []):
        k = (item.get("key") or "").strip()
        v = _to_clean_text(item.get("value"))
        try:
            conf = float(item.get("confidence") or 0.0)
        except Exception:
            conf = 0.0
        if k and v and conf >= min_confidence:
            if not mem.get(k):
                mem[k] = v

    for item in (patch.get("memories_to_update") or []):
        k = (item.get("key") or "").strip()
        new_v = _to_clean_text(item.get("new_value"))
        try:
            conf = float(item.get("confidence") or 0.0)
        except Exception:
            conf = 0.0
        if k and new_v and conf >= min_confidence:
            mem[k] = new_v

    for item in (patch.get("memories_to_forget") or []):
        k = (item.get("key") or "").strip()
        if k in mem:
            mem.pop(k, None)

    return mem


# ============ Memory Summary (Incremental) ============

def summarize_existing_memory_incremental(
        client,
        model: str,
        memory: Dict[str, Any],
        previous_summary: str = "",
) -> str:
    """
    Incrementally update a human-readable summary using the previous summary + latest memory.
    Glue code intentionally minimal; constraints live in the prompt.
    """
    if not memory or not isinstance(memory, dict):
        return previous_summary.strip() or "(no long-term memory yet)"

    prev = (previous_summary or "").strip()
    payload = json.dumps(memory, ensure_ascii=False)

    user_content = (
        f"[PREVIOUS SUMMARY]\n{prev if prev else '(none)'}\n\n"
        f"[LATEST LONG-TERM MEMORY JSON]\n{payload}"
    )

    try:
        resp = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": INCREMENTAL_MEMORY_SUMMARY_PROMPT},
                {"role": "user", "content": user_content},
            ],
            temperature=0.2,
        )
        text = (resp.choices[0].message.content or "").strip()
        return text or (prev or "(no summary generated)")
    except Exception:
        return prev or "(memory summary unavailable)"


def update_memory_summary_field(
        client,
        model: str,
        user_id: str,
        memory: Dict[str, Any],
        memory_dir: str = DEFAULT_MEMORY_DIR,
        persist: bool = True,
) -> Dict[str, Any]:
    """
    Store an incremental summary into memory under: meta.memory_summary
    """
    mem = dict(memory or {})

    previous_summary = ""
    if isinstance(mem.get("meta.memory_summary"), str):
        previous_summary = mem.get("meta.memory_summary", "").strip()

    summary = summarize_existing_memory_incremental(
        client=client,
        model=model,
        memory=mem,
        previous_summary=previous_summary,
    )

    mem["meta.memory_summary"] = summary

    if persist:
        save_memory(user_id, mem, memory_dir=memory_dir)

    return mem


# ============ Session-level Extraction + Persistence ============

def extract_and_update_memory_session_level(
        client,
        model: str,
        user_id: str,
        messages: List[dict],
        current_memory: Optional[Dict[str, Any]] = None,
        max_turns: int = 80,
) -> Dict[str, Any]:
    """
    Session-level memory update:
    - Build session snippet
    - Call LLM once at end of session to extract patch
    - Apply patch and persist
    - Update incremental memory summary and persist
    """
    memory = dict(current_memory or {})
    snippet = build_session_snippet(messages, max_turns=max_turns)
    if not snippet:
        return memory

    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": KEY_MEMORY_EXTRACTION_PROMPT},
            {"role": "user", "content": snippet},
        ],
        temperature=0.2,
    )

    raw = (resp.choices[0].message.content or "").strip()
    patch = safe_json_load(raw)

    if not patch or not patch.get("should_write_memory"):
        # Even if no new memory, we keep whatever exists (no summary refresh)
        return memory

    new_memory = apply_memory_patch(memory, patch)
    save_memory(user_id, new_memory)

    # Incremental summary (kept inside memory as meta.memory_summary)
    new_memory = update_memory_summary_field(
        client=client,
        model=model,
        user_id=user_id,
        memory=new_memory,
        persist=True,
    )

    return new_memory