Part D: Memory

"""Module 7: session memory and a long-term memory store for the support assistant.

Long-term memories are facts about a customer that outlive one ticket
("prefers English", "is the billing admin"). This file extracts them from a
conversation with a model call, validates and redacts them, stores them with a
timestamp and a source, retrieves the relevant ones for a new ticket, resolves
conflicts by source trust and recency, forgets on request, and honors opt-out.
Run:  PYTHONPATH=. python examples/m07_memory.py
"""
from __future__ import annotations

import json
import re
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field, ValidationError

from supportdesk.kb_search import tokenize
from supportdesk.stand_in import ScriptedLLM

Source = Literal["crm", "customer_said", "agent_note", "model_inferred"]
TRUST: dict[str, int] = {"crm": 3, "customer_said": 2, "agent_note": 2, "model_inferred": 1}
PINNED_KEYS = {"reply_language", "role"}          # always injected when present
TTL = {"open_issue": timedelta(days=30)}          # keys that go stale; others never expire on their own

EXTRACT_PROMPT = """Extract durable facts about this customer from the conversation below.
Return a JSON list. Each item: {"key": snake_case slot, "value": short value, "text": one sentence,
"source": "customer_said" if the customer stated it, "model_inferred" if you are guessing}.
Only facts useful for future tickets: preferences, role, plan details, open issues.
Never include passwords, card numbers, 2FA codes, or health, religion, or other sensitive data.
Return [] if nothing qualifies.

Conversation:
"""

REDACT = [
    (re.compile(r"\b\d(?:[ -]?\d){12,18}\b"), "[card]"),
    (re.compile(r"(?i)\b(password|passwort|contraseña)\b\s*(is|:)?\s*\S+"), "[secret]"),
    (re.compile(r"(?<![-\w])\d{6}\b"), "[code]"),  # 2FA codes, but not the tail of INV-2026-004977
]


class ExtractedFact(BaseModel):
    """What we accept from the extractor model. Anything else is rejected."""
    key: str = Field(pattern=r"^[a-z][a-z0-9_]{1,40}$")
    value: str = Field(min_length=1, max_length=120)
    text: str = Field(min_length=3, max_length=240)
    source: Literal["customer_said", "model_inferred"]


@dataclass
class Memory:
    customer_id: str
    key: str
    value: str
    text: str
    source: str
    created_at: datetime
    ticket_id: str
    active: bool = True

    def line(self) -> str:
        return f"{self.text} ({self.source}, {self.created_at:%Y-%m-%d}, {self.ticket_id})"


@dataclass
class SessionMemory:
    """Scratch state for one ticket. Cleared when the ticket closes; never shared across customers."""
    ticket_id: str
    notes: list[str] = field(default_factory=list)
    confirmed: dict[str, str] = field(default_factory=dict)

    def as_state(self) -> dict[str, Any]:
        return {"ticket": self.ticket_id, "confirmed": self.confirmed, "notes": self.notes[-5:]}


def redact(text: str) -> tuple[str, int]:
    count = 0
    for pattern, replacement in REDACT:
        text, n = pattern.subn(replacement, text)
        count += n
    return text, count


