Added message uningest to reverse Markov data when recently ingested messages are deleted.

crabstero/bot.py

@@ -26,6 +26,7 @@ from discord import app_commands
 from discord.ext import commands
 
 from crabstero import __version__ as crabstero_version
+from crabstero.cache import IngestCache
 from crabstero.database import Database
 from crabstero.listeners import interaction, message, server_events
 from crabstero.metrics import (
@@ -83,6 +84,7 @@ class Crabstero(commands.Bot):
         ] = asyncio.Queue()
         self._ingestion_workers: list[asyncio.Task[None]] = []
         self.db: Database
+        self.ingest_cache = IngestCache()
 
         self._metrics_address = metrics_address
 
@@ -96,6 +98,7 @@ class Crabstero(commands.Bot):
     async def setup_hook(self) -> None:
         """Open the database, start ingestion workers, and load all cogs."""
         self.db = await Database.connect(self._database_path)
+        self.ingest_cache.start()
         self._start_ingestion_workers()
 
         if self._metrics_address is not None:
@@ -157,6 +160,7 @@ class Crabstero(commands.Bot):
             worker.cancel()
         await asyncio.gather(*self._ingestion_workers, return_exceptions=True)
         self._ingestion_workers.clear()
+        self.ingest_cache.stop()
         if hasattr(self, "_metrics_server"):
             await self._metrics_server.stop()
         if hasattr(self, "db"):

crabstero/cache.py

@@ -0,0 +1,114 @@
+# Copyright 2026 Logan Fick
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TTL cache for recently ingested messages, enabling uningest on delete."""
+
+import asyncio
+import time
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class CachedMessage:
+    """Data stored for a recently ingested message.
+
+    :param channel_id: The Discord channel ID the message was ingested into.
+    :param user_id: The Discord user ID of the message author.
+    :param content: The text content of the message, if any.
+    :param embed_texts: Titles and descriptions extracted from embeds.
+    :param image_urls: Image URLs extracted from embeds.
+    """
+
+    channel_id: int
+    user_id: int
+    content: str | None
+    embed_texts: list[str]
+    image_urls: list[str]
+
+
+class IngestCache:
+    """In-memory TTL cache mapping message IDs to their ingested data.
+
+    Entries expire after ``ttl_seconds`` and are evicted by a periodic
+    background task.
+
+    :param ttl_seconds: How long entries remain valid, in seconds.
+    :param cleanup_interval_seconds: How often the background task runs.
+    """
+
+    def __init__(
+        self,
+        ttl_seconds: float = 120,
+        cleanup_interval_seconds: float = 30,
+    ) -> None:
+        """Create a new cache with the given TTL and cleanup interval.
+
+        :param ttl_seconds: How long entries remain valid, in seconds.
+        :param cleanup_interval_seconds: How often the background task runs.
+        """
+        self._ttl = ttl_seconds
+        self._cleanup_interval = cleanup_interval_seconds
+        self._entries: dict[int, tuple[float, CachedMessage]] = {}
+        self._task: asyncio.Task[None] | None = None
+
+    def put(self, message_id: int, entry: CachedMessage) -> None:
+        """Store a cache entry for a message.
+
+        :param message_id: The Discord message ID.
+        :param entry: The cached message data.
+        """
+        self._entries[message_id] = (time.monotonic(), entry)
+
+    def pop(self, message_id: int) -> CachedMessage | None:
+        """Remove and return a cache entry if it exists and has not expired.
+
+        :param message_id: The Discord message ID.
+        :return: The cached message data, or None.
+        """
+        pair = self._entries.pop(message_id, None)
+        if pair is None:
+            return None
+        stored_at, entry = pair
+        if time.monotonic() - stored_at > self._ttl:
+            return None
+        return entry
+
+    def _cleanup(self) -> None:
+        """Remove all expired entries from the cache."""
+        now = time.monotonic()
+        expired = [
+            mid
+            for mid, (stored_at, _) in self._entries.items()
+            if now - stored_at > self._ttl
+        ]
+        for mid in expired:
+            del self._entries[mid]
+
+    def start(self) -> None:
+        """Start the periodic background cleanup task."""
+        if self._task is not None:
+            return
+        self._task = asyncio.create_task(self._cleanup_loop())
+
+    def stop(self) -> None:
+        """Stop the periodic background cleanup task."""
+        if self._task is not None:
+            self._task.cancel()
+            self._task = None
+
+    async def _cleanup_loop(self) -> None:
+        """Run cleanup on a fixed interval until cancelled."""
+        while True:
+            await asyncio.sleep(self._cleanup_interval)
+            self._cleanup()

crabstero/database.py

@@ -134,6 +134,42 @@ class Database:
         if start_words or transitions:
             await self._connection.commit()
 
+    async def remove_markov_data(
+        self,
+        start_words: list[tuple[int, int, str]],
+        transitions: list[tuple[int, int, str, str]],
+    ) -> None:
+        """Remove one matching row per entry from the Markov tables.
+
+        Each entry removes at most one duplicate row, preserving remaining
+        frequency weight.
+
+        :param start_words: A list of (channel_id, user_id, word) tuples.
+        :param transitions: A list of (channel_id, user_id, word, next_word) tuples.
+        """
+        for channel_id, user_id, word in start_words:
+            await self._connection.execute(
+                "DELETE FROM markov_start_words"
+                " WHERE rowid = ("
+                "   SELECT rowid FROM markov_start_words"
+                "   WHERE channel_id = ? AND user_id = ? AND word = ?"
+                "   LIMIT 1"
+                " )",
+                (channel_id, user_id, word),
+            )
+        for channel_id, user_id, word, next_word in transitions:
+            await self._connection.execute(
+                "DELETE FROM markov_transitions"
+                " WHERE rowid = ("
+                "   SELECT rowid FROM markov_transitions"
+                "   WHERE channel_id = ? AND user_id = ? AND word = ? AND next_word = ?"
+                "   LIMIT 1"
+                " )",
+                (channel_id, user_id, word, next_word),
+            )
+        if start_words or transitions:
+            await self._connection.commit()
+
     async def get_random_start_word(self, channel_id: int) -> str | None:
         """Return a random starting word for a channel.
 
