feat: thread-aware memory with /flush and /recall commands
This commit is contained in:
committed by
GitHub
parent
8759bea963
commit
c74c8b0d8b
+194
-24
@@ -15,13 +15,41 @@ from telegram.ext import (
|
||||
|
||||
from steward.config import Settings
|
||||
from steward.llm.client import LLMClient
|
||||
from steward.memory.thread_store import ThreadMemoryStore, ThreadSummary
|
||||
from steward.proposals.generator import Proposal, ProposalGenerator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-user conversation history (in-memory for the MVP)
|
||||
# Per-user conversation history for non-threaded messages (capped at _MAX_HISTORY turns).
|
||||
_history: dict[int, list[dict[str, str]]] = defaultdict(list)
|
||||
_MAX_HISTORY = 20 # keep the last N turns per user
|
||||
_MAX_HISTORY = 20
|
||||
|
||||
# Per-thread conversation history: (chat_id, thread_id) → full history (unbounded).
|
||||
# Messages belonging to a Telegram message thread are kept in their entirety here
|
||||
# until explicitly flushed by the /flush command.
|
||||
_thread_history: dict[tuple[int, int], list[dict[str, str]]] = defaultdict(list)
|
||||
|
||||
_FLUSH_SYSTEM_PROMPT = (
|
||||
"You are Steward. The following is a complete Telegram message thread conversation. "
|
||||
"Produce a concise but comprehensive summary that captures:\n"
|
||||
"- The main topics discussed\n"
|
||||
"- Key decisions or conclusions reached\n"
|
||||
"- Any outstanding actions or open questions\n"
|
||||
"- Important context that would help recall this conversation later\n\n"
|
||||
"Be precise. Omit pleasantries."
|
||||
)
|
||||
|
||||
|
||||
def _thread_key(update: Update) -> tuple[int, int] | None:
|
||||
"""Return the (chat_id, thread_id) key if the message is part of a thread, else None."""
|
||||
msg = update.message
|
||||
chat = update.effective_chat
|
||||
if msg is None or chat is None:
|
||||
return None
|
||||
thread_id = msg.message_thread_id
|
||||
if thread_id is None:
|
||||
return None
|
||||
return (chat.id, thread_id)
|
||||
|
||||
|
||||
def _is_allowed(user_id: int, settings: Settings) -> bool:
|
||||
@@ -32,7 +60,7 @@ def _is_allowed(user_id: int, settings: Settings) -> bool:
|
||||
|
||||
|
||||
async def _send_long(update: Update, text: str) -> None:
|
||||
"""Send text, splitting if it exceeds Telegram's 4096-char limit."""
|
||||
"""Send text, splitting across messages if it exceeds Telegram's 4096-char limit."""
|
||||
limit = 4096
|
||||
for i in range(0, len(text), limit):
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
@@ -53,7 +81,9 @@ async def start_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> N
|
||||
"I'm your AI-assisted personal operations platform.\n"
|
||||
"Talk to me naturally, or use:\n"
|
||||
"/help – show available commands\n"
|
||||
"/clear – reset our conversation history\n"
|
||||
"/clear – reset conversation history\n"
|
||||
"/flush – summarise and archive this thread's memory\n"
|
||||
"/recall – retrieve archived thread summaries\n"
|
||||
"/analyse – run a manual API analysis right now",
|
||||
parse_mode=ParseMode.MARKDOWN,
|
||||
)
|
||||
@@ -70,21 +100,150 @@ async def help_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> No
|
||||
"*Steward commands*\n\n"
|
||||
"/start – greeting\n"
|
||||
"/help – this message\n"
|
||||
"/clear – reset conversation history\n"
|
||||
"/clear – reset conversation history for this context\n"
|
||||
"/flush – summarise the current thread, store the summary, and compress memory\n"
|
||||
" _(only available inside a message thread)_\n"
|
||||
"/recall – show the stored summary for this thread, or list all summaries\n"
|
||||
"/analyse – trigger an immediate API analysis and proposal",
|
||||
parse_mode=ParseMode.MARKDOWN,
|
||||
)
|
||||
|
||||
|
||||
async def clear_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle /clear – wipe conversation history for this user."""
|
||||
"""Handle /clear – wipe conversation history for this context.
|
||||
|
||||
Inside a message thread: clears the thread's unbounded history.
|
||||
Outside a thread: clears the per-user capped history.
|
||||
"""
|
||||
settings: Settings = context.bot_data["settings"]
|
||||
user = update.effective_user
|
||||
if user is None or not _is_allowed(user.id, settings):
|
||||
return
|
||||
|
||||
_history[user.id].clear()
|
||||
await update.message.reply_text("Conversation history cleared.") # type: ignore[union-attr]
|
||||
key = _thread_key(update)
|
||||
if key is not None:
|
||||
_thread_history[key].clear()
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"Thread conversation history cleared."
|
||||
)
|
||||
else:
|
||||
_history[user.id].clear()
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"Conversation history cleared."
|
||||
)
|
||||
|
||||
|
||||
async def flush_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle /flush – summarise thread memory, persist it, and compress in-memory history.
|
||||
|
||||
Steps:
|
||||
1. Verify the command is issued inside a message thread.
|
||||
2. Summarise the full thread history via the LLM.
|
||||
3. Persist the summary in the ThreadMemoryStore (keyed by chat_id + thread_id).
|
||||
4. Replace the in-memory thread history with a single compressed context message
|
||||
so conversation can continue with the summary as background.
|
||||
"""
|
||||
settings: Settings = context.bot_data["settings"]
|
||||
llm: LLMClient = context.bot_data["llm"]
|
||||
store: ThreadMemoryStore = context.bot_data["thread_store"]
|
||||
user = update.effective_user
|
||||
|
||||
if user is None or not _is_allowed(user.id, settings):
|
||||
return
|
||||
|
||||
key = _thread_key(update)
|
||||
if key is None:
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"\u26a0\ufe0f /flush can only be used inside a message thread."
|
||||
)
|
||||
return
|
||||
|
||||
chat_id, thread_id = key
|
||||
history = _thread_history[key]
|
||||
|
||||
if not history:
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"This thread has no conversation history to flush."
|
||||
)
|
||||
return
|
||||
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"\U0001f4be Summarising thread memory\u2026"
|
||||
)
|
||||
|
||||
# Build a readable transcript for the LLM to summarise
|
||||
transcript_lines = []
|
||||
for msg in history:
|
||||
role_label = "User" if msg["role"] == "user" else "Steward"
|
||||
transcript_lines.append(f"{role_label}: {msg['content']}")
|
||||
transcript = "\n".join(transcript_lines)
|
||||
|
||||
summary_text = await llm.chat(
|
||||
f"Thread transcript:\n\n{transcript}",
|
||||
system_prompt=_FLUSH_SYSTEM_PROMPT,
|
||||
)
|
||||
|
||||
message_count = sum(1 for m in history if m["role"] == "user")
|
||||
thread_summary = ThreadSummary(
|
||||
chat_id=chat_id,
|
||||
thread_id=thread_id,
|
||||
summary=summary_text,
|
||||
message_count=message_count,
|
||||
)
|
||||
store.save(thread_summary)
|
||||
|
||||
# Compress: replace history with a single system-context entry so the thread
|
||||
# can continue with the summary as background knowledge.
|
||||
_thread_history[key] = [
|
||||
{"role": "system", "content": f"Summary of earlier conversation:\n{summary_text}"}
|
||||
]
|
||||
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
f"\u2705 Thread memory flushed and stored "
|
||||
f"(thread `{thread_id}`, {message_count} messages summarised).",
|
||||
parse_mode=ParseMode.MARKDOWN,
|
||||
)
|
||||
|
||||
|
||||
async def recall_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle /recall – retrieve stored thread summaries.
|
||||
|
||||
Inside a thread: shows the stored summary for this thread (if any).
|
||||
Outside a thread: lists all stored summaries (newest first).
|
||||
"""
|
||||
settings: Settings = context.bot_data["settings"]
|
||||
store: ThreadMemoryStore = context.bot_data["thread_store"]
|
||||
user = update.effective_user
|
||||
|
||||
if user is None or not _is_allowed(user.id, settings):
|
||||
return
|
||||
|
||||
key = _thread_key(update)
|
||||
if key is not None:
|
||||
chat_id, thread_id = key
|
||||
stored = store.get(chat_id, thread_id)
|
||||
if stored is None:
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"No stored summary for this thread yet. Use /flush to create one."
|
||||
)
|
||||
else:
|
||||
await _send_long(update, stored.format_for_telegram())
|
||||
return
|
||||
|
||||
# Outside a thread: list all stored summaries
|
||||
all_summaries = store.all()
|
||||
if not all_summaries:
|
||||
await update.message.reply_text( # type: ignore[union-attr]
|
||||
"No thread summaries stored yet. Use /flush inside a message thread."
|
||||
)
|
||||
return
|
||||
|
||||
lines = ["\U0001f4da *Stored thread summaries*\n"]
|
||||
for s in all_summaries:
|
||||
date = s.flushed_at[:10]
|
||||
first_line = s.summary.split("\n")[0][:80]
|
||||
lines.append(f"\u2022 Thread `{s.thread_id}` ({date}): {first_line}\u2026")
|
||||
await _send_long(update, "\n".join(lines))
|
||||
|
||||
|
||||
async def analyse_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
@@ -109,7 +268,11 @@ async def analyse_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
|
||||
|
||||
|
||||
async def message_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handle plain text messages – forward to LLM and reply."""
|
||||
"""Handle plain text messages – forward to LLM and reply.
|
||||
|
||||
Thread messages: history is stored unbounded under the (chat_id, thread_id) key.
|
||||
Non-thread messages: history is capped at _MAX_HISTORY turns per user.
|
||||
"""
|
||||
settings: Settings = context.bot_data["settings"]
|
||||
llm: LLMClient = context.bot_data["llm"]
|
||||
user = update.effective_user
|
||||
@@ -120,32 +283,39 @@ async def message_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
|
||||
if not text:
|
||||
return
|
||||
|
||||
history = _history[user.id]
|
||||
reply = await llm.chat(text, history=history)
|
||||
|
||||
# Update history
|
||||
history.append({"role": "user", "content": text})
|
||||
history.append({"role": "assistant", "content": reply})
|
||||
# Trim to keep only the most recent turns (2 messages per turn)
|
||||
if len(history) > _MAX_HISTORY * 2:
|
||||
_history[user.id] = history[-( _MAX_HISTORY * 2):]
|
||||
key = _thread_key(update)
|
||||
if key is not None:
|
||||
# Thread message: unbounded history
|
||||
history = _thread_history[key]
|
||||
reply = await llm.chat(text, history=history)
|
||||
history.append({"role": "user", "content": text})
|
||||
history.append({"role": "assistant", "content": reply})
|
||||
else:
|
||||
# Non-thread message: capped history per user
|
||||
history = _history[user.id]
|
||||
reply = await llm.chat(text, history=history)
|
||||
history.append({"role": "user", "content": text})
|
||||
history.append({"role": "assistant", "content": reply})
|
||||
if len(history) > _MAX_HISTORY * 2:
|
||||
_history[user.id] = history[-(_MAX_HISTORY * 2) :]
|
||||
|
||||
await _send_long(update, reply)
|
||||
|
||||
|
||||
def build_application(settings: Settings, llm: LLMClient) -> Application: # type: ignore[type-arg]
|
||||
def build_application(
|
||||
settings: Settings, llm: LLMClient, thread_store: ThreadMemoryStore | None = None
|
||||
) -> Application: # type: ignore[type-arg]
|
||||
"""Build and return the Telegram Application."""
|
||||
app = (
|
||||
Application.builder()
|
||||
.token(settings.telegram_bot_token)
|
||||
.build()
|
||||
)
|
||||
app = Application.builder().token(settings.telegram_bot_token).build()
|
||||
app.bot_data["settings"] = settings
|
||||
app.bot_data["llm"] = llm
|
||||
app.bot_data["thread_store"] = thread_store or ThreadMemoryStore(settings.thread_memory_path)
|
||||
|
||||
app.add_handler(CommandHandler("start", start_handler))
|
||||
app.add_handler(CommandHandler("help", help_handler))
|
||||
app.add_handler(CommandHandler("clear", clear_handler))
|
||||
app.add_handler(CommandHandler("flush", flush_handler))
|
||||
app.add_handler(CommandHandler("recall", recall_handler))
|
||||
app.add_handler(CommandHandler("analyse", analyse_handler))
|
||||
app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, message_handler))
|
||||
|
||||
|
||||
@@ -29,6 +29,9 @@ class Settings(BaseSettings):
|
||||
analysis_cron_hour: int = 8
|
||||
analysis_cron_minute: int = 0
|
||||
|
||||
# Thread memory
|
||||
thread_memory_path: str = "thread_memory.json"
|
||||
|
||||
|
||||
def get_settings() -> Settings:
|
||||
"""Return application settings singleton."""
|
||||
|
||||
+3
-1
@@ -8,6 +8,7 @@ from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from steward.bot.telegram import build_application, send_proposal
|
||||
from steward.config import get_settings
|
||||
from steward.llm.client import LLMClient
|
||||
from steward.memory.thread_store import ThreadMemoryStore
|
||||
from steward.proposals.generator import ProposalGenerator
|
||||
|
||||
logging.basicConfig(
|
||||
@@ -44,7 +45,8 @@ def main() -> None:
|
||||
sys.exit(1)
|
||||
|
||||
llm = LLMClient(settings)
|
||||
app = build_application(settings, llm)
|
||||
thread_store = ThreadMemoryStore(settings.thread_memory_path)
|
||||
app = build_application(settings, llm, thread_store)
|
||||
generator = ProposalGenerator(settings, llm)
|
||||
|
||||
scheduler = AsyncIOScheduler()
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user