nonebot-plugin-group-chat-memory 0.1.0__py3-none-any.whl

nonebot_plugin_group_chat_memory/replies.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+from collections import deque
+from collections.abc import Awaitable, Callable
+
+from .models import Message
+from .store import MemoryStore
+
+logger = logging.getLogger(__name__)
+
+ReplyCallable = Callable[[int, Message], Awaitable[None]]
+RandomValue = Callable[[], float]
+
+WORKER_NAME_PREFIX = "group-chat-memory-reply"
+
+
+class _GroupState:
+    """单群的待回复槽位：点名请求 FIFO，随机请求只保留最新一条。"""
+
+    __slots__ = ("explicit", "latest_random", "task", "wake")
+
+    def __init__(self) -> None:
+        self.explicit: deque[Message] = deque()
+        self.latest_random: Message | None = None
+        self.wake = asyncio.Event()
+        self.task: asyncio.Task[None] | None = None
+
+    def take_next(self) -> Message | None:
+        """取出下一条要回复的消息：先点名队列，再随机槽位。"""
+        if self.explicit:
+            return self.explicit.popleft()
+        message, self.latest_random = self.latest_random, None
+        return message
+
+    def has_pending(self) -> bool:
+        return bool(self.explicit) or self.latest_random is not None
+
+    def discard_random(self) -> bool:
+        """丢弃等待中的随机请求；点名请求保留给 worker 排空。"""
+        had_random = self.latest_random is not None
+        self.latest_random = None
+        return had_random
+
+
+class ReplyCoordinator:
+    """按群串行、跨群并行的回复调度器。
+
+    每个群一个 worker：`explicit` 请求进入 FIFO，保证逐条且全部被回复；
+    非点名（随机）请求只保留最新一条，等待中的旧请求被直接替换，避免突发消息
+    造成回复风暴。不同群各自独立执行，一个群的慢回复不会阻塞其他群。
+
+    提交时只保存消息本身，不复制上下文；`reply` 回调在实际执行时才构建上下文，
+    因此回复使用的永远是执行时刻的群聊状态。单个请求的异常被隔离并记录日志，
+    不会中断该群后续请求；`reply` 抛出的 `CancelledError` 会结束 worker，并由
+    监督回调为剩余请求重新拉起 worker。`submit` 必须在事件循环中调用；关闭后
+    不再接受请求，已接受的点名请求仍会全部执行完毕。
+    """
+
+    def __init__(
+        self,
+        store: MemoryStore,
+        reply: ReplyCallable,
+        *,
+        random_value: RandomValue = random.random,
+    ) -> None:
+        self._store = store
+        self._reply = reply
+        self._random_value = random_value
+        self._groups: dict[int, _GroupState] = {}
+        self._closing = False
+
+    async def should_random_reply(self, group_id: int) -> bool:
+        """按当前群概率判断是否随机回复：0% 永不、100% 必定。
+
+        概率每次从存储读取，管理命令修改后立即生效。0 与 100 的边界不消耗随机值。
+        """
+        probability = await self._store.get_probability(group_id)
+        if probability <= 0.0:
+            return False
+        if probability >= 100.0:
+            return True
+        return self._random_value() < probability / 100.0
+
+    def submit(self, group_id: int, message: Message, explicit: bool) -> None:
+        """提交一条待回复消息；`explicit` 表示点名或回复触发的必须回复请求。"""
+        if self._closing:
+            logger.debug("回复协调器已关闭，忽略群 %s 的回复请求。", group_id)
+            return
+        state = self._groups.get(group_id)
+        if state is None:
+            state = _GroupState()
+            self._groups[group_id] = state
+        if explicit:
+            state.explicit.append(message)
+        else:
+            state.latest_random = message
+        state.wake.set()
+        self._ensure_worker(group_id, state)
+
+    async def close(self) -> None:
+        """停止接受新请求，并在所有已接受的点名请求回复完毕后关闭 worker。
+
+        已接受的 explicit（点名/回复）请求全部执行完毕才返回；等待中的随机请求
+        被丢弃（只记录群号，不含正文），在途回复也会执行完毕，因此关闭后不会
+        遗留任何 worker 或群状态。已结束 worker 的 done callback 可能仍在队列里
+        并因排队的点名请求拉起新 worker，所以这里循环等待到完全静默。
+        可重复调用；调用方取消 `close()` 不会级联取消 worker，取消异常照常向上
+        抛出，可再次调用完成收尾。
+        """
+        first = not self._closing
+        self._closing = True
+        for group_id, state in list(self._groups.items()):
+            if first and state.discard_random():
+                logger.warning("回复协调器关闭，群 %s 的等待中的随机回复请求被丢弃。", group_id)
+            state.wake.set()
+        while True:
+            # 先让已结束 worker 的 done callback 执行: 它可能为排队的点名请求拉起新 worker.
+            await asyncio.sleep(0)
+            tasks = [
+                state.task
+                for state in self._groups.values()
+                if state.task is not None and not state.task.done()
+            ]
+            if tasks:
+                await asyncio.gather(
+                    *(asyncio.shield(task) for task in tasks), return_exceptions=True
+                )
+                continue
+            pending = [
+                (group_id, state) for group_id, state in self._groups.items() if state.explicit
+            ]
+            if not pending:
+                break
+            # 兜底: worker 的 done callback 尚未执行时, 由 close 自己续跑点名请求.
+            for group_id, state in pending:
+                state.wake.set()
+                self._ensure_worker(group_id, state)
+        self._groups.clear()
+
+    def _ensure_worker(self, group_id: int, state: _GroupState) -> None:
+        task = state.task
+        if task is not None and not task.done():
+            return
+        task = asyncio.create_task(
+            self._run_group(group_id, state), name=f"{WORKER_NAME_PREFIX}-{group_id}"
+        )
+        state.task = task
+        task.add_done_callback(
+            lambda done, gid=group_id, st=state: self._on_worker_done(gid, st, done)
+        )
+
+    async def _run_group(self, group_id: int, state: _GroupState) -> None:
+        while True:
+            message = state.take_next()
+            if message is None:
+                if self._closing:
+                    return
+                state.wake.clear()
+                if state.has_pending():
+                    continue
+                await state.wake.wait()
+                continue
+            try:
+                await self._reply(group_id, message)
+            except Exception:
+                logger.exception("群 %s 的回复失败，已跳过该请求。", group_id)
+
+    def _on_worker_done(self, group_id: int, state: _GroupState, task: asyncio.Task[None]) -> None:
+        if task.cancelled():
+            if not self._closing:
+                logger.warning("群 %s 的回复 worker 被取消。", group_id)
+        else:
+            error = task.exception()
+            if error is not None and not self._closing:
+                logger.error("群 %s 的回复 worker 意外退出：%r", group_id, error, exc_info=error)
+        if state.task is task:
+            state.task = None
+        if state.task is not None and not state.task.done():
+            return
+        if self._needs_recovery(state):
+            state.wake.set()
+            self._ensure_worker(group_id, state)
+        elif self._closing:
+            self._groups.pop(group_id, None)
+
+    def _needs_recovery(self, state: _GroupState) -> bool:
+        """是否还有必须继续处理的请求: 关闭后仍要排空已接受的点名请求。"""
+        if state.explicit:
+            return True
+        return not self._closing and state.latest_random is not None