@@ -192,18 +228,40 @@ class Database:
             row = await cursor.fetchone()
             return row[0] if row else None
 
-    async def add_image(self, channel_id: int, user_id: int, url: str) -> None:
-        """Store an image URL for a given channel.
+    async def add_images(self, images: list[tuple[int, int, str]]) -> None:
+        """Store image URLs for a given channel.
 
-        :param channel_id: The Discord channel ID.
-        :param user_id: The Discord user ID of the contributor.
-        :param url: The image URL to store.
+        :param images: A list of (channel_id, user_id, url) tuples.
         """
-        await self._connection.execute(
-            "INSERT INTO channel_images (channel_id, user_id, url) VALUES (?, ?, ?)",
-            (channel_id, user_id, url),
-        )
-        await self._connection.commit()
+        if images:
+            await self._connection.executemany(
+                "INSERT INTO channel_images"
+                " (channel_id, user_id, url)"
+                " VALUES (?, ?, ?)",
+                images,
+            )
+            await self._connection.commit()
+
+    async def remove_images(self, images: list[tuple[int, int, str]]) -> None:
+        """Remove one matching row per entry from the images table.
+
+        Each entry removes at most one duplicate row, preserving remaining
+        frequency weight.
+
+        :param images: A list of (channel_id, user_id, url) tuples.
+        """
+        for channel_id, user_id, url in images:
+            await self._connection.execute(
+                "DELETE FROM channel_images"
+                " WHERE rowid = ("
+                "   SELECT rowid FROM channel_images"
+                "   WHERE channel_id = ? AND user_id = ? AND url = ?"
+                "   LIMIT 1"
+                " )",
+                (channel_id, user_id, url),
+            )
+        if images:
+            await self._connection.commit()
 
     async def get_random_image(self, channel_id: int) -> str | None:
         """Return a random image URL for a given channel.

crabstero/listeners/message.py

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Handles created messages.
+"""Handles message creation and deletion events.
 
-Responds to mentions and ingests normal text messages as a Cog with an
-on_message listener.
+Responds to mentions, ingests normal text messages, and reverses
+ingestion when messages are deleted.
 """
 
 from typing import TYPE_CHECKING
@@ -24,7 +24,7 @@ import discord
 from discord.ext import commands
 
 from crabstero import metrics
-from crabstero.messages import ingest_message, reply_to_message
+from crabstero.messages import ingest_message, reply_to_message, uningest_message
 
 if TYPE_CHECKING:
     from crabstero.bot import Crabstero
@@ -63,7 +63,7 @@ class MessageCog(commands.Cog):
             if message.type == discord.MessageType.default and not isinstance(
                 channel, discord.Thread
             ):
-                await ingest_message(self.bot.db, message)
+                await ingest_message(self.bot.db, message, self.bot.ingest_cache)
             return
 
         if self.bot.user in message.mentions:
@@ -74,7 +74,28 @@ class MessageCog(commands.Cog):
         if message.type == discord.MessageType.default and not isinstance(
             channel, discord.Thread
         ):
