Rewrote Crabstero from Java to Python.
- Replaced Javacord with discord.py. - Replaced Redis backend with SQLite via aiosqlite. - Replaced Gradle build with pyproject.toml and uv. - Added setuptools-scm for automatic versioning from git tags. - Added argparse CLI with systemd credential support for the bot token. - Replaced per-channel ingestion tasks with a bounded queue and worker pool. - Removed Dockerfile and Gitea Actions workflow.
This commit is contained in:
@@ -0,0 +1,130 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Entry point for the Crabstero Discord bot.
|
||||
|
||||
Provides an argparse CLI with environment variable and systemd credential fallbacks for --token and --database-path.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
|
||||
from crabstero._version import version as crabstero_version
|
||||
from crabstero.bot import Crabstero
|
||||
|
||||
logger = logging.getLogger("crabstero")
|
||||
|
||||
|
||||
def _read_credential(name: str) -> str | None:
|
||||
"""
|
||||
Read a value from a systemd credential file.
|
||||
|
||||
Looks for a file named *name* inside the directory pointed to by the
|
||||
``CREDENTIALS_DIRECTORY`` environment variable (set automatically by
|
||||
systemd when ``LoadCredential=`` or ``SetCredential=`` is used).
|
||||
|
||||
:param name: Credential name to look up.
|
||||
:return: The credential value, or ``None`` if unavailable.
|
||||
"""
|
||||
credentials_dir = os.environ.get("CREDENTIALS_DIRECTORY")
|
||||
if credentials_dir is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return (
|
||||
open(os.path.join(credentials_dir, name)).read().strip() # noqa: SIM115
|
||||
)
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
"""
|
||||
Parses command-line arguments with environment variable fallbacks.
|
||||
|
||||
:param argv: Optional argument list (defaults to sys.argv[1:]).
|
||||
:return: Parsed arguments namespace.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="crabstero",
|
||||
description="Crabstero - the simple nonversation Discord bot.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--token",
|
||||
default=os.environ.get("TOKEN") or _read_credential("token"),
|
||||
help="Discord bot token (default: TOKEN environment variable or systemd credential 'token').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database-path",
|
||||
"--database",
|
||||
default=os.environ.get("DATABASE_PATH", "crabstero.db"),
|
||||
help='Path to the SQLite database file (default: DATABASE_PATH environment variable or "crabstero.db").',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ingestion-workers",
|
||||
type=int,
|
||||
default=int(os.environ.get("INGESTION_WORKERS", "4")),
|
||||
help="Number of concurrent ingestion workers (default: INGESTION_WORKERS environment variable or 4).",
|
||||
)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.token is None:
|
||||
parser.error(
|
||||
"a Discord bot token is required via --token, the TOKEN environment variable, or a systemd credential named 'token'"
|
||||
)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Main entry point. Parses arguments, configures logging, and starts the bot.
|
||||
"""
|
||||
args = _parse_args()
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s",
|
||||
)
|
||||
|
||||
logger.info("Starting Crabstero %s...", crabstero_version)
|
||||
|
||||
bot = Crabstero(
|
||||
token=args.token,
|
||||
database_path=args.database_path,
|
||||
ingestion_workers=args.ingestion_workers,
|
||||
)
|
||||
|
||||
async def _run() -> None:
|
||||
async with bot:
|
||||
await bot.start()
|
||||
|
||||
# Make SIGTERM behave like SIGINT so systemd stop triggers the same
|
||||
# clean shutdown path (context manager __aexit__ -> bot.close()).
|
||||
signal.signal(signal.SIGTERM, signal.default_int_handler)
|
||||
|
||||
try:
|
||||
asyncio.run(_run())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Interrupted.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,153 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
The simple nonversation Discord bot.
|
||||
|
||||
Provides the Crabstero subclass that owns the full bot lifecycle: database connection,
|
||||
cog loading, ingestion worker pool, and graceful shutdown.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
|
||||
from crabstero._version import version as crabstero_version
|
||||
from crabstero.database import Database
|
||||
from crabstero.tasks.ingestion import ingest_channel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Crabstero(commands.Bot):
|
||||
"""
|
||||
Central bot subclass that owns all lifecycle state.
|
||||
|
||||
The database is opened in setup_hook and closed in close(). Background
|
||||
ingestion is handled by a bounded queue and a fixed worker pool.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, token: str, database_path: str, ingestion_workers: int = 4
|
||||
) -> None:
|
||||
"""
|
||||
Configures intents, stores configuration, and prepares ingestion queue state.
|
||||
|
||||
:param token: The Discord bot token.
|
||||
:param database_path: The file path to the SQLite database.
|
||||
:param ingestion_workers: The number of concurrent ingestion workers.
|
||||
"""
|
||||
intents = discord.Intents.default()
|
||||
intents.guilds = True
|
||||
intents.guild_messages = True
|
||||
intents.message_content = True
|
||||
|
||||
super().__init__(
|
||||
command_prefix=[],
|
||||
intents=intents,
|
||||
max_messages=None, # Disables the message cache.
|
||||
)
|
||||
|
||||
self._token = token
|
||||
self._database_path = database_path
|
||||
self._ingestion_worker_count = ingestion_workers
|
||||
self._ingestion_queue: asyncio.Queue[
|
||||
discord.TextChannel | discord.VoiceChannel
|
||||
] = asyncio.Queue()
|
||||
self._ingestion_workers: list[asyncio.Task[None]] = []
|
||||
self.db: Database
|
||||
|
||||
self.http.user_agent = f"DiscordBot (https://git.logal.dev/LogalDeveloper/Crabstero, {crabstero_version})"
|
||||
|
||||
async def setup_hook(self) -> None:
|
||||
"""Opens the database, starts ingestion workers, loads all cogs, and syncs slash commands if changed."""
|
||||
self.db = await Database.connect(self._database_path)
|
||||
self._start_ingestion_workers()
|
||||
|
||||
from crabstero.listeners import interaction, message, server_events
|
||||
|
||||
await interaction.setup(self)
|
||||
await message.setup(self)
|
||||
await server_events.setup(self)
|
||||
|
||||
# Only sync slash commands if the registered commands differ from local definitions.
|
||||
local_commands = {
|
||||
cmd.name: cmd.description
|
||||
for cmd in self.tree.get_commands()
|
||||
if isinstance(cmd, (app_commands.Command, app_commands.Group))
|
||||
}
|
||||
try:
|
||||
remote_commands = {
|
||||
cmd.name: cmd.description for cmd in await self.tree.fetch_commands()
|
||||
}
|
||||
except discord.HTTPException:
|
||||
remote_commands = {}
|
||||
|
||||
if local_commands != remote_commands:
|
||||
logger.info("Slash command tree has changed, syncing with Discord.")
|
||||
await self.tree.sync()
|
||||
|
||||
async def on_ready(self) -> None:
|
||||
"""Logs that the bot has started successfully."""
|
||||
logger.info("Crabstero started!")
|
||||
|
||||
async def start(self, token: str = "", *, reconnect: bool = True) -> None:
|
||||
"""
|
||||
Starts the bot using the stored token by default.
|
||||
|
||||
:param token: Optional token override. Falls back to the stored token if empty.
|
||||
:param reconnect: Whether to automatically reconnect on disconnect.
|
||||
"""
|
||||
await super().start(token or self._token, reconnect=reconnect)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Cancels ingestion workers, closes the database, and then the bot connection."""
|
||||
if self.is_closed():
|
||||
return
|
||||
logger.info("Shutting down Crabstero...")
|
||||
for worker in self._ingestion_workers:
|
||||
worker.cancel()
|
||||
await asyncio.gather(*self._ingestion_workers, return_exceptions=True)
|
||||
self._ingestion_workers.clear()
|
||||
if hasattr(self, "db"):
|
||||
await self.db.close()
|
||||
await super().close()
|
||||
|
||||
def queue_channel_for_ingestion(
|
||||
self, channel: discord.TextChannel | discord.VoiceChannel
|
||||
) -> None:
|
||||
"""
|
||||
Enqueues a single channel for background message history ingestion.
|
||||
|
||||
:param channel: The channel to enqueue.
|
||||
"""
|
||||
self._ingestion_queue.put_nowait(channel)
|
||||
|
||||
def _start_ingestion_workers(self) -> None:
|
||||
"""Spawns the fixed pool of ingestion worker tasks."""
|
||||
for _ in range(self._ingestion_worker_count):
|
||||
task = asyncio.create_task(self._ingestion_worker())
|
||||
self._ingestion_workers.append(task)
|
||||
|
||||
async def _ingestion_worker(self) -> None:
|
||||
"""Loops forever pulling channels from the ingestion queue and ingesting them."""
|
||||
while True:
|
||||
channel = await self._ingestion_queue.get()
|
||||
try:
|
||||
await ingest_channel(channel, self.db)
|
||||
finally:
|
||||
self._ingestion_queue.task_done()
|
||||
@@ -0,0 +1,327 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Provides async SQLite database access for all of Crabstero's persistent storage needs.
|
||||
|
||||
Uses aiosqlite for native async access. All methods are async def.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Self
|
||||
|
||||
import aiosqlite
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AUTO_COMMIT_WRITE_THRESHOLD = 100 # Commit after this many DB write operations.
|
||||
AUTO_COMMIT_TIMEOUT_SECONDS = (
|
||||
60.0 # Commit after this many seconds since first uncommitted write.
|
||||
)
|
||||
|
||||
# SQL statements for creating the database schema.
|
||||
_SCHEMA = """
|
||||
-- Markov chain starting words.
|
||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||
CREATE TABLE IF NOT EXISTS markov_start_words (
|
||||
channel_id INTEGER NOT NULL,
|
||||
word TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_start_channel ON markov_start_words(channel_id, word);
|
||||
|
||||
-- Markov chain word transitions.
|
||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||
CREATE TABLE IF NOT EXISTS markov_transitions (
|
||||
channel_id INTEGER NOT NULL,
|
||||
word TEXT NOT NULL,
|
||||
next_word TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_transitions_channel_word ON markov_transitions(channel_id, word);
|
||||
|
||||
-- Image URLs per channel.
|
||||
CREATE TABLE IF NOT EXISTS channel_images (
|
||||
channel_id INTEGER NOT NULL,
|
||||
url TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_images_channel ON channel_images(channel_id);
|
||||
|
||||
-- Flags for channels, servers, and users.
|
||||
CREATE TABLE IF NOT EXISTS flags (
|
||||
entity_type TEXT NOT NULL,
|
||||
entity_id TEXT NOT NULL,
|
||||
flag_name TEXT NOT NULL,
|
||||
PRIMARY KEY (entity_type, entity_id, flag_name)
|
||||
);
|
||||
|
||||
-- Tracks which channels have been bulk-ingested.
|
||||
CREATE TABLE IF NOT EXISTS ingested_channels (
|
||||
channel_id INTEGER NOT NULL PRIMARY KEY
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
class Database:
|
||||
"""
|
||||
Manages all SQLite database operations for Crabstero.
|
||||
|
||||
Uses aiosqlite for native async access. A single connection is held open for the lifetime of
|
||||
the bot process with WAL mode enabled for concurrent read performance.
|
||||
"""
|
||||
|
||||
def __init__(self, connection: aiosqlite.Connection) -> None:
|
||||
"""
|
||||
Initializes the Database wrapper with an already-opened aiosqlite connection.
|
||||
|
||||
:param connection: An open aiosqlite connection.
|
||||
"""
|
||||
self._connection = connection
|
||||
self._pending_writes = 0
|
||||
self._flush_task: asyncio.Task[None] | None = None
|
||||
|
||||
@classmethod
|
||||
async def connect(cls, path: str) -> Self:
|
||||
"""
|
||||
Opens a new SQLite database at the given path, configures it for performance, and creates
|
||||
the schema if it does not already exist.
|
||||
|
||||
:param path: The file path to the SQLite database.
|
||||
:return: A new Database instance ready for use.
|
||||
"""
|
||||
connection = await aiosqlite.connect(path)
|
||||
|
||||
# Enable WAL mode for better concurrent read performance.
|
||||
await connection.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
# Set synchronous to NORMAL for a balance between safety and speed.
|
||||
await connection.execute("PRAGMA synchronous=NORMAL")
|
||||
|
||||
# Create the schema tables and indexes if they do not already exist.
|
||||
# executescript auto-commits, so no explicit commit is needed.
|
||||
await connection.executescript(_SCHEMA)
|
||||
|
||||
return cls(connection)
|
||||
|
||||
async def commit(self) -> None:
|
||||
"""Commits pending writes and resets the flush timer."""
|
||||
await self._connection.commit()
|
||||
self._pending_writes = 0
|
||||
if self._flush_task is not None:
|
||||
self._flush_task.cancel()
|
||||
self._flush_task = None
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Cancels the flush timer, commits any pending writes, then closes the database connection."""
|
||||
if self._flush_task is not None:
|
||||
self._flush_task.cancel()
|
||||
try:
|
||||
await self._flush_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._flush_task = None
|
||||
if self._pending_writes > 0:
|
||||
await self._connection.commit()
|
||||
self._pending_writes = 0
|
||||
await self._connection.close()
|
||||
|
||||
async def add_start_words_batch(self, rows: list[tuple[int, str]]) -> None:
|
||||
"""
|
||||
Inserts a batch of starting words into the markov_start_words table.
|
||||
|
||||
:param rows: A list of (channel_id, word) tuples to insert.
|
||||
"""
|
||||
await self._connection.executemany(
|
||||
"INSERT INTO markov_start_words (channel_id, word) VALUES (?, ?)", rows
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_start_word(self, channel_id: int) -> str | None:
|
||||
"""
|
||||
Returns a random starting word for a given channel, weighted by occurrence frequency.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: A random starting word, or None if no starting words exist for this channel.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT word FROM markov_start_words WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def add_transitions_batch(self, rows: list[tuple[int, str, str]]) -> None:
|
||||
"""
|
||||
Inserts a batch of word transitions into the markov_transitions table.
|
||||
|
||||
:param rows: A list of (channel_id, word, next_word) tuples to insert.
|
||||
"""
|
||||
await self._connection.executemany(
|
||||
"INSERT INTO markov_transitions (channel_id, word, next_word) VALUES (?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_next_word(self, channel_id: int, word: str) -> str | None:
|
||||
"""
|
||||
Returns a random next word for a given word in a given channel, weighted by occurrence
|
||||
frequency.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param word: The current word to find a transition for.
|
||||
:return: A random next word, or None if no transitions exist.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id, word),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def get_random_completing_next_word(
|
||||
self, channel_id: int, word: str
|
||||
) -> str | None:
|
||||
"""
|
||||
Returns a random next word that ends a sentence for a given word in a given channel,
|
||||
weighted by occurrence frequency.
|
||||
|
||||
A completing word is one whose last character is '.', '!', '?', or '§'.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param word: The current word to find a completing transition for.
|
||||
:return: A random completing next word, or None if no completing transitions exist.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? AND SUBSTR(next_word, -1, 1) IN ('.', '!', '?', '§') ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id, word),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def add_image(self, channel_id: int, url: str) -> None:
|
||||
"""
|
||||
Stores an image URL for a given channel.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param url: The image URL to store.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT INTO channel_images (channel_id, url) VALUES (?, ?)",
|
||||
(channel_id, url),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_image(self, channel_id: int) -> str | None:
|
||||
"""
|
||||
Returns a random image URL for a given channel.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: A random image URL, or None if no images exist for this channel.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT url FROM channel_images WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def set_flag(self, entity_type: str, entity_id: str, flag_name: str) -> None:
|
||||
"""
|
||||
Sets a flag on a given entity. If the flag is already set, this is a no-op.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to set.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT OR IGNORE INTO flags (entity_type, entity_id, flag_name) VALUES (?, ?, ?)",
|
||||
(entity_type, entity_id, flag_name),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def clear_flag(
|
||||
self, entity_type: str, entity_id: str, flag_name: str
|
||||
) -> None:
|
||||
"""
|
||||
Clears a flag on a given entity. If the flag is not set, this is a no-op.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to clear.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"DELETE FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
||||
(entity_type, entity_id, flag_name),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def is_flag_set(
|
||||
self, entity_type: str, entity_id: str, flag_name: str
|
||||
) -> bool:
|
||||
"""
|
||||
Checks whether a flag is set on a given entity.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to check.
|
||||
:return: True if the flag is set, False otherwise.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT 1 FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
||||
(entity_type, entity_id, flag_name),
|
||||
) as cursor:
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
async def is_channel_ingested(self, channel_id: int) -> bool:
|
||||
"""
|
||||
Checks whether a channel has already been bulk-ingested.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: True if the channel has been ingested, False otherwise.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT 1 FROM ingested_channels WHERE channel_id = ?",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
async def mark_channel_ingested(self, channel_id: int) -> None:
|
||||
"""
|
||||
Marks a channel as having been bulk-ingested.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)",
|
||||
(channel_id,),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def _maybe_commit(self) -> None:
|
||||
"""Tracks a pending write. Commits if the threshold is reached, otherwise starts a flush timer."""
|
||||
self._pending_writes += 1
|
||||
if self._pending_writes >= AUTO_COMMIT_WRITE_THRESHOLD:
|
||||
await self.commit()
|
||||
elif self._flush_task is None:
|
||||
self._flush_task = asyncio.create_task(self._flush_after_timeout())
|
||||
|
||||
async def _flush_after_timeout(self) -> None:
|
||||
"""Background task that commits after the timeout elapses."""
|
||||
try:
|
||||
await asyncio.sleep(AUTO_COMMIT_TIMEOUT_SECONDS)
|
||||
if self._pending_writes > 0:
|
||||
await self.commit()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Error in database flush timer.")
|
||||
@@ -0,0 +1,106 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Provides async convenience wrappers around the database flag methods.
|
||||
|
||||
Accepts discord.py objects or raw integer IDs and translates them into the
|
||||
entity_type/entity_id pairs used by the database layer.
|
||||
"""
|
||||
|
||||
import enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import discord
|
||||
|
||||
from crabstero.database import Database
|
||||
|
||||
|
||||
class Flag(enum.Enum):
|
||||
"""Enum of all supported flag names."""
|
||||
|
||||
NO_REPLY = "noReply"
|
||||
NO_INGEST = "noIngest"
|
||||
ALLOW_PINGS = "allowPings"
|
||||
|
||||
|
||||
class EntityType(enum.Enum):
|
||||
"""Enum of entity types that can have flags."""
|
||||
|
||||
CHANNEL = "channel"
|
||||
SERVER = "server"
|
||||
USER = "user"
|
||||
|
||||
|
||||
def _entity_id(entity: discord.abc.Snowflake | int) -> str:
|
||||
"""
|
||||
Extracts a string entity ID from a Discord object or raw integer ID.
|
||||
|
||||
:param entity: A Discord entity or raw integer ID.
|
||||
:return: The entity ID as a string.
|
||||
"""
|
||||
return str(entity) if isinstance(entity, int) else str(entity.id)
|
||||
|
||||
|
||||
async def set_flag(
|
||||
db: Database,
|
||||
entity: discord.abc.Snowflake | int,
|
||||
entity_type: EntityType,
|
||||
flag: Flag,
|
||||
) -> None:
|
||||
"""
|
||||
Sets a flag on a given entity.
|
||||
|
||||
:param db: The database instance.
|
||||
:param entity: The Discord entity or raw integer ID to set the flag on.
|
||||
:param entity_type: The type of the entity.
|
||||
:param flag: The flag to set.
|
||||
"""
|
||||
await db.set_flag(entity_type.value, _entity_id(entity), flag.value)
|
||||
|
||||
|
||||
async def clear_flag(
|
||||
db: Database,
|
||||
entity: discord.abc.Snowflake | int,
|
||||
entity_type: EntityType,
|
||||
flag: Flag,
|
||||
) -> None:
|
||||
"""
|
||||
Clears a flag on a given entity.
|
||||
|
||||
:param db: The database instance.
|
||||
:param entity: The Discord entity or raw integer ID to clear the flag on.
|
||||
:param entity_type: The type of the entity.
|
||||
:param flag: The flag to clear.
|
||||
"""
|
||||
await db.clear_flag(entity_type.value, _entity_id(entity), flag.value)
|
||||
|
||||
|
||||
async def is_flag_set(
|
||||
db: Database,
|
||||
entity: discord.abc.Snowflake | int,
|
||||
entity_type: EntityType,
|
||||
flag: Flag,
|
||||
) -> bool:
|
||||
"""
|
||||
Checks whether a flag is set on a given entity.
|
||||
|
||||
:param db: The database instance.
|
||||
:param entity: The Discord entity or raw integer ID to check.
|
||||
:param entity_type: The type of the entity.
|
||||
:param flag: The flag to check for.
|
||||
:return: True if the flag is set, False otherwise.
|
||||
"""
|
||||
return await db.is_flag_set(entity_type.value, _entity_id(entity), flag.value)
|
||||
@@ -0,0 +1,97 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Handles responding to interactions.
|
||||
|
||||
Provides the /pingme slash command as a Cog with an app command.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
|
||||
from crabstero import flags
|
||||
from crabstero.flags import EntityType, Flag
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.bot import Crabstero
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class InteractionCog(commands.Cog):
|
||||
"""
|
||||
Cog for handling slash command interactions.
|
||||
"""
|
||||
|
||||
def __init__(self, bot: Crabstero) -> None:
|
||||
"""
|
||||
Creates a new interaction handler cog.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
self.bot = bot
|
||||
|
||||
@app_commands.command(
|
||||
name="pingme",
|
||||
description="Opts into (or back out of) receiving pings for generated message which mention you.",
|
||||
)
|
||||
async def pingme(self, interaction: discord.Interaction) -> None:
|
||||
"""
|
||||
Toggles the allowPings flag for the user who ran the command.
|
||||
|
||||
:param interaction: The interaction event.
|
||||
"""
|
||||
# Wrap in try/except to always send a response even if the database fails.
|
||||
try:
|
||||
if await flags.is_flag_set(
|
||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
||||
):
|
||||
await flags.clear_flag(
|
||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
||||
)
|
||||
await interaction.response.send_message(
|
||||
"I will no longer ping you for messages which mention you. If you decide to opt back in, run `/pingme` any time.",
|
||||
ephemeral=True,
|
||||
)
|
||||
else:
|
||||
await flags.set_flag(
|
||||
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
|
||||
)
|
||||
await interaction.response.send_message(
|
||||
"I will now ping you for messages which mention you. If you change your mind, run `/pingme` any time.",
|
||||
ephemeral=True,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"An exception occurred while attempting to execute slash command responder for command 'pingme'."
|
||||
)
|
||||
await interaction.response.send_message(
|
||||
"An error occurred while executing your command. Please try again later.",
|
||||
ephemeral=True,
|
||||
)
|
||||
|
||||
|
||||
async def setup(bot: Crabstero) -> None:
|
||||
"""
|
||||
Adds the InteractionCog to the bot.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
await bot.add_cog(InteractionCog(bot))
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Handles created messages.
|
||||
|
||||
Responds to mentions and ingests normal text messages as a Cog with an on_message listener.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
|
||||
from crabstero.messages import ingest_message, reply_to_message
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.bot import Crabstero
|
||||
|
||||
|
||||
class MessageCog(commands.Cog):
|
||||
"""
|
||||
Cog for handling message creation events.
|
||||
"""
|
||||
|
||||
def __init__(self, bot: Crabstero) -> None:
|
||||
"""
|
||||
Creates a new message creation handler cog.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
self.bot = bot
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_message(self, message: discord.Message) -> None:
|
||||
"""
|
||||
Responds to mentions and ingests normal text messages.
|
||||
|
||||
:param message: The message event.
|
||||
"""
|
||||
channel = message.channel
|
||||
|
||||
if not isinstance(
|
||||
channel, (discord.TextChannel, discord.Thread, discord.VoiceChannel)
|
||||
):
|
||||
return
|
||||
|
||||
if message.author.bot or message.author == self.bot.user:
|
||||
return
|
||||
|
||||
if self.bot.user in message.mentions:
|
||||
await reply_to_message(self.bot.db, message)
|
||||
return
|
||||
|
||||
# Only ingest non-thread messages; threads share their parent channel's chain.
|
||||
if message.type == discord.MessageType.default and not isinstance(
|
||||
channel, discord.Thread
|
||||
):
|
||||
await ingest_message(self.bot.db, message)
|
||||
|
||||
|
||||
async def setup(bot: Crabstero) -> None:
|
||||
"""
|
||||
Adds the MessageCog to the bot.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
await bot.add_cog(MessageCog(bot))
|
||||
@@ -0,0 +1,130 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Handles server-level events which trigger channel history ingestion.
|
||||
|
||||
All of these events share the same outcome: queuing all textable channels for message history
|
||||
ingestion when something changes that may grant new read permissions.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
|
||||
from crabstero.tasks.ingestion import queue_channels_for_ingestion
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.bot import Crabstero
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ServerEventsCog(commands.Cog):
|
||||
"""
|
||||
Cog for handling server-level events that trigger channel history ingestion.
|
||||
"""
|
||||
|
||||
def __init__(self, bot: Crabstero) -> None:
|
||||
"""
|
||||
Creates a new server events handler cog.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
self.bot = bot
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_guild_join(self, guild: discord.Guild) -> None:
|
||||
"""
|
||||
Queues all text channels for message history ingestion when joining a new server.
|
||||
Also logs the join and sends an embed to the bot owner.
|
||||
|
||||
:param guild: The guild that was joined.
|
||||
"""
|
||||
logger.info('Joined new server! (Name: "%s" | ID: %s)', guild.name, guild.id)
|
||||
|
||||
embed = discord.Embed(
|
||||
title="Joined New Server", color=discord.Color.from_rgb(255, 165, 0)
|
||||
)
|
||||
embed.add_field(
|
||||
name=f"{guild.name} ({guild.id})",
|
||||
value=f"{guild.member_count} members",
|
||||
)
|
||||
|
||||
if guild.icon:
|
||||
embed.set_image(url=guild.icon.url)
|
||||
|
||||
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
|
||||
|
||||
app_info = await self.bot.application_info()
|
||||
if app_info.owner:
|
||||
try:
|
||||
await app_info.owner.send(embed=embed)
|
||||
except discord.HTTPException:
|
||||
pass
|
||||
|
||||
queue_channels_for_ingestion(guild, self.bot)
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_guild_available(self, guild: discord.Guild) -> None:
|
||||
"""
|
||||
Queues all textable channels for message history ingestion when a server becomes available.
|
||||
|
||||
:param guild: The guild that became available.
|
||||
"""
|
||||
queue_channels_for_ingestion(guild, self.bot)
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_guild_role_update(
|
||||
self, before: discord.Role, after: discord.Role
|
||||
) -> None:
|
||||
"""
|
||||
Queues all text channels for message history ingestion when role permissions change,
|
||||
but only if the bot is a member of the updated role.
|
||||
|
||||
:param before: The role before the update.
|
||||
:param after: The role after the update.
|
||||
"""
|
||||
if before.permissions == after.permissions:
|
||||
return
|
||||
|
||||
if after.guild.me in after.members:
|
||||
queue_channels_for_ingestion(after.guild, self.bot)
|
||||
|
||||
@commands.Cog.listener()
|
||||
async def on_guild_channel_update(
|
||||
self, before: discord.abc.GuildChannel, after: discord.abc.GuildChannel
|
||||
) -> None:
|
||||
"""
|
||||
Queues all text channels for message history ingestion when channel override permissions
|
||||
change.
|
||||
|
||||
:param before: The channel before the update.
|
||||
:param after: The channel after the update.
|
||||
"""
|
||||
if before.overwrites == after.overwrites:
|
||||
return
|
||||
|
||||
queue_channels_for_ingestion(after.guild, self.bot)
|
||||
|
||||
|
||||
async def setup(bot: Crabstero) -> None:
|
||||
"""
|
||||
Adds the ServerEventsCog to the bot.
|
||||
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
await bot.add_cog(ServerEventsCog(bot))
|
||||
@@ -0,0 +1,152 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Ingests sentences and generates new ones using a Markov chain backed by SQLite storage.
|
||||
|
||||
The Markov chain stores word transitions per channel, using duplicate rows to represent frequency
|
||||
weight. Random selection via ORDER BY RANDOM() LIMIT 1 naturally preserves this weighting.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.database import Database
|
||||
|
||||
# Section sign (§), used internally to mark the end of sentences that lack punctuation,
|
||||
# as is common on Discord.
|
||||
DEFAULT_SENTENCE_END = "\u00a7"
|
||||
|
||||
|
||||
def is_complete_sentence(sentence: str) -> bool:
|
||||
"""
|
||||
Checks whether a given sentence ends with a default sentence end character, a period, an
|
||||
exclamation mark, or a question mark.
|
||||
|
||||
:param sentence: The sentence to test.
|
||||
:return: True if the sentence ends with a valid terminator, False otherwise.
|
||||
"""
|
||||
if not sentence:
|
||||
return False
|
||||
|
||||
return sentence[-1] in (DEFAULT_SENTENCE_END, ".", "!", "?")
|
||||
|
||||
|
||||
async def ingest(db: Database, channel_id: int, paragraph: str) -> None:
|
||||
"""
|
||||
Ingests a string potentially containing multiple smaller sentences into the Markov chain
|
||||
for a given channel.
|
||||
|
||||
:param db: The database instance.
|
||||
:param channel_id: The Discord channel ID to associate with this data.
|
||||
: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:
|
||||
await _ingest_sentence(db, channel_id, sentence)
|
||||
|
||||
|
||||
async def _ingest_sentence(db: Database, channel_id: int, sentence: str) -> None:
|
||||
"""
|
||||
Ingests a string containing a single sentence into the Markov chain for a given channel.
|
||||
|
||||
:param db: The database instance.
|
||||
:param channel_id: The Discord channel ID to associate with this data.
|
||||
: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, str]] = []
|
||||
transitions: list[tuple[int, str, str]] = []
|
||||
|
||||
for i in range(len(words) - 1):
|
||||
if i == 0:
|
||||
start_words.append((channel_id, words[i]))
|
||||
transitions.append((channel_id, words[i], words[i + 1]))
|
||||
|
||||
if start_words:
|
||||
await db.add_start_words_batch(start_words)
|
||||
if transitions:
|
||||
await db.add_transitions_batch(transitions)
|
||||
|
||||
|
||||
async def generate(
|
||||
db: Database, channel_id: int, soft_limit: int = 750, hard_limit: int = 1000
|
||||
) -> str:
|
||||
"""
|
||||
Generates a new sentence using words learned from previously ingested sentences for a given
|
||||
channel.
|
||||
|
||||
:param db: The database instance.
|
||||
:param channel_id: The Discord channel ID to generate from.
|
||||
:param soft_limit: The amount of characters to try and limit sentence length around.
|
||||
:param hard_limit: The amount of characters to cut off the sentence at if it gets too long.
|
||||
:return: A new generated sentence.
|
||||
"""
|
||||
word = await db.get_random_start_word(channel_id)
|
||||
|
||||
# Seed the chain with a fallback sentence if the channel has no data yet.
|
||||
if word is None:
|
||||
await _ingest_sentence(db, channel_id, "Hello world!")
|
||||
word = await db.get_random_start_word(channel_id)
|
||||
if word is None:
|
||||
return ""
|
||||
|
||||
parts: list[str] = []
|
||||
parts.append(word)
|
||||
current_length = len(word)
|
||||
|
||||
# The loop is skipped if the starting word already ends a sentence (e.g. "Yes.").
|
||||
while not is_complete_sentence(word):
|
||||
# Past the soft limit, prefer a sentence-ending word to wrap up.
|
||||
if current_length >= soft_limit:
|
||||
next_word = await db.get_random_completing_next_word(channel_id, word)
|
||||
if next_word is None:
|
||||
next_word = await db.get_random_next_word(channel_id, word)
|
||||
else:
|
||||
next_word = await db.get_random_next_word(channel_id, word)
|
||||
|
||||
if next_word is None:
|
||||
break
|
||||
|
||||
word = next_word
|
||||
|
||||
parts.append(word)
|
||||
current_length += 1 + len(word) # +1 for the joining space.
|
||||
|
||||
if current_length >= hard_limit:
|
||||
result = " ".join(parts)[:hard_limit]
|
||||
# Strip the internal sentence-end marker if it ended up at the boundary.
|
||||
if result and result[-1] == DEFAULT_SENTENCE_END:
|
||||
return result[:-1]
|
||||
return result
|
||||
|
||||
result = " ".join(parts)
|
||||
|
||||
# Strip the internal sentence-end marker so it never appears in output.
|
||||
if result and result[-1] == DEFAULT_SENTENCE_END:
|
||||
return result[:-1]
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,157 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Assists with generating Discord messages in response to other users and ingesting raw messages.
|
||||
|
||||
Orchestrates Markov chain generation and ingestion in the context of Discord messages, handling
|
||||
reply logic, embed generation, mention filtering, and message ingestion with flag checks.
|
||||
"""
|
||||
|
||||
import re
|
||||
import secrets
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import discord
|
||||
|
||||
from crabstero import flags, markov
|
||||
from crabstero.flags import EntityType, Flag
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.database import Database
|
||||
|
||||
# Copied from discordjs/discord-api-types:
|
||||
# https://github.com/discordjs/discord-api-types/blob/7fe434114e91c80ed79f0204ae6c73047672d55d/globals.ts#L30
|
||||
MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
|
||||
|
||||
|
||||
async def reply_to_message(db: Database, message: discord.Message) -> None:
|
||||
"""
|
||||
Sends a new message in Discord in response to a given message.
|
||||
|
||||
:param db: The database instance.
|
||||
:param message: The message prompting the response.
|
||||
"""
|
||||
channel = message.channel
|
||||
guild = message.guild
|
||||
if guild is None:
|
||||
return
|
||||
|
||||
if not channel.permissions_for(guild.me).send_messages:
|
||||
return
|
||||
|
||||
if (
|
||||
await flags.is_flag_set(db, channel, EntityType.CHANNEL, Flag.NO_REPLY)
|
||||
or await flags.is_flag_set(db, guild, EntityType.SERVER, Flag.NO_REPLY)
|
||||
or await flags.is_flag_set(db, message.author, EntityType.USER, Flag.NO_REPLY)
|
||||
):
|
||||
return
|
||||
|
||||
# Threads share their parent channel's Markov chain.
|
||||
if isinstance(channel, discord.Thread):
|
||||
channel_id = channel.parent_id
|
||||
else:
|
||||
channel_id = channel.id
|
||||
|
||||
body = await markov.generate(db, channel_id, 750, 1000)
|
||||
|
||||
embed = None
|
||||
|
||||
# 5% chance to include an embed, if the bot has permission.
|
||||
if secrets.randbelow(100) >= 95 and channel.permissions_for(guild.me).embed_links:
|
||||
embed = discord.Embed(
|
||||
title=await markov.generate(db, channel_id, 200, 300),
|
||||
description=await markov.generate(db, channel_id, 300, 500),
|
||||
)
|
||||
|
||||
random_image = await db.get_random_image(channel_id)
|
||||
if random_image is not None:
|
||||
embed.set_image(url=random_image)
|
||||
|
||||
# Suppress all mentions by default; only ping users who opted in.
|
||||
allowed_user_ids: list[int] = []
|
||||
|
||||
for match in MENTION_PATTERN.finditer(body):
|
||||
user_id = int(match.group("id"))
|
||||
|
||||
if await flags.is_flag_set(db, user_id, EntityType.USER, Flag.ALLOW_PINGS):
|
||||
allowed_user_ids.append(user_id)
|
||||
|
||||
allowed_mentions = discord.AllowedMentions(
|
||||
everyone=False,
|
||||
roles=False,
|
||||
users=[discord.Object(id=uid) for uid in allowed_user_ids],
|
||||
)
|
||||
|
||||
if embed is not None:
|
||||
await message.reply(
|
||||
content=body,
|
||||
embed=embed,
|
||||
allowed_mentions=allowed_mentions,
|
||||
mention_author=False,
|
||||
)
|
||||
else:
|
||||
await message.reply(
|
||||
content=body,
|
||||
allowed_mentions=allowed_mentions,
|
||||
mention_author=False,
|
||||
)
|
||||
|
||||
|
||||
async def ingest_message(db: Database, message: discord.Message) -> None:
|
||||
"""
|
||||
Ingests a given message into its channel's Markov chain.
|
||||
|
||||
:param db: The database instance.
|
||||
:param message: The message to ingest.
|
||||
"""
|
||||
guild = message.guild
|
||||
if guild is None:
|
||||
return
|
||||
|
||||
if message.author.bot or guild.me in message.mentions:
|
||||
return
|
||||
|
||||
if (
|
||||
await flags.is_flag_set(db, message.channel, EntityType.CHANNEL, Flag.NO_INGEST)
|
||||
or await flags.is_flag_set(db, guild, EntityType.SERVER, Flag.NO_INGEST)
|
||||
or await flags.is_flag_set(db, message.author, EntityType.USER, Flag.NO_INGEST)
|
||||
):
|
||||
return
|
||||
|
||||
channel_id = message.channel.id
|
||||
|
||||
if message.content:
|
||||
await markov.ingest(db, channel_id, message.content)
|
||||
|
||||
for embed in message.embeds:
|
||||
await _ingest_embed(db, channel_id, embed)
|
||||
|
||||
|
||||
async def _ingest_embed(db: Database, channel_id: int, embed: discord.Embed) -> None:
|
||||
"""
|
||||
Ingests a given embed into a given channel's Markov chain.
|
||||
|
||||
:param db: The database instance.
|
||||
:param channel_id: The ID of the channel to use for the Markov chain.
|
||||
:param embed: The embed to ingest.
|
||||
"""
|
||||
if embed.title:
|
||||
await markov.ingest(db, channel_id, embed.title)
|
||||
|
||||
if embed.description:
|
||||
await markov.ingest(db, channel_id, embed.description)
|
||||
|
||||
if embed.image and embed.image.url:
|
||||
await db.add_image(channel_id, embed.image.url)
|
||||
@@ -0,0 +1,92 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Bulk-ingests the message history of channels.
|
||||
|
||||
Channels are enqueued for processing by the bot's fixed ingestion worker pool.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import discord
|
||||
|
||||
from crabstero.messages import ingest_message
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crabstero.bot import Crabstero
|
||||
from crabstero.database import Database
|
||||
|
||||
MAXIMUM_MESSAGES_PER_CHANNEL = (
|
||||
50000 # The maximum amount of historical messages to ingest per channel.
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
|
||||
"""
|
||||
Enqueues every textable channel in a guild for message history ingestion.
|
||||
|
||||
:param guild: The Discord guild whose channels should be ingested.
|
||||
:param bot: The bot instance.
|
||||
"""
|
||||
for channel in guild.channels:
|
||||
if isinstance(channel, (discord.TextChannel, discord.VoiceChannel)):
|
||||
bot.queue_channel_for_ingestion(channel)
|
||||
|
||||
|
||||
async def ingest_channel(
|
||||
channel: discord.TextChannel | discord.VoiceChannel,
|
||||
db: Database,
|
||||
) -> None:
|
||||
"""
|
||||
Bulk-ingests the message history of a given channel. The task will be ended early if
|
||||
permissions do not allow ingesting this channel or if it has already been ingested in the past.
|
||||
|
||||
:param channel: The channel to ingest.
|
||||
:param db: The database instance.
|
||||
"""
|
||||
try:
|
||||
if not channel.permissions_for(channel.guild.me).read_message_history:
|
||||
logger.warning(
|
||||
"[%s] Unable to ingest textable channel history due to lacking permissions. Ignoring.",
|
||||
channel.id,
|
||||
)
|
||||
return
|
||||
|
||||
if await db.is_channel_ingested(channel.id):
|
||||
return
|
||||
|
||||
await db.mark_channel_ingested(channel.id)
|
||||
|
||||
logger.info("[%s] Starting ingestion of textable channel history.", channel.id)
|
||||
|
||||
count = 0
|
||||
async for message in channel.history(limit=MAXIMUM_MESSAGES_PER_CHANNEL):
|
||||
count += 1
|
||||
await ingest_message(db, message)
|
||||
|
||||
logger.info(
|
||||
"[%s] Ingestion of textable channel history complete. %d messages ingested.",
|
||||
channel.id,
|
||||
count,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"[%s] An error occurred while ingesting textable channel history!",
|
||||
channel.id,
|
||||
)
|
||||
Reference in New Issue
Block a user