62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
"""
|
|
Per-user conversation history store.
|
|
|
|
Each Discord user gets their own isolated message list. Only the last
|
|
`max_history` messages are kept — older ones are silently dropped so the
|
|
LLM context stays small.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Dict, List
|
|
|
|
logger = logging.getLogger("bot.conversation")
|
|
|
|
# role we assign to user messages inside the OpenAI-style message list
|
|
_USER_ROLE = "user"
|
|
# role we assign to bot responses
|
|
_ASSISTANT_ROLE = "assistant"
|
|
|
|
|
|
class ConversationStore:
|
|
"""Thread-safe-ish in-memory store keyed by Discord user ID (int)."""
|
|
|
|
def __init__(self, max_history: int = 7) -> None:
|
|
self._max = max_history
|
|
self._store: Dict[int, List[dict]] = {}
|
|
|
|
# ------------------------------------------------------------------
|
|
# public API
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_history(self, user_id: int) -> list[dict]:
|
|
"""Return the last *max_history* messages for *user_id*."""
|
|
return list(self._store.get(user_id, []))
|
|
|
|
def append(self, user_id: int, user_msg: str, assistant_reply: str) -> None:
|
|
"""Store the user message + assistant reply, then trim to max."""
|
|
if user_id not in self._store:
|
|
self._store[user_id] = []
|
|
|
|
history = self._store[user_id]
|
|
history.append({"role": _USER_ROLE, "content": user_msg})
|
|
history.append({"role": _ASSISTANT_ROLE, "content": assistant_reply})
|
|
|
|
# Trim oldest messages if we exceeded the limit
|
|
while len(history) > self._max:
|
|
history.pop(0)
|
|
|
|
def clear(self, user_id: int) -> None:
|
|
"""Wipe the conversation for a user."""
|
|
self._store.pop(user_id, None)
|
|
logger.info("Cleared conversation for user %s", user_id)
|
|
|
|
# ------------------------------------------------------------------
|
|
# debug / introspection
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def user_count(self) -> int:
|
|
return len(self._store)
|