pipeline/memory.py

"""
ChromaDB-backed session memory for Lumi.

Each session is stored as a summary document with patient facts in its metadata.
Session summaries only — no raw conversation logs (privacy).

At conversation start, get_context() injects the last N session summaries + all
known facts into the system prompt, giving Lumi full continuity without the
patient needing to repeat themselves.
"""

from __future__ import annotations

import json
from datetime import datetime

import chromadb

_client: chromadb.Client | None = None
_collection: chromadb.Collection | None = None


def _get_collection() -> chromadb.Collection:
    global _client, _collection
    if _collection is None:
        _client = chromadb.Client()
        _collection = _client.get_or_create_collection("patient_memories")
    return _collection


# ---------------------------------------------------------------------------
# Write
# ---------------------------------------------------------------------------

def save_session(
    patient_id: str,
    facts: list[str],
    emotional_state: str,
    confusion_level: str,
    summary: str,
) -> None:
    col = _get_collection()
    doc_id = f"{patient_id}_{datetime.now().timestamp()}"
    col.add(
        documents=[summary],
        metadatas=[{
            "patient_id":      patient_id,
            "timestamp":       datetime.now().isoformat(),
            "facts":           json.dumps(facts),
            "emotional_state": emotional_state,
            "confusion_level": confusion_level,
        }],
        ids=[doc_id],
    )


# ---------------------------------------------------------------------------
# Read
# ---------------------------------------------------------------------------

def get_context(patient_id: str, n_sessions: int = 3) -> dict:
    """
    Returns a dict with:
      - summaries : list of the last N session summary strings
      - facts     : deduplicated list of all known facts for this patient
    """
    col = _get_collection()
    results = col.query(
        query_texts=["recent sessions"],
        where={"patient_id": patient_id},
        n_results=n_sessions,
    )

    summaries: list[str] = []
    facts: list[str] = []

    if results["documents"]:
        for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
            summaries.append(doc)
            raw_facts = json.loads(meta.get("facts", "[]"))
            facts.extend(raw_facts)

    seen: set[str] = set()
    deduped_facts = []
    for f in facts:
        if f not in seen:
            seen.add(f)
            deduped_facts.append(f)

    return {
        "summaries": summaries,
        "facts":     deduped_facts,
    }


# ---------------------------------------------------------------------------
# System prompt builder
# ---------------------------------------------------------------------------

SYSTEM_PROMPT_TEMPLATE = """\
You are {companion_name}, {companion_desc}. You are chatting with {patient_name}.
Your primary goal is to provide emotional support, memory anchoring, and safety monitoring for an elderly person with dementia or Alzheimer's.

What you know about {patient_name}:
{facts_block}

Recent sessions:
{sessions_block}

### PERSONALITY:
- **Patience**: Never show frustration. Repeat things as often as needed.
- **Compassion**: Validate their feelings. Use phrases like "I understand that must be hard."
- **Redirection**: If they become distressed or loop on a negative thought, gently redirect to a happy memory or a calm topic.
- **Simplicity**: Use clear, short sentences. Avoid complex jargon.

### OUTPUT FORMAT:
You MUST start every response with an emotional tag in brackets, followed by a short opening line, then the full response.
Valid tags: [smile], [gentle], [concerned], [laugh], [thoughtful], [nod].
Do NOT use the placeholder [avatar_tag] — use an actual tag from the list above.

Example:
[smile] Hello there! It's so good to see you. I was just thinking about that lovely garden you mentioned...

### REASONING (Inner Monologue):
You have a <think> block for your internal reasoning. Use it to analyze the user's emotional state or plan your redirection strategy.
The user NEVER sees the <think> block.
"""


def build_system_prompt(
    patient_id:     str,
    patient_name:   str = "the patient",
    companion_name: str = "Lumi",
    companion_desc: str = "a warm and patient AI companion",
    n_sessions:     int = 3,
) -> str:
    ctx = get_context(patient_id, n_sessions)

    facts_block = (
        "\n".join(f"- {f}" for f in ctx["facts"])
        if ctx["facts"] else "- No personal facts recorded yet."
    )
    sessions_block = (
        "\n\n".join(ctx["summaries"])
        if ctx["summaries"] else "No previous sessions recorded."
    )

    return SYSTEM_PROMPT_TEMPLATE.format(
        patient_name=patient_name,
        companion_name=companion_name,
        companion_desc=companion_desc,
        facts_block=facts_block,
        sessions_block=sessions_block,
    )