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

nonebot_plugin_group_chat_memory/handlers.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from datetime import datetime, timezone
+from typing import Annotated
+from urllib.parse import urlparse
+from weakref import WeakKeyDictionary
+
+from nonebot import get_driver, on_command, on_message
+from nonebot import logger as nonebot_logger
+from nonebot.adapters.onebot.v11 import (
+    Bot,
+    GroupMessageEvent,
+    Message,
+    MessageEvent,
+    MessageSegment,
+    PrivateMessageEvent,
+)
+from nonebot.params import CommandArg
+from nonebot.permission import SUPERUSER
+
+from .api import ChatClient, EmbeddingClient
+from .config import get_settings
+from .context import ContextBuilder
+from .memory import MemoryWriter
+from .models import Message as MemoryMessage
+from .replies import ReplyCoordinator
+from .store import MemoryStore
+
+logger = logging.getLogger(__name__)
+
+GENERIC_REPLY_ERROR = "暂时无法生成回复，请稍后再试。"
+PERMISSION_DENIED = "仅 SUPERUSER 可执行此命令。"
+INVALID_PROBABILITY = "概率需为 0～100 的数字，最多一位小数。"
+PROBABILITY_PATTERN = re.compile(r"(?:100(?:\.0)?|\d{1,2}(?:\.\d)?)")
+SECRET_COMMAND_LOG_PATTERN = re.compile(r"(设置聊天密钥|设置向量密钥)\s+[^\r\n]*")
+
+
+def _redact_secret_command_logs(record) -> None:
+    record["message"] = SECRET_COMMAND_LOG_PATTERN.sub(r"\1 ****", record["message"])
+
+
+nonebot_logger.configure(patcher=_redact_secret_command_logs)
+
+config = get_settings()
+store = MemoryStore(config)
+embedding_client = EmbeddingClient(config)
+writer = MemoryWriter(config, store, embedding_client)
+context_builder = ContextBuilder(config, store, embedding_client)
+chat_client = ChatClient(config)
+_bots: dict[int, Bot] = {}
+
+
+async def startup() -> None:
+    await store.open()
+    config.apply_runtime_config(await store.get_runtime_config())
+    writer.start()
+
+
+async def shutdown() -> None:
+    try:
+        try:
+            try:
+                await writer.flush(timeout=5)
+            except asyncio.TimeoutError:
+                logger.warning("关闭插件时等待记忆写入超时，将保留未完成数据供下次启动恢复。")
+        finally:
+            await writer.close()
+    finally:
+        try:
+            await coordinator.close()
+        finally:
+            await store.close()
+
+
+async def reply(group_id: int, message: MemoryMessage) -> None:
+    bot = _bots[group_id]
+    try:
+        context = await context_builder.build(group_id, message.text)
+        response = await chat_client.complete(context, message.text)
+    except Exception as exc:
+        logger.error("群 %s 的聊天回复失败：%s", group_id, type(exc).__name__)
+        await bot.send_group_msg(group_id=group_id, message=GENERIC_REPLY_ERROR)
+        return
+
+    await bot.send_group_msg(group_id=group_id, message=MessageSegment.text(response))
+    await writer.record(
+        MemoryMessage(
+            group_id=group_id,
+            user_id=int(bot.self_id),
+            display_name="Bot",
+            text=response,
+            is_bot=True,
+        )
+    )
+
+
+coordinator = ReplyCoordinator(store, reply)
+
+
+async def _eligible_group_message(bot: Bot, event: GroupMessageEvent) -> bool:
+    if not isinstance(event, GroupMessageEvent):
+        return False
+    if event.anonymous is not None:
+        return False
+    if event.user_id in {event.self_id, int(bot.self_id)}:
+        return False
+    return event.user_id not in config.group_chat_memory_ignored_qq
+
+
+settings_locks: WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = WeakKeyDictionary()
+
+
+def _settings_lock() -> asyncio.Lock:
+    loop = asyncio.get_running_loop()
+    if (lock := settings_locks.get(loop)) is None:
+        lock = asyncio.Lock()
+        settings_locks[loop] = lock
+    return lock
+
+
+async def apply_setting(key: str, value: str) -> None:
+    async with _settings_lock():
+        values = {
+            "chat_base_url": config.group_chat_memory_chat_base_url,
+            "chat_api_key": config.group_chat_memory_chat_api_key.get_secret_value(),
+            "chat_model": config.group_chat_memory_chat_model,
+            "embedding_api_key": config.group_chat_memory_embedding_api_key.get_secret_value(),
+        }
+        values[key] = value
+        await store.set_runtime_config(key, value)
+        config.apply_runtime_config(values)
+
+
+chat_url_matcher = on_command("设置聊天地址", priority=5, block=True)
+chat_key_matcher = on_command("设置聊天密钥", priority=5, block=True)
+chat_model_matcher = on_command("设置聊天模型", priority=5, block=True)
+embedding_key_matcher = on_command("设置向量密钥", priority=5, block=True)
+model_config_matcher = on_command("模型配置", priority=5, block=True)
+
+
+def _valid_http_url(value: str) -> bool:
+    if any(
+        character.isspace() or ord(character) < 32 or ord(character) == 127
+        for character in value
+    ):
+        return False
+    try:
+        parsed = urlparse(value)
+        if parsed.scheme not in {"http", "https"} or not parsed.hostname:
+            return False
+        _ = parsed.port
+    except ValueError:
+        return False
+    return True
+
+
+async def _handle_setting(matcher, bot: Bot, event: MessageEvent, args: Message, key: str) -> None:
+    if not isinstance(event, PrivateMessageEvent):
+        await matcher.finish("请在私聊中设置模型配置。")
+    if not await SUPERUSER(bot, event):
+        await matcher.finish(PERMISSION_DENIED)
+
+    value = args.extract_plain_text().strip()
+    if not value:
+        await matcher.finish("配置参数不能为空。")
+    if key == "chat_base_url" and not _valid_http_url(value):
+        await matcher.finish("聊天 API 地址需为有效的 http(s) URL。")
+    try:
+        await apply_setting(key, value)
+    except Exception as exc:
+        logger.error("保存模型配置失败：%s", type(exc).__name__)
+        await matcher.finish("配置保存失败，请稍后再试。")
+    await matcher.finish("模型配置已更新。")
+
+
+@chat_url_matcher.handle()
+async def handle_chat_url(
+    bot: Bot, event: MessageEvent, args: Annotated[Message, CommandArg()]
+) -> None:
+    await _handle_setting(chat_url_matcher, bot, event, args, "chat_base_url")
+
+
+@chat_key_matcher.handle()
+async def handle_chat_key(
+    bot: Bot, event: MessageEvent, args: Annotated[Message, CommandArg()]
+) -> None:
+    await _handle_setting(chat_key_matcher, bot, event, args, "chat_api_key")
+
+
+@chat_model_matcher.handle()
+async def handle_chat_model(
+    bot: Bot, event: MessageEvent, args: Annotated[Message, CommandArg()]
+) -> None:
+    await _handle_setting(chat_model_matcher, bot, event, args, "chat_model")
+
+
+@embedding_key_matcher.handle()
+async def handle_embedding_key(
+    bot: Bot, event: MessageEvent, args: Annotated[Message, CommandArg()]
+) -> None:
+    await _handle_setting(embedding_key_matcher, bot, event, args, "embedding_api_key")
+
+
+@model_config_matcher.handle()
+async def handle_model_config(bot: Bot, event: MessageEvent) -> None:
+    if not await SUPERUSER(bot, event):
+        await model_config_matcher.finish(PERMISSION_DENIED)
+
+    def status(secret: str) -> str:
+        return "已配置（****）" if secret else "未配置"
+
+    await model_config_matcher.finish(
+        f"聊天地址：{config.group_chat_memory_chat_base_url or '未配置'}\n"
+        f"聊天模型：{config.group_chat_memory_chat_model}\n"
+        f"聊天密钥：{status(config.group_chat_memory_chat_api_key.get_secret_value())}\n"
+        "Embedding 密钥："
+        f"{status(config.group_chat_memory_embedding_api_key.get_secret_value())}"
+    )
+
+
+probability_matcher = on_command(
+    "回复概率", rule=_eligible_group_message, priority=5, block=True
+)
+clear_matcher = on_command("清空群记忆", rule=_eligible_group_message, priority=5, block=True)
+group_message_matcher = on_message(rule=_eligible_group_message, priority=20, block=False)
+
+
+@probability_matcher.handle()
+async def handle_probability(
+    bot: Bot,
+    event: GroupMessageEvent,
+    args: Annotated[Message, CommandArg()],
+) -> None:
+    if not await SUPERUSER(bot, event):
+        await probability_matcher.finish(PERMISSION_DENIED)
+
+    text = args.extract_plain_text().strip()
+    if not text:
+        probability = await store.get_probability(event.group_id)
+        await probability_matcher.finish(f"当前群随机回复概率：{probability:.1f}%")
+    if PROBABILITY_PATTERN.fullmatch(text) is None:
+        await probability_matcher.finish(INVALID_PROBABILITY)
+
+    probability = float(text)
+    await store.set_probability(event.group_id, probability)
+    await probability_matcher.finish(f"已设置当前群随机回复概率：{probability:.1f}%")
+
+
+@clear_matcher.handle()
+async def handle_clear(bot: Bot, event: GroupMessageEvent) -> None:
+    if not await SUPERUSER(bot, event):
+        await clear_matcher.finish(PERMISSION_DENIED)
+    await store.clear_group(event.group_id)
+    await clear_matcher.finish("已清空当前群记忆。")
+
+
+@group_message_matcher.handle()
+async def handle_group_message(bot: Bot, event: GroupMessageEvent) -> None:
+    explicit_at = any(
+        segment.type == "at" and segment.data.get("qq") == str(bot.self_id)
+        for segment in event.message
+    )
+    directed = explicit_at or (
+        event.reply is not None and event.reply.sender.user_id == int(bot.self_id)
+    )
+    allowed_types = {"text", "at", "reply"}
+    if any(segment.type not in allowed_types for segment in event.message):
+        return
+
+    text = event.message.extract_plain_text().strip()
+    if not text:
+        return
+
+    memory_message = MemoryMessage(
+        group_id=event.group_id,
+        user_id=event.user_id,
+        display_name=event.sender.card or event.sender.nickname or str(event.user_id),
+        text=text,
+        created_at=datetime.fromtimestamp(event.time, timezone.utc),
+    )
+    _bots[event.group_id] = bot
+    await writer.record(memory_message)
+    if directed:
+        coordinator.submit(event.group_id, memory_message, explicit=True)
+    elif (
+        event.reply is None
+        and all(segment.type == "text" for segment in event.message)
+        and await coordinator.should_random_reply(event.group_id)
+    ):
+        coordinator.submit(event.group_id, memory_message, explicit=False)
+
+
+driver = get_driver()
+driver.on_startup(startup)
+driver.on_shutdown(shutdown)

