File size: 2,485 Bytes
7b46e89 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
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]),
}
|