-            await ingest_message(self.bot.db, message)
+            await ingest_message(self.bot.db, message, self.bot.ingest_cache)
+
+    @commands.Cog.listener()
+    async def on_raw_message_delete(
+        self, payload: discord.RawMessageDeleteEvent
+    ) -> None:
+        """Uningest a recently ingested message when it is deleted.
+
+        :param payload: The raw message delete event.
+        """
+        await uningest_message(self.bot.db, self.bot.ingest_cache, payload.message_id)
+
+    @commands.Cog.listener()
+    async def on_raw_bulk_message_delete(
+        self, payload: discord.RawBulkMessageDeleteEvent
+    ) -> None:
+        """Uningest recently ingested messages when they are bulk-deleted.
+
+        :param payload: The raw bulk message delete event.
+        """
+        for message_id in payload.message_ids:
+            await uningest_message(self.bot.db, self.bot.ingest_cache, message_id)
 
 
 async def setup(bot: Crabstero) -> None:

crabstero/markov.py

@@ -47,6 +47,39 @@ def is_complete_sentence(sentence: str) -> bool:
     return sentence[-1] in (DEFAULT_SENTENCE_END, ".", "!", "?")
 
 
+def _split_sentences(paragraph: str) -> list[str]:
+    """Normalize whitespace and split a paragraph into sentences.
+
+    :param paragraph: The raw paragraph text.
+    :return: A list of individual sentences.
+    """
+    if not is_complete_sentence(paragraph):
+        paragraph += DEFAULT_SENTENCE_END
+    normalized = re.sub(r" +", " ", paragraph.strip().replace("\n", " "))
+    return re.split(r"(?<=[.!?]) ", normalized)
+
+
+def _tokenize_sentence(
+    sentence: str,
+) -> tuple[list[str], list[tuple[str, str]]]:
+    """Tokenize a sentence into start words and transition pairs.
+
+    :param sentence: A single sentence.
+    :return: A tuple of (start_words, transitions) using bare word strings.
+    """
+    if not is_complete_sentence(sentence):
+        sentence += DEFAULT_SENTENCE_END
+    words = re.sub(r" +", " ", sentence.strip()).split(" ")
+
+    start_words: list[str] = []
+    transitions: list[tuple[str, str]] = []
+    for i in range(len(words) - 1):
+        if i == 0:
+            start_words.append(words[i])
+        transitions.append((words[i], words[i + 1]))
+    return start_words, transitions
+
+
 async def ingest(db: Database, channel_id: int, user_id: int, paragraph: str) -> None:
     """Ingest a paragraph into the Markov chain for a given channel.
 
@@ -58,15 +91,7 @@ async def ingest(db: Database, channel_id: int, user_id: int, paragraph: str) ->
     :param user_id: The Discord user ID of the contributor.
     :param paragraph: The paragraph of sentences to ingest.
     """
-    if not is_complete_sentence(paragraph):
-        paragraph += DEFAULT_SENTENCE_END
-
-    # Normalize whitespace, then split on sentence-ending punctuation
-    # followed by a space.
-    normalized = re.sub(r" +", " ", paragraph.strip().replace("\n", " "))
-    sentences = re.split(r"(?<=[.!?]) ", normalized)
-
-    for sentence in sentences:
+    for sentence in _split_sentences(paragraph):
         await _ingest_sentence(db, channel_id, user_id, sentence)
 
 