nonebot_plugin_group_chat_memory/store.py

@@ -0,0 +1,425 @@
+from __future__ import annotations
+
+import asyncio
+import math
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
+from datetime import datetime, timezone
+
+import aiosqlite
+import sqlite_vec
+
+from .config import RUNTIME_CONFIG_DEFAULTS, Config
+from .models import Chunk, ChunkStatus, Message
+
+DIMENSIONS_META_KEY = "embedding_dimensions"
+
+_SCHEMA_STATEMENTS = (
+    """
+    CREATE TABLE IF NOT EXISTS meta (
+        key TEXT PRIMARY KEY,
+        value TEXT NOT NULL
+    )
+    """,
+    """
+    CREATE TABLE IF NOT EXISTS runtime_config (
+        key TEXT PRIMARY KEY,
+        value TEXT NOT NULL
+    )
+    """,
+    """
+    CREATE TABLE IF NOT EXISTS groups (
+        group_id INTEGER PRIMARY KEY,
+        probability REAL NOT NULL DEFAULT 0
+            CHECK (probability >= 0 AND probability <= 100),
+        updated_at TEXT NOT NULL
+    )
+    """,
+    """
+    CREATE TABLE IF NOT EXISTS chunks (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        group_id INTEGER NOT NULL REFERENCES groups(group_id) ON DELETE CASCADE,
+        text TEXT NOT NULL,
+        status TEXT NOT NULL CHECK (status IN ('pending', 'ready')),
+        created_at TEXT NOT NULL
+    )
+    """,
+    "CREATE INDEX IF NOT EXISTS idx_chunks_group ON chunks(group_id)",
+    "CREATE INDEX IF NOT EXISTS idx_chunks_status ON chunks(status, id)",
+    """
+    CREATE TABLE IF NOT EXISTS messages (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        group_id INTEGER NOT NULL REFERENCES groups(group_id) ON DELETE CASCADE,
+        user_id INTEGER NOT NULL,
+        display_name TEXT NOT NULL,
+        text TEXT NOT NULL,
+        created_at TEXT NOT NULL,
+        is_bot INTEGER NOT NULL DEFAULT 0 CHECK (is_bot IN (0, 1)),
+        chunk_id INTEGER REFERENCES chunks(id) ON DELETE SET NULL
+    )
+    """,
+    """
+    CREATE INDEX IF NOT EXISTS idx_messages_group_time
+        ON messages(group_id, created_at, id)
+    """,
+    """
+    CREATE INDEX IF NOT EXISTS idx_messages_unchunked
+        ON messages(group_id, id) WHERE chunk_id IS NULL
+    """,
+)
+
+_VECTOR_TABLE_STATEMENT = """
+    CREATE VIRTUAL TABLE IF NOT EXISTS memory_vectors USING vec0(
+        chunk_id INTEGER PRIMARY KEY,
+        group_id INTEGER PARTITION KEY,
+        embedding FLOAT[{dimensions}] distance_metric=cosine
+    )
+"""
+
+
+class MemoryStore:
+    """按群隔离的 SQLite 与 sqlite-vec 持久化存储。"""
+
+    def __init__(self, config: Config) -> None:
+        self._config = config
+        self._dimensions = config.group_chat_memory_embedding_dimensions
+        self._default_probability = config.group_chat_memory_default_probability
+        self._db: aiosqlite.Connection | None = None
+        self._lock = asyncio.Lock()
+
+    async def open(self) -> None:
+        """打开数据库、加载 sqlite-vec 并创建缺失的表结构。"""
+        async with self._lock:
+            if self._db is not None:
+                return
+
+            path = self._config.database_path
+            path.parent.mkdir(parents=True, exist_ok=True)
+            db = await aiosqlite.connect(path, isolation_level=None)
+            try:
+                db.row_factory = aiosqlite.Row
+                await db.enable_load_extension(True)
+                await db.load_extension(sqlite_vec.loadable_path())
+                await db.enable_load_extension(False)
+                await db.execute("PRAGMA journal_mode = WAL")
+                await db.execute("PRAGMA foreign_keys = ON")
+                await db.execute("PRAGMA busy_timeout = 5000")
+                await db.execute("PRAGMA synchronous = NORMAL")
+                self._db = db
+                async with self._transaction_on(db) as tx:
+                    for statement in _SCHEMA_STATEMENTS:
+                        await tx.execute(statement)
+                    await self._ensure_dimensions(tx)
+                    await tx.execute(_VECTOR_TABLE_STATEMENT.format(dimensions=self._dimensions))
+            except BaseException:
+                self._db = None
+                await db.close()
+                raise
+
+    async def close(self) -> None:
+        """关闭数据库连接，可重复调用。"""
+        async with self._lock:
+            db, self._db = self._db, None
+            if db is not None:
+                await db.close()
+
+    async def add_message(self, message: Message) -> int:
+        """持久化一条群消息，返回数据库内的消息 ID。"""
+        async with self._transaction() as db:
+            await self._ensure_group(db, message.group_id)
+            cursor = await db.execute(
+                "INSERT INTO messages"
+                "(group_id, user_id, display_name, text, created_at, is_bot) "
+                "VALUES (?, ?, ?, ?, ?, ?)",
+                (
+                    message.group_id,
+                    message.user_id,
+                    message.display_name,
+                    message.text,
+                    _format_timestamp(message.created_at),
+                    int(message.is_bot),
+                ),
+            )
+            message_id = cursor.lastrowid
+            await cursor.close()
+        if message_id is None:
+            raise RuntimeError("消息写入失败：数据库未返回消息 ID。")
+        return int(message_id)
+
+    async def recent_messages(self, group_id: int, limit: int) -> list[Message]:
+        """返回当前群最近 limit 条消息，按时间从旧到新排列。"""
+        if limit < 1:
+            return []
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute(
+                "SELECT group_id, user_id, display_name, text, created_at, is_bot "
+                "FROM messages WHERE group_id = ? "
+                "ORDER BY created_at DESC, id DESC LIMIT ?",
+                (group_id, limit),
+            )
+            rows = await cursor.fetchall()
+            await cursor.close()
+        return [_row_to_message(row) for row in reversed(rows)]
+
+    async def get_probability(self, group_id: int) -> float:
+        """返回当前群的随机回复概率；未配置的群使用插件默认值。"""
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute(
+                "SELECT probability FROM groups WHERE group_id = ?", (group_id,)
+            )
+            row = await cursor.fetchone()
+            await cursor.close()
+        if row is None:
+            return self._default_probability
+        return float(row["probability"])
+
+    async def set_probability(self, group_id: int, value: float) -> None:
+        """设置当前群的随机回复概率，取值必须在 0 到 100 之间。"""
+        probability = float(value)
+        if not math.isfinite(probability) or not 0.0 <= probability <= 100.0:
+            raise ValueError("随机回复概率必须是 0 到 100 之间的数值。")
+        async with self._transaction() as db:
+            await db.execute(
+                "INSERT INTO groups(group_id, probability, updated_at) VALUES (?, ?, ?) "
+                "ON CONFLICT(group_id) DO UPDATE SET "
+                "probability = excluded.probability, updated_at = excluded.updated_at",
+                (group_id, probability, _format_timestamp(_utc_now())),
+            )
+
+    async def get_runtime_config(self) -> dict[str, str]:
+        """读取全 Bot 共用的模型运行时配置。"""
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute("SELECT key, value FROM runtime_config")
+            rows = await cursor.fetchall()
+            await cursor.close()
+        return {row["key"]: row["value"] for row in rows}
+
+    async def set_runtime_config(self, key: str, value: str) -> None:
+        """持久化一项白名单内的全局模型配置。"""
+        if key not in RUNTIME_CONFIG_DEFAULTS:
+            raise ValueError(f"不支持的运行时配置项：{key}")
+        async with self._transaction() as db:
+            await db.execute(
+                "INSERT INTO runtime_config(key, value) VALUES (?, ?) "
+                "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+                (key, value),
+            )
+
+    async def create_ready_chunk(self, group_id: int, size: int) -> Chunk | None:
+        """用当前群最早的 size 条未分块消息组成片段，并原子关联这些消息。"""
+        if size < 1:
+            raise ValueError("分块大小必须至少为 1。")
+        async with self._transaction() as db:
+            await self._ensure_group(db, group_id)
+            cursor = await db.execute(
+                "SELECT id, user_id, display_name, text, created_at, is_bot "
+                "FROM messages WHERE group_id = ? AND chunk_id IS NULL "
+                "ORDER BY created_at, id LIMIT ?",
+                (group_id, size),
+            )
+            rows = await cursor.fetchall()
+            await cursor.close()
+            if len(rows) < size:
+                return None
+
+            text = _format_chunk_text(group_id, rows)
+            cursor = await db.execute(
+                "INSERT INTO chunks(group_id, text, status, created_at) VALUES (?, ?, ?, ?)",
+                (
+                    group_id,
+                    text,
+                    ChunkStatus.PENDING.value,
+                    _format_timestamp(_utc_now()),
+                ),
+            )
+            chunk_id = cursor.lastrowid
+            await cursor.close()
+            if chunk_id is None:
+                raise RuntimeError("片段写入失败：数据库未返回片段 ID。")
+            placeholders = ", ".join("?" for _ in rows)
+            await db.execute(
+                f"UPDATE messages SET chunk_id = ? WHERE id IN ({placeholders})",
+                (chunk_id, *(row["id"] for row in rows)),
+            )
+        return Chunk(id=chunk_id, group_id=group_id, text=text)
+
+    async def pending_chunks(self) -> list[Chunk]:
+        """返回所有群中尚未写入向量的片段。"""
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute(
+                "SELECT id, group_id, text, status FROM chunks WHERE status = ? ORDER BY id",
+                (ChunkStatus.PENDING.value,),
+            )
+            rows = await cursor.fetchall()
+            await cursor.close()
+        return [_row_to_chunk(row) for row in rows]
+
+    async def pending_groups(self, size: int) -> list[int]:
+        """返回未分块消息数量至少为 size 的群，按群号升序。"""
+        if size < 1:
+            raise ValueError("分块大小必须至少为 1。")
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute(
+                "SELECT group_id FROM messages WHERE chunk_id IS NULL "
+                "GROUP BY group_id HAVING count(*) >= ? ORDER BY group_id",
+                (size,),
+            )
+            rows = await cursor.fetchall()
+            await cursor.close()
+        return [int(row["group_id"]) for row in rows]
+
+    async def complete_chunk(self, chunk_id: int, vector: list[float]) -> None:
+        """写入片段向量并标记为 ready；片段已被清空时静默忽略。"""
+        embedding = _validated_vector(vector, self._dimensions)
+        async with self._transaction() as db:
+            cursor = await db.execute("SELECT group_id FROM chunks WHERE id = ?", (chunk_id,))
+            row = await cursor.fetchone()
+            await cursor.close()
+            if row is None:
+                return
+            await db.execute("DELETE FROM memory_vectors WHERE chunk_id = ?", (chunk_id,))
+            await db.execute(
+                "INSERT INTO memory_vectors(chunk_id, group_id, embedding) VALUES (?, ?, ?)",
+                (chunk_id, row["group_id"], sqlite_vec.serialize_float32(embedding)),
+            )
+            await db.execute(
+                "UPDATE chunks SET status = ? WHERE id = ?",
+                (ChunkStatus.READY.value, chunk_id),
+            )
+
+    async def search(self, group_id: int, vector: list[float], limit: int) -> list[Chunk]:
+        """在当前群已完成的片段中检索最接近的 limit 个片段。"""
+        embedding = _validated_vector(vector, self._dimensions)
+        if limit < 1:
+            return []
+        async with self._lock:
+            db = self._require_db()
+            cursor = await db.execute(
+                "SELECT c.id, c.group_id, c.text, c.status FROM ("
+                "SELECT chunk_id, distance FROM memory_vectors "
+                "WHERE embedding MATCH ? AND k = ? AND group_id = ?"
+                ") AS v JOIN chunks AS c ON c.id = v.chunk_id ORDER BY v.distance",
+                (sqlite_vec.serialize_float32(embedding), limit, group_id),
+            )
+            rows = await cursor.fetchall()
+            await cursor.close()
+        return [_row_to_chunk(row) for row in rows]
+
+    async def clear_group(self, group_id: int) -> None:
+        """删除当前群的消息、片段与向量，保留概率配置。"""
+        async with self._transaction() as db:
+            await db.execute("DELETE FROM memory_vectors WHERE group_id = ?", (group_id,))
+            await db.execute("DELETE FROM messages WHERE group_id = ?", (group_id,))
+            await db.execute("DELETE FROM chunks WHERE group_id = ?", (group_id,))
+
+    def _require_db(self) -> aiosqlite.Connection:
+        if self._db is None:
+            raise RuntimeError("数据库尚未打开，请先调用 MemoryStore.open()。")
+        return self._db
+
+    @asynccontextmanager
+    async def _transaction(self) -> AsyncIterator[aiosqlite.Connection]:
+        db = self._require_db()
+        async with self._lock, self._transaction_on(db) as tx:
+            yield tx
+
+    @asynccontextmanager
+    async def _transaction_on(
+        self, db: aiosqlite.Connection
+    ) -> AsyncIterator[aiosqlite.Connection]:
+        await db.execute("BEGIN IMMEDIATE")
+        try:
+            yield db
+        except BaseException:
+            await db.rollback()
+            raise
+        else:
+            await db.commit()
+
+    async def _ensure_group(self, db: aiosqlite.Connection, group_id: int) -> None:
+        await db.execute(
+            "INSERT INTO groups(group_id, probability, updated_at) VALUES (?, ?, ?) "
+            "ON CONFLICT(group_id) DO NOTHING",
+            (group_id, self._default_probability, _format_timestamp(_utc_now())),
+        )
+
+    async def _ensure_dimensions(self, db: aiosqlite.Connection) -> None:
+        cursor = await db.execute("SELECT value FROM meta WHERE key = ?", (DIMENSIONS_META_KEY,))
+        row = await cursor.fetchone()
+        await cursor.close()
+        if row is None:
+            await db.execute(
+                "INSERT INTO meta(key, value) VALUES (?, ?)",
+                (DIMENSIONS_META_KEY, str(self._dimensions)),
+            )
+            return
+        stored = int(row["value"])
+        if stored != self._dimensions:
+            raise RuntimeError(
+                f"数据库向量维度为 {stored}，与配置的 {self._dimensions} 不一致；"
+                "请恢复原配置或删除数据库文件后重建。"
+            )
+
+
+def _row_to_message(row: aiosqlite.Row) -> Message:
+    return Message(
+        group_id=row["group_id"],
+        user_id=row["user_id"],
+        display_name=row["display_name"],
+        text=row["text"],
+        created_at=_parse_timestamp(row["created_at"]),
+        is_bot=bool(row["is_bot"]),
+    )
+
+
+def _row_to_chunk(row: aiosqlite.Row) -> Chunk:
+    return Chunk(
+        id=row["id"],
+        group_id=row["group_id"],
+        text=row["text"],
+        status=row["status"],
+    )
+
+
+def _format_chunk_text(group_id: int, rows: Sequence[aiosqlite.Row]) -> str:
+    first = _parse_timestamp(rows[0]["created_at"])
+    last = _parse_timestamp(rows[-1]["created_at"])
+    lines = [f"群 {group_id} 聊天记录（{first:%Y-%m-%d %H:%M} ~ {last:%Y-%m-%d %H:%M} UTC）"]
+    for row in rows:
+        created_at = _parse_timestamp(row["created_at"])
+        name = row["display_name"].strip() or str(row["user_id"])
+        suffix = "(bot)" if row["is_bot"] else ""
+        lines.append(f"[{created_at:%Y-%m-%d %H:%M}] {name}{suffix}: {row['text']}")
+    return "\n".join(lines)
+
+
+def _validated_vector(vector: list[float], dimensions: int) -> list[float]:
+    if len(vector) != dimensions:
+        raise ValueError(f"向量维度为 {len(vector)}，配置要求 {dimensions}。")
+    values = [float(item) for item in vector]
+    if not all(math.isfinite(item) for item in values):
+        raise ValueError("向量包含非有限数值。")
+    return values
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _format_timestamp(value: datetime) -> str:
+    return _as_utc(value).isoformat()
+
+
+def _parse_timestamp(value: str) -> datetime:
+    return _as_utc(datetime.fromisoformat(value))
+
+
+def _as_utc(value: datetime) -> datetime:
+    if value.tzinfo is None:
+        return value.replace(tzinfo=timezone.utc)
+    return value.astimezone(timezone.utc)