nonebot_plugin_group_chat_memory/memory.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import Awaitable, Callable
+from typing import NamedTuple
+
+from .api import EmbeddingClient
+from .config import Config
+from .models import Chunk, Message
+from .store import MemoryStore
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_RETRY_DELAYS: tuple[float, ...] = (1.0, 2.0, 4.0)
+DEFAULT_QUEUE_SIZE = 1024
+
+
+class _Wake(NamedTuple):
+    """提示后台任务为某个群尝试分块；同群唤醒会合并。"""
+
+    group_id: int
+
+
+class _Stop:
+    """关闭后台任务的哨兵。"""
+
+
+_STOP = _Stop()
+
+_Item = _Wake | Chunk | _Stop
+
+
+class MemoryWriter:
+    """在后台把已持久化的群消息分块并写入长期记忆向量。
+
+    `record` 只负责持久化原文并唤醒后台任务；分块、Embedding 与写入向量都在
+    单一后台任务中完成，任何失败都不会抛回消息处理流程。后台任务被意外取消或
+    抛出 `BaseException` 时由监督回调记录并允许重新 `start()`，`flush()` 不会
+    无限等待。关闭前应先 `flush()`，`close()` 会取消退避休眠并等待在途请求
+    收尾。
+    """
+
+    def __init__(
+        self,
+        config: Config,
+        store: MemoryStore,
+        embedding: EmbeddingClient,
+        *,
+        retry_delays: tuple[float, ...] = DEFAULT_RETRY_DELAYS,
+        queue_size: int = DEFAULT_QUEUE_SIZE,
+        sleep: Callable[[float], Awaitable[None]] | None = None,
+    ) -> None:
+        self._store = store
+        self._embedding = embedding
+        self._chunk_size = config.group_chat_memory_chunk_size
+        self._retry_delays = tuple(retry_delays)
+        self._sleep_impl = sleep
+        self._queue: asyncio.Queue[_Item] = asyncio.Queue(maxsize=max(1, queue_size))
+        self._stop_event = asyncio.Event()
+        self._recovered = asyncio.Event()
+        self._queued_wakes: set[int] = set()
+        self._dirty_groups: set[int] = set()
+        self._overflow_warned = False
+        self._closing = False
+        self._task: asyncio.Task[None] | None = None
+
+    @property
+    def running(self) -> bool:
+        """后台任务是否仍在运行。"""
+        return self._is_running()
+
+    def start(self) -> None:
+        """启动后台任务并恢复待处理记忆；已在运行时不做任何事。"""
+        if self._is_running():
+            return
+        self._task = None
+        self._closing = False
+        self._stop_event.clear()
+        self._recovered.clear()
+        self._queued_wakes.clear()
+        task = asyncio.create_task(self._run(), name="group-chat-memory-writer")
+        self._task = task
+        task.add_done_callback(self._on_task_done)
+
+    async def record(self, message: Message) -> None:
+        """先持久化消息原文，再交给后台任务分块；数据库失败只记录日志。"""
+        try:
+            await self._store.add_message(message)
+        except Exception:
+            logger.exception("群 %s 的消息写入数据库失败，已跳过长期记忆。", message.group_id)
+            return
+        self._dirty_groups.add(message.group_id)
+        if self._closing:
+            return
+        self._wake(message.group_id)
+
+    async def flush(self, timeout: float | None = None) -> None:
+        """等待已排队的分块与向量写入全部结束。
+
+        传入 `timeout`（秒）时最多等待该时长，超时抛出 `asyncio.TimeoutError`；
+        该异常只表示本次等待超时，后台任务与后续 `close()` 都不受影响。后台任务
+        未运行时立即返回，未完成的群聊消息由下次启动恢复。
+        """
+        if timeout is None:
+            await self._wait_until_idle()
+            return
+        await asyncio.wait_for(self._wait_until_idle(), timeout)
+
+    async def close(self) -> None:
+        """停止后台任务：取消退避休眠，等待在途请求结束，可重复调用。
+
+        等待时使用 `asyncio.shield`，调用方取消 `close()` 不会级联取消后台任务，
+        取消异常照常向上抛出；只有后台任务本身被取消时才按关闭处理并只记录日志。
+        """
+        task = self._task
+        if task is None or task.done():
+            return
+        first = not self._closing
+        self._closing = True
+        self._stop_event.set()
+        if first:
+            with contextlib.suppress(asyncio.QueueFull):
+                self._queue.put_nowait(_STOP)
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            if not task.cancelled():
+                raise
+            if first:
+                logger.warning("记忆后台任务在关闭时被取消，未完成的消息由下次启动补做。")
+        except Exception:
+            if first:
+                logger.exception("记忆后台任务异常退出。")
+
+    def _is_running(self) -> bool:
+        return self._task is not None and not self._task.done()
+
+    def _on_task_done(self, task: asyncio.Task[None]) -> None:
+        if self._task is task:
+            self._task = None
+            self._recovered.set()
+        if task.cancelled():
+            if not self._closing:
+                logger.warning("记忆后台任务被意外取消，未处理的消息将在下次启动时补做。")
+            return
+        error = task.exception()
+        if error is not None and not self._closing:
+            logger.error("记忆后台任务意外退出：%r", error, exc_info=error)
+
+    def _wake(self, group_id: int) -> None:
+        if not self._is_running() or group_id in self._queued_wakes:
+            return
+        try:
+            self._queue.put_nowait(_Wake(group_id))
+        except asyncio.QueueFull:
+            if not self._overflow_warned:
+                self._overflow_warned = True
+                logger.warning("记忆队列已满，未完成的分块将在下次启动时补做。")
+            return
+        self._queued_wakes.add(group_id)
+        self._overflow_warned = False
+
+    async def _wait_until_idle(self) -> None:
+        if not self._is_running():
+            self._warn_stranded()
+            return
+        await self._recovered.wait()
+        if not self._is_running():
+            self._warn_stranded()
+            return
+        await self._queue.join()
+
+    def _warn_stranded(self) -> None:
+        if not self._queue.empty():
+            logger.warning(
+                "记忆后台任务未运行，%s 项排队工作将在下次启动时补做。", self._queue.qsize()
+            )
+
+    async def _run(self) -> None:
+        try:
+            try:
+                await self._recover_pending()
+            finally:
+                self._recovered.set()
+            while True:
+                item = await self._queue.get()
+                try:
+                    if isinstance(item, _Wake):
+                        self._queued_wakes.discard(item.group_id)
+                    if isinstance(item, _Stop) or self._closing:
+                        return
+                    if isinstance(item, _Wake):
+                        await self._chunk_group(item.group_id)
+                    else:
+                        await self._embed_chunk(item)
+                finally:
+                    self._queue.task_done()
+        finally:
+            self._drain_queue()
+
+    def _drain_queue(self) -> None:
+        while True:
+            try:
+                item = self._queue.get_nowait()
+            except asyncio.QueueEmpty:
+                return
+            if isinstance(item, _Wake):
+                self._queued_wakes.discard(item.group_id)
+                self._dirty_groups.add(item.group_id)
+            self._queue.task_done()
+
+    async def _recover_pending(self) -> None:
+        try:
+            chunks = await self._store.pending_chunks()
+        except Exception:
+            logger.exception("读取待处理记忆片段失败，已跳过片段恢复。")
+            chunks = []
+        try:
+            groups = await self._store.pending_groups(self._chunk_size)
+        except Exception:
+            logger.exception("读取待处理群聊消息失败，已跳过补分块恢复。")
+            groups = []
+        for chunk in chunks:
+            if self._closing:
+                return
+            await self._embed_chunk(chunk)
+        for group_id in dict.fromkeys([*groups, *sorted(self._dirty_groups)]):
+            if self._closing:
+                return
+            self._dirty_groups.add(group_id)
+            await self._chunk_group(group_id)
+
+    async def _chunk_group(self, group_id: int) -> None:
+        while not self._closing:
+            try:
+                chunk = await self._store.create_ready_chunk(group_id, self._chunk_size)
+            except Exception:
+                logger.exception("群 %s 的消息分块失败，原文已保留。", group_id)
+                return
+            if chunk is None:
+                self._dirty_groups.discard(group_id)
+                return
+            await self._embed_chunk(chunk)
+
+    async def _embed_chunk(self, chunk: Chunk) -> None:
+        attempts = len(self._retry_delays) + 1
+        for attempt in range(attempts):
+            try:
+                vectors = await self._embedding.embed([chunk.text])
+            except Exception as exc:
+                if attempt + 1 >= attempts:
+                    logger.error("片段 %s 向量化失败，已保留待处理状态：%s", chunk.id, exc)
+                    return
+                logger.warning("片段 %s 向量化失败，将重试：%s", chunk.id, exc)
+                if not await self._sleep(self._retry_delays[attempt]):
+                    return
+                continue
+            if len(vectors) != 1:
+                logger.error(
+                    "片段 %s 的 Embedding 返回 %s 条向量，期望 1 条。", chunk.id, len(vectors)
+                )
+                return
+            try:
+                await self._store.complete_chunk(chunk.id, vectors[0])
+            except Exception:
+                logger.exception("片段 %s 写入向量失败，已保留待处理状态。", chunk.id)
+            return
+
+    async def _sleep(self, delay: float) -> bool:
+        """等待退避时间；返回 False 表示等待期间收到关闭请求。"""
+        if self._sleep_impl is not None:
+            await self._sleep_impl(delay)
+            return not self._closing
+        if delay <= 0:
+            await asyncio.sleep(0)
+            return not self._closing
+        try:
+            await asyncio.wait_for(self._stop_event.wait(), timeout=delay)
+        except asyncio.TimeoutError:
+            return True
+        return False

nonebot_plugin_group_chat_memory/models.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class Message(BaseModel):
+    """群聊中的一条文本消息。"""
+
+    group_id: int
+    user_id: int
+    display_name: str
+    text: str
+    created_at: datetime = Field(default_factory=_utc_now)
+    is_bot: bool = False
+
+
+class ChunkStatus(str, Enum):
+    PENDING = "pending"
+    READY = "ready"
+
+
+class Chunk(BaseModel):
+    """一段待向量化或已向量化的群聊记忆。"""
+
+    id: int
+    group_id: int
+    text: str
+    status: ChunkStatus = ChunkStatus.PENDING
+
+
+class ChatContext(BaseModel):
+    """一次模型请求所需的最近消息与长期记忆。"""
+
+    recent_messages: list[Message] = Field(default_factory=list)
+    long_term_chunks: list[Chunk] = Field(default_factory=list)