85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""Persistent storage for flushed thread summaries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ThreadSummary:
|
|
"""A persisted summary of a flushed Telegram message thread."""
|
|
|
|
chat_id: int
|
|
thread_id: int
|
|
summary: str
|
|
message_count: int
|
|
flushed_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
|
|
|
|
@property
|
|
def key(self) -> str:
|
|
return f"{self.chat_id}:{self.thread_id}"
|
|
|
|
def format_for_telegram(self) -> str:
|
|
"""Return a concise Telegram-formatted recall card."""
|
|
flushed = self.flushed_at[:19].replace("T", " ")
|
|
header = (
|
|
f"\U0001f9e0 *Thread summary* (thread `{self.thread_id}`)\n"
|
|
f"_Flushed: {flushed} UTC \u2014 {self.message_count} messages_"
|
|
)
|
|
return f"{header}\n\n{self.summary}"
|
|
|
|
|
|
class ThreadMemoryStore:
|
|
"""JSON-backed store for flushed thread summaries.
|
|
|
|
The store is intentionally lightweight for the MVP. Each flush overwrites
|
|
any previous summary for the same (chat_id, thread_id) pair.
|
|
"""
|
|
|
|
def __init__(self, path: str | Path = "thread_memory.json") -> None:
|
|
self._path = Path(path)
|
|
self._data: dict[str, dict[str, int | str]] = self._load()
|
|
|
|
def _load(self) -> dict[str, dict[str, int | str]]:
|
|
if self._path.exists():
|
|
try:
|
|
raw = json.loads(self._path.read_text(encoding="utf-8"))
|
|
if isinstance(raw, dict):
|
|
return raw # type: ignore[return-value]
|
|
except (json.JSONDecodeError, OSError):
|
|
logger.warning(
|
|
"Could not read thread memory store at %s; starting fresh", self._path
|
|
)
|
|
return {}
|
|
|
|
def _save(self) -> None:
|
|
try:
|
|
self._path.write_text(
|
|
json.dumps(self._data, indent=2, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
except OSError:
|
|
logger.exception("Failed to write thread memory store to %s", self._path)
|
|
|
|
def save(self, summary: ThreadSummary) -> None:
|
|
"""Persist a thread summary, replacing any previous entry for this thread."""
|
|
self._data[summary.key] = asdict(summary)
|
|
self._save()
|
|
|
|
def get(self, chat_id: int, thread_id: int) -> ThreadSummary | None:
|
|
"""Return the stored summary for a thread, or None if not found."""
|
|
raw = self._data.get(f"{chat_id}:{thread_id}")
|
|
if raw is None:
|
|
return None
|
|
return ThreadSummary(**raw) # type: ignore[arg-type]
|
|
|
|
def all(self) -> list[ThreadSummary]:
|
|
"""Return all stored summaries, newest first."""
|
|
entries = [ThreadSummary(**v) for v in self._data.values()] # type: ignore[arg-type]
|
|
return sorted(entries, key=lambda s: s.flushed_at, reverse=True)
|