class MemoryStore:
    def __init__(self, path: Path | None = None) -> None:
        self.path = path
        self.memories: list[Memory] = []
        self.opted_out: set[str] = set()
        self.audit: list[dict[str, str]] = []  # what happened, never the forgotten content

    # --- writing ----------------------------------------------------------------------------------
    def add(self, memory: Memory) -> Memory:
        """Store one memory, then re-resolve the active winner for its key."""
        if memory.customer_id in self.opted_out:
            self._log("skipped_opt_out", memory.customer_id, memory.key)
            return memory
        memory.text, n = redact(memory.text)
        memory.value, m = redact(memory.value)
        if n + m:
            self._log("redacted", memory.customer_id, memory.key)
        self.memories.append(memory)
        self._resolve(memory.customer_id, memory.key)
        return memory

    def extract(self, conversation: list[dict[str, str]], customer_id: str, ticket_id: str, when: datetime,
                chat: Callable[..., Any]) -> tuple[list[Memory], list[str]]:
        """Ask a model for facts, validate each one, store the valid ones. Returns (stored, rejected reasons)."""
        if customer_id in self.opted_out:
            return [], ["customer opted out of memory"]
        transcript = "\n".join(f"{m['role']}: {redact(m['content'])[0]}" for m in conversation)
        reply = chat([{"role": "user", "content": EXTRACT_PROMPT + transcript}], temperature=0.0, max_tokens=400)
        try:
            items = json.loads(reply.text)
            if not isinstance(items, list):
                raise ValueError("not a list")
        except (json.JSONDecodeError, ValueError) as err:
            return [], [f"unparseable extractor output: {err}"]
        stored, rejected = [], []
        for item in items:
            try:
                fact = ExtractedFact.model_validate(item)
            except ValidationError as err:
                rejected.append(f"{str(item)[:60]}: {err.errors()[0]['msg']}")
                continue
            if redact(fact.value)[1]:
                rejected.append(f"{fact.key}: value is sensitive, not stored")
                self._log("refused_sensitive", customer_id, fact.key)
                continue
            stored.append(self.add(Memory(customer_id, fact.key, fact.value, fact.text, fact.source, when, ticket_id)))
        return stored, rejected

    def _resolve(self, customer_id: str, key: str) -> None:
        """For one slot, the most trusted source wins; among equals, the most recent wins."""
        same = [m for m in self.memories if m.customer_id == customer_id and m.key == key]
        winner = max(same, key=lambda m: (TRUST[m.source], m.created_at))
        for m in same:
            m.active = m is winner

    # --- reading ----------------------------------------------------------------------------------
    def retrieve(self, customer_id: str, query: str, now: datetime, k: int = 3) -> list[Memory]:
        """Pinned slots always, plus the k active memories that share the most words with the query."""
        if customer_id in self.opted_out:
            return []
        live = [m for m in self.memories if m.customer_id == customer_id and m.active
                and not (m.key in TTL and now - m.created_at > TTL[m.key])]
        pinned = [m for m in live if m.key in PINNED_KEYS]
        q = set(tokenize(query))

        def score(m: Memory) -> tuple[float, datetime]:
            words = set(tokenize(m.text + " " + m.key.replace("_", " ")))
            return (len(q & words) / (len(words) ** 0.5 or 1), m.created_at)

        relevant = sorted((m for m in live if m.key not in PINNED_KEYS), key=score, reverse=True)
        relevant = [m for m in relevant if score(m)[0] > 0][:k]
        return pinned + relevant

    # --- user control -------------------------------------------------------------------------------
    def forget(self, customer_id: str, key: str | None = None) -> int:
        """Delete one slot or everything for a customer. Deleted means gone, including superseded versions."""
        before = len(self.memories)
        self.memories = [m for m in self.memories
                         if not (m.customer_id == customer_id and (key is None or m.key == key))]
        removed = before - len(self.memories)
        self._log("forgot", customer_id, key or "*", str(removed))
        return removed

    def opt_out(self, customer_id: str) -> int:
        self.opted_out.add(customer_id)
        return self.forget(customer_id)

    def export(self, customer_id: str) -> list[dict[str, Any]]:
        """Everything stored about one customer, for an access request."""
        return [{**asdict(m), "created_at": m.created_at.isoformat()} for m in self.memories if m.customer_id == customer_id]

    def _log(self, action: str, customer_id: str, key: str, detail: str = "") -> None:
        self.audit.append({"action": action, "customer": customer_id, "key": key, "detail": detail})

    def save(self) -> None:
        if self.path:
            data = {"memories": [{**asdict(m), "created_at": m.created_at.isoformat()} for m in self.memories],
                    "opted_out": sorted(self.opted_out), "audit": self.audit}
            self.path.write_text(json.dumps(data, ensure_ascii=False, indent=1))