@@ -80,23 +105,43 @@ async def _ingest_sentence(
     :param user_id: The Discord user ID of the contributor.
     :param sentence: The sentence to ingest.
     """
-    if not is_complete_sentence(sentence):
-        sentence += DEFAULT_SENTENCE_END
-
-    # Normalize whitespace and split into individual words.
-    words = re.sub(r" +", " ", sentence.strip()).split(" ")
-
-    start_words: list[tuple[int, int, str]] = []
-    transitions: list[tuple[int, int, str, str]] = []
-
-    for i in range(len(words) - 1):
-        if i == 0:
-            start_words.append((channel_id, user_id, words[i]))
-        transitions.append((channel_id, user_id, words[i], words[i + 1]))
-
+    raw_starts, raw_transitions = _tokenize_sentence(sentence)
+    start_words = [(channel_id, user_id, w) for w in raw_starts]
+    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
     await db.add_markov_data(start_words, transitions)
 
 
+async def uningest(db: Database, channel_id: int, user_id: int, paragraph: str) -> None:
+    """Remove a paragraph's Markov data from the chain for a given channel.
+
+    Mirrors :func:`ingest` but deletes one matching row per entry instead of
+    inserting.
+
+    :param db: The database instance.
+    :param channel_id: The Discord channel ID.
+    :param user_id: The Discord user ID of the contributor.
+    :param paragraph: The paragraph of sentences to uningest.
+    """
+    for sentence in _split_sentences(paragraph):
+        await _uningest_sentence(db, channel_id, user_id, sentence)
+
+
+async def _uningest_sentence(
+    db: Database, channel_id: int, user_id: int, sentence: str
+) -> None:
+    """Remove a single sentence's Markov data from the chain.
+
+    :param db: The database instance.
+    :param channel_id: The Discord channel ID.
+    :param user_id: The Discord user ID of the contributor.
+    :param sentence: The sentence to uningest.
+    """
+    raw_starts, raw_transitions = _tokenize_sentence(sentence)
+    start_words = [(channel_id, user_id, w) for w in raw_starts]
+    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
+    await db.remove_markov_data(start_words, transitions)
+
+
 async def generate(
     db: Database, channel_id: int, soft_limit: int = 750, hard_limit: int = 1000
 ) -> str:

crabstero/messages.py

@@ -26,6 +26,7 @@ from typing import TYPE_CHECKING
 import discord
 
 from crabstero import flags, markov, metrics
+from crabstero.cache import CachedMessage, IngestCache
 from crabstero.flags import EntityType, Flag
 
 if TYPE_CHECKING:
@@ -115,11 +116,18 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
     metrics.REPLIES_SENT.inc()
 
 
-async def ingest_message(db: Database, message: discord.Message) -> None:
+async def ingest_message(
+    db: Database, message: discord.Message, cache: IngestCache | None = None
+) -> None:
     """Ingest a given message into its channel's Markov chain.
 
+    If a cache is provided, the ingested data is recorded so it can be
+    reversed by :func:`uningest_message` if the message is deleted shortly
+    after.
+
     :param db: The database instance.
     :param message: The message to ingest.
+    :param cache: Optional ingest cache for tracking recently ingested data.
     """
     guild = message.guild
     if guild is None:
@@ -141,31 +149,64 @@ async def ingest_message(db: Database, message: discord.Message) -> None:
     if not message.content and not message.embeds:
         return
 
+    embed_texts: list[str] = []
+    image_urls: list[str] = []
+
     with metrics.MESSAGE_INGESTION_DURATION.time():
         if message.content:
             await markov.ingest(db, channel_id, user_id, message.content)
 
         for embed in message.embeds:
-            await _ingest_embed(db, channel_id, user_id, embed)
+            if embed.title:
+                embed_texts.append(embed.title)
+                await markov.ingest(db, channel_id, user_id, embed.title)
+            if embed.description:
+                embed_texts.append(embed.description)
+                await markov.ingest(db, channel_id, user_id, embed.description)
+            if embed.image and embed.image.url:
+                image_urls.append(embed.image.url)
+
+        if image_urls:
+            await db.add_images([(channel_id, user_id, url) for url in image_urls])
 
     metrics.MESSAGES_INGESTED.inc()
 
+    if cache is not None:
+        cache.put(
+            message.id,
+            CachedMessage(
+                channel_id=channel_id,
+                user_id=user_id,
+                content=message.content or None,
+                embed_texts=embed_texts,
+                image_urls=image_urls,
+            ),
+        )
 
-async def _ingest_embed(
-    db: Database, channel_id: int, user_id: int, embed: discord.Embed
-) -> None:
-    """Ingest a given embed into a given channel's Markov chain.
+
+async def uningest_message(db: Database, cache: IngestCache, message_id: int) -> None:
+    """Reverse ingestion for a recently deleted message.
+
+    Looks up the message in the cache. If found, removes all Markov data
+    and images that were added during ingestion.
 
     :param db: The database instance.
-    :param channel_id: The ID of the channel to use for the Markov chain.
-    :param user_id: The Discord user ID of the contributor.
-    :param embed: The embed to ingest.
+    :param cache: The ingest cache.
+    :param message_id: The Discord message ID that was deleted.
     """
-    if embed.title:
-        await markov.ingest(db, channel_id, user_id, embed.title)
+    entry = cache.pop(message_id)
+    if entry is None:
+        return
 
-    if embed.description:
-        await markov.ingest(db, channel_id, user_id, embed.description)
+    if entry.content:
+        await markov.uningest(db, entry.channel_id, entry.user_id, entry.content)
 
-    if embed.image and embed.image.url:
-        await db.add_image(channel_id, user_id, embed.image.url)
+    for text in entry.embed_texts:
+        await markov.uningest(db, entry.channel_id, entry.user_id, text)
+
+    if entry.image_urls:
+        await db.remove_images(
+            [(entry.channel_id, entry.user_id, url) for url in entry.image_urls]
+        )
+
+    metrics.MESSAGES_UNINGESTED.inc()

crabstero/metrics.py

@@ -41,6 +41,10 @@ MESSAGES_INGESTED = Counter(
     "crabstero_messages_ingested_total",
     "Messages ingested into a Markov chain",
 )
+MESSAGES_UNINGESTED = Counter(
+    "crabstero_messages_uningested_total",
+    "Messages uningested from a Markov chain after deletion",
+)
 REPLIES_SENT = Counter(
     "crabstero_replies_sent_total",
     "Markov replies sent",