Memory.py

"""
STEW Memory Engine
Stores per-user conversation history and session context in memory.
On HF Spaces this resets on restart; for persistence plug in a Redis/DB backend.
"""
from collections import defaultdict, deque
from datetime import datetime
from config import Config


class ConversationMemory:
    """Lightweight in-process memory for multi-turn conversations."""

    def __init__(self):
        # { user_id: deque of {"role": ..., "content": ..., "ts": ...} }
        self._history: dict[str, deque] = defaultdict(
            lambda: deque(maxlen=Config.MAX_MEMORY_TURNS * 2)
        )
        # { user_id: dict of arbitrary facts }
        self._context: dict[str, dict] = defaultdict(dict)

    # ── History ──────────────────────────────────────────────────────────────

    def add(self, user_id: str, user_msg: str, assistant_msg: str) -> None:
        """Record a completed exchange."""
        ts = datetime.utcnow().isoformat()
        self._history[user_id].append(
            {"role": "user",      "content": user_msg,      "ts": ts}
        )
        self._history[user_id].append(
            {"role": "assistant", "content": assistant_msg, "ts": ts}
        )

    def get_messages(self, user_id: str) -> list[dict]:
        """Return history formatted for the OpenAI messages array."""
        return [
            {"role": m["role"], "content": m["content"]}
            for m in self._history[user_id]
        ]

    def clear(self, user_id: str) -> None:
        """Wipe history for a user."""
        self._history[user_id].clear()
        self._context[user_id].clear()

    # ── Context facts ─────────────────────────────────────────────────────────

    def set_fact(self, user_id: str, key: str, value: str) -> None:
        self._context[user_id][key] = value

    def get_context_string(self, user_id: str) -> str:
        ctx = self._context[user_id]
        if not ctx:
            return ""
        lines = "\n".join(f"- {k}: {v}" for k, v in ctx.items())
        return f"Known facts about this user:\n{lines}"

    def summary(self, user_id: str) -> dict:
        return {
            "turns": len(self._history[user_id]) // 2,
            "facts": len(self._context[user_id]),
        }