# Scripted extractor replies. These are hand-written stand-ins for what a model might return,
# including two bad items, so the validation and redaction paths run. NOT model output.
SCRIPTED_EXTRACTIONS = [
    json.dumps([
        {"key": "role", "value": "billing admin", "text": "The customer is the workspace billing admin.", "source": "customer_said"},
        {"key": "reply_language", "value": "en", "text": "Prefers replies in English.", "source": "customer_said"},
        {"key": "open_issue", "value": "annual renewal refund INV-2026-004977", "text": "Asked for a refund of the 16 September annual renewal, invoice INV-2026-004977.", "source": "customer_said"},
        {"key": "finance_contact_language", "value": "es", "text": "Their finance lead Carla prefers Spanish.", "source": "model_inferred"},
    ]),
    json.dumps([
        {"key": "reply_language", "value": "es", "text": "Seems to prefer Spanish.", "source": "model_inferred"},
        {"key": "payment_card", "value": "4111 1111 1111 1111", "text": "Pays with card 4111 1111 1111 1111.", "source": "customer_said"},
        {"key": "Mood", "value": "annoyed", "text": "Customer is annoyed.", "source": "model_inferred"},
        {"key": "slack_channel", "value": "#finance-ops", "text": "Posts Brightlane notifications to Slack channel #finance-ops.", "source": "customer_said"},
    ]),
    json.dumps([
        {"key": "reply_language", "value": "es", "text": "Asked to switch replies to Spanish from now on.", "source": "customer_said"},
    ]),
    "Sure! Here are the facts I found: the customer likes English.",
]

CONVERSATIONS = [
    ("T-1004", datetime(2026, 9, 18, 10, 0), [
        {"role": "user", "content": "We renewed our annual Business plan by mistake, invoice INV-2026-004977, renewed 16 September. I'm the billing admin. Please write to me in English even though our finance lead Carla prefers Spanish."}]),
    ("T-2210", datetime(2026, 9, 19, 9, 30), [
        {"role": "user", "content": "Hola, the Slack notifications to #finance-ops stopped. My card 4111 1111 1111 1111 was charged too, whatever."}]),
    ("T-2231", datetime(2026, 9, 20, 16, 5), [
        {"role": "user", "content": "From now on please answer me in Spanish, my colleague will read the replies."}]),
    ("T-2240", datetime(2026, 9, 21, 8, 0), [
        {"role": "user", "content": "Another question about invoices."}]),
]


if __name__ == "__main__":
    store = MemoryStore()
    extractor = ScriptedLLM(replies=list(SCRIPTED_EXTRACTIONS))  # plumbing test, not a model
    customer = "cust-acme"
    for ticket_id, when, convo in CONVERSATIONS:
        stored, rejected = store.extract(convo, customer, ticket_id, when, extractor)
        print(f"{ticket_id} {when:%Y-%m-%d}: stored {len(stored)}, rejected {len(rejected)}")
        for m in stored:
            print(f"  + {m.key} = {m.value!r} [{m.source}]{'' if m.active else ' (inactive: outranked)'}")
        for r in rejected:
            print(f"  - rejected: {r}")

    print("\nActive vs superseded for reply_language:")
    for m in store.memories:
        if m.key == "reply_language":
            print(f"  {'ACTIVE ' if m.active else 'old    '} {m.value} {m.source:14} {m.created_at:%Y-%m-%d %H:%M} {m.ticket_id}")

    for now, query in [(datetime(2026, 9, 21, 9, 0), "Where do I download last month's invoice?"),
                       (datetime(2026, 9, 21, 9, 0), "Slack stopped posting notifications"),
                       (datetime(2026, 11, 2, 9, 0), "Did my refund go through?")]:
        print(f"\nretrieve at {now:%Y-%m-%d} for {query!r}:")
        for m in store.retrieve(customer, query, now):
            print(f"  {m.key:15} {m.line()}")

    print("\nWhat the extractor model was shown for T-2210 (redacted before the call):")
    print(" ", extractor.calls[1]["messages"][0]["content"].split("Conversation:\n")[1])

    print("\nCustomer asks: 'forget my Slack setup'")
    print("  removed:", store.forget(customer, "slack_channel"))
    print("Customer opts out of memory entirely")
    print("  removed:", store.opt_out(customer))
    stored, rejected = store.extract(CONVERSATIONS[0][2], customer, "T-2250", datetime(2026, 9, 21, 12, 0),
                                     ScriptedLLM(replies=[SCRIPTED_EXTRACTIONS[0]]))
    print("  extract after opt-out:", len(stored), rejected)
    print("  export:", store.export(customer))
    print("\nAudit log (no forgotten content is kept):")
    for entry in store.audit:
        print(" ", entry)

  + open_issue = 'annual renewal refund INV-2026-[code]' [customer_said]