✨ use msg_id to cache generated unimsg

Author: RF-Tar-Railt

example/plugins/demo.py

@@ -1,6 +1,7 @@
 from typing import Union, Literal
 
 from nonebot import require
+from nonebot.internal.adapter import Event
 from nonebot.adapters.onebot.v12 import Bot
 from importlib_metadata import distributions
 from nonebot.adapters.onebot.v12.event import GroupMessageDeleteEvent
@@ -516,3 +517,43 @@ async def setu_h(count: int, tags: Match[tuple[str, ...]], r18: bool):
 @preview.handle()
 async def preview_h(content: UniMessage):
     await preview.finish("rendering preview:\n" + content)
+
+
+from arclet.alconna.tools.debug import analyse_header
+
+
+class TestExtension1(Extension):
+    @property
+    def priority(self) -> int:
+        return 12
+
+    @property
+    def id(self) -> str:
+        return "test1"
+
+    def post_init(self, alc: Alconna) -> None:
+        def test_func(msg: str):
+            test_res = analyse_header(alc.prefixes, alc.command, msg, raise_exception=False)
+            return test_res and test_res.matched
+
+        self.test_func = test_func
+
+    async def call_llm(self, msg: str):
+        return "calculate " + "".join(c for c in msg if ord(c) < 256)
+
+    async def receive_wrapper(self, bot: Bot, event: Event, command: Alconna, receive: UniMessage) -> UniMessage:
+        msg = receive.extract_plain_text()
+        if self.test_func(msg):
+            return receive
+        output = await self.call_llm(msg)
+        return UniMessage(output)
+
+
+calculate = Command("calculate <expression>", "计算").build(
+    auto_send_output=True, extensions=[TestExtension1()], use_cmd_start=False
+)
+
+
+@calculate.handle()
+async def calculate_h(expression: str):
+    await calculate.finish(expression)

src/nonebot_plugin_alconna/builtins/extensions/discord.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from tarina import lang
+from tarina import LRU, lang
 from arclet.alconna import Alconna
 from nonebot.adapters import Event
 from nonebot.typing import T_State
@@ -71,6 +71,7 @@ def __init__(
         self.nsfw = nsfw
         super().__init__()
         self.using = False
+        self.cache: "LRU[int, UniMessage]" = LRU(20)  # noqa: UP037
 
     def post_init(self, alc: Alconna) -> None:
         if "/" not in alc.prefixes or (
@@ -123,6 +124,8 @@ def validate(self, bot, event: Event) -> bool:
     async def message_provider(self, event: Event, state: T_State, bot, use_origin: bool = False):
         if not isinstance(event, ApplicationCommandInteractionEvent):
             return None
+        if event.id in self.cache:
+            return self.cache[event.id]
         data = event.data
         cmd = f"/{data.name}"
 
@@ -142,7 +145,9 @@ def _handle_options(options: list[ApplicationCommandInteractionDataOption]):
             cmd += " "
             cmd += " ".join(_handle_options(data.options))
 
-        return UniMessage(cmd.rstrip())
+        res = UniMessage(cmd.rstrip())
+        self.cache[event.id] = res
+        return res
 
     @classmethod
     async def send_deferred_response(cls) -> None:

src/nonebot_plugin_alconna/builtins/extensions/reply.py

@@ -76,7 +76,7 @@ def __init__(self, add_left: bool = False, sep: str = " "):
         self.add_left = add_left
         self.sep = sep
 
-    cache: "LRU[int, UniMessage]" = LRU(20)
+    cache: "LRU[str, UniMessage]" = LRU(20)
 
     @property
     def priority(self) -> int:
@@ -87,17 +87,17 @@ def id(self) -> str:
         return "builtins.extensions.reply:ReplyMergeExtension"
 
     async def message_provider(self, event, state, bot, use_origin: bool = False):
-        event_id = id(event)
-        if event_id in self.cache:
-            return self.cache[event_id]
         if event.get_type() != "message":
             return None
         try:
             msg = event.get_message()
         except (NotImplementedError, ValueError):
             return None
+        msg_id = UniMessage.get_message_id(event, bot)
+        if msg_id in self.cache:
+            return self.cache[msg_id]
         uni_msg = UniMessage.generate_sync(message=msg, bot=bot)
-        self.cache[event_id] = uni_msg
+        self.cache[msg_id] = uni_msg
         if not (reply := await reply_fetch(event, bot)):
             return uni_msg
         if not reply.msg:
@@ -109,11 +109,11 @@ async def message_provider(self, event, state, bot, use_origin: bool = False):
         if self.add_left:
             uni_msg_reply += self.sep
             uni_msg_reply.extend(uni_msg)
-            self.cache[event_id] = uni_msg_reply
+            self.cache[msg_id] = uni_msg_reply
             return uni_msg_reply
         uni_msg += self.sep
         uni_msg.extend(uni_msg_reply)
-        self.cache[event_id] = uni_msg
+        self.cache[msg_id] = uni_msg
         return uni_msg
 
 

src/nonebot_plugin_alconna/builtins/plugins/with/extension.py

@@ -23,21 +23,21 @@ def id(self) -> str:
     prefixes: list[str]
     command: str
     sep: str
-    cache: "LRU[int, UniMessage]" = LRU(20)
+    cache: "LRU[str, UniMessage]" = LRU(20)
 
     def post_init(self, alc: Alconna) -> None:
         self.prefixes = [pf for pf in alc.prefixes if isinstance(pf, str)]
         self.command = alc.header_display
         self.sep = alc.separators[0]
 
     async def receive_wrapper(self, bot: Bot, event: Event, command: Alconna, receive: UniMessage) -> UniMessage:
-        event_id = id(event)
-        if event_id in self.cache:
-            return self.cache[event_id]
+        msg_id = UniMessage.get_message_id(event, bot)
+        if msg_id in self.cache:
+            return self.cache[msg_id]
         target = UniMessage.get_target(event, bot)
         prefix = self.supplier(target)
         if not prefix or not command.header_display.endswith(prefix):
             return receive
         res = UniMessage.text(random.choice(self.prefixes) + prefix + self.sep) + receive
-        self.cache[event_id] = res
+        self.cache[msg_id] = res
         return res

src/nonebot_plugin_alconna/extension.py

@@ -125,7 +125,7 @@ def id(self) -> str:
 
 _callbacks = set()
 
-unimsg_cache: LRU[int, UniMessage] = LRU(16)
+unimsg_cache: LRU[str, UniMessage] = LRU(16)
 
 
 class ExtensionExecutor:
@@ -208,24 +208,17 @@ async def output_converter(self, output_type: OutputType, content: str) -> UniMe
     async def message_provider(
         self, event: Event, state: T_State, bot: Bot, use_origin: bool = False
     ) -> UniMessage | None:
-        event_id = id(event)
-        if (uni_msg := unimsg_cache.get(event_id)) is not None:
-            msg = uni_msg
-        elif event.get_type() != "message":
-            msg = None
-        else:
-            try:
-                msg = event.get_message()
-            except (NotImplementedError, ValueError):
-                msg = None
+        if event.get_type().startswith("message"):
+            msg = event.get_message()
             if use_origin:
-                try:
-                    msg = getattr(event, "original_message", msg)  # type: ignore
-                except (NotImplementedError, ValueError):
-                    pass
-            if msg is not None:
+                msg = getattr(event, "original_message", None) or msg  # type: ignore
+            msg_id = UniMessage.get_message_id(event, bot)
+            if (uni_msg := unimsg_cache.get(msg_id)) is not None:
+                msg = uni_msg
+            else:
                 msg = UniMessage.generate_without_reply(message=msg, bot=bot)
-                unimsg_cache[event_id] = msg
+                unimsg_cache[msg_id] = msg
+            return msg
         exc = None
         for ext in self.context:
             if not ext._overrides["message_provider"]:
@@ -238,7 +231,7 @@ async def message_provider(
         if exc is not None:
             raise exc
 
-        return msg
+        return None
 
     async def receive_wrapper(self, bot: Bot, event: Event, command: Alconna, receive: UniMessage) -> UniMessage:
         res = receive

src/nonebot_plugin_alconna/rule.py

@@ -2,6 +2,7 @@
 import weakref
 from typing import Any, Union, Literal, Optional
 
+import nonebot
 from nonebot.typing import T_State
 from tarina import lang, init_spec
 from nonebot.matcher import Matcher
@@ -13,9 +14,6 @@
 from arclet.alconna.exceptions import SpecialOptionTriggered
 from arclet.alconna import Alconna, Arparma, CompSession, output_manager, command_manager
 
-require("nonebot_plugin_waiter")
-from nonebot_plugin_waiter import waiter
-
 from .i18n import Lang
 from .config import Config
 from .uniseg import UniMsg, UniMessage
@@ -24,7 +22,14 @@
 from .extension import Extension, ExtensionExecutor
 from .consts import ALCONNA_RESULT, ALCONNA_EXTENSION, ALCONNA_EXEC_RESULT, log
 
-_modules = set()
+try:
+    if nonebot._driver:
+        require("nonebot_plugin_waiter")
+        from nonebot_plugin_waiter import waiter
+    else:
+        waiter = None
+except (RuntimeError, ValueError):
+    waiter = None
 
 
 def check_self_send(bot: Bot, event: Event) -> bool:
@@ -232,7 +237,7 @@ async def handle(
         def _checker(_event: Event):
             return session_id == _event.get_session_id()
 
-        w = waiter(["message"], Matcher, keep_session=True, block=False, rule=Rule(_checker))(self._waiter)
+        w = waiter(["message"], Matcher, keep_session=True, block=False, rule=Rule(_checker))(self._waiter)  # type: ignore
 
         while interface.available:
 

src/nonebot_plugin_alconna/uniseg/adapters/qq/exporter.py

@@ -376,7 +376,11 @@ async def send_to(self, target: Union[Target, Event], bot: Bot, message: Message
         message = message.exclude("mention_channel", "mention_user", "mention_everyone", "reference")
         if target.private:
             res = await bot.send_to_c2c(
-                openid=target.id, message=message, msg_id=target.source, msg_seq=target.extra.get("qq.reply_seq"), **kwargs
+                openid=target.id,
+                message=message,
+                msg_id=target.source,
+                msg_seq=target.extra.get("qq.reply_seq"),
+                **kwargs,
             )
         elif target.extra.get("qq.interaction", False):
             return await bot.send_to_group(group_openid=target.id, message=message, event_id=target.source, **kwargs)

src/nonebot_plugin_alconna/uniseg/message.py

@@ -1268,6 +1268,8 @@ def get_message_id(event: Event | None = None, bot: Bot | None = None, adapter:
                 event = current_event.get()
             except LookupError as e:
                 raise SerializeFailed(lang.require("nbp-uniseg", "event_missing")) from e
+        if hasattr(event, "__uniseg_message_id__"):
+            return event.__uniseg_message_id__  # type: ignore
         if not adapter:
             if not bot:
                 try:
@@ -1277,7 +1279,8 @@ def get_message_id(event: Event | None = None, bot: Bot | None = None, adapter:
             _adapter = bot.adapter
             adapter = _adapter.get_name()
         if fn := alter_get_exporter(adapter):
-            return fn.get_message_id(event)
+            setattr(event, "__uniseg_message_id__", msg_id := fn.get_message_id(event))
+            return msg_id
         raise SerializeFailed(lang.require("nbp-uniseg", "unsupported").format(adapter=adapter))
 
     @staticmethod

tests/fake.py

@@ -12,6 +12,13 @@
     from nonebot.adapters.discord.event import GuildMessageCreateEvent as DiscordMessageEvent
 
 
+_msg_ids = iter(range(1000000))
+
+
+def get_msg_id() -> int:
+    return next(_msg_ids)
+
+
 def fake_group_message_event_v11(**field) -> "GroupMessageEventV11":
     from pydantic import create_model
     from nonebot.adapters.onebot.v11.event import Sender
@@ -27,7 +34,6 @@ class FakeEvent(_fake):
         user_id: int = 10
         message_type: Literal["group"] = "group"
         group_id: int = 10000
-        message_id: int = 1
         message: Message = Message("test")
         raw_message: str = "test"
         font: int = 0
@@ -41,7 +47,7 @@ class FakeEvent(_fake):
         class Config:
             extra = "allow"
 
-    return FakeEvent(**field)
+    return FakeEvent(message_id=get_msg_id(), **field)
 
 
 def fake_private_message_event_v11(**field) -> "PrivateMessageEventV11":
@@ -58,7 +64,6 @@ class FakeEvent(_fake):
         sub_type: str = "friend"
         user_id: int = 10
         message_type: Literal["private"] = "private"
-        message_id: int = 1
         message: Message = Message("test")
         raw_message: str = "test"
         font: int = 0
@@ -68,7 +73,7 @@ class FakeEvent(_fake):
         class Config:
             extra = "forbid"
 
-    return FakeEvent(**field)
+    return FakeEvent(message_id=get_msg_id(), **field)
 
 
 def fake_discord_interaction_event(**field) -> "ApplicationCommandInteractionEvent":
@@ -77,7 +82,7 @@ def fake_discord_interaction_event(**field) -> "ApplicationCommandInteractionEve
 
     _fake = create_model("_fake", __base__=ApplicationCommandInteractionEvent)
     field["type"] = 2
-    field["id"] = 123456
+    field["id"] = get_msg_id() + 123456
     field["application_id"] = 123456789
     field["token"] = "sometoken"  # noqa: S105
     field["version"] = 1
@@ -95,7 +100,7 @@ def fake_message_event_discord(content: str) -> "DiscordMessageEvent":
     return type_validate_python(
         GuildMessageCreateEvent,
         {
-            "id": 11234,
+            "id": get_msg_id() + 11234,
             "channel_id": 5566,
             "guild_id": 6677,
             "author": {
@@ -153,7 +158,7 @@ class Config:
             extra = "allow"
 
     _message = field.pop("message", Message("test"))
-    event = FakeEvent(message={"id": "1", "content": "text"}, **field)  # type: ignore
+    event = FakeEvent(message={"id": str(get_msg_id()), "content": "text"}, **field)  # type: ignore
     event._message = _message
     event.original_message = _message
     return event
@@ -168,7 +173,6 @@ def fake_message_event_guild(**field) -> "MessageCreateEvent":
     _fake = create_model("_fake", __base__=MessageCreateEvent)
 
     class FakeEvent(_fake):
-        id: str = "1234"
         channel_id: str = "abcd"
         guild_id: str = "efgh"
         content: str = "test"
@@ -178,7 +182,7 @@ class FakeEvent(_fake):
         class Config:
             extra = "forbid"
 
-    return FakeEvent(**field)
+    return FakeEvent(id=str(get_msg_id() + 5555), **field)
 
 
 def fake_satori_bot_params(self_id: str = "test", platform: str = "test") -> dict:

tests/test_buttons.py

@@ -110,7 +110,7 @@ async def _(row: int):
 
         adapter = get_adapter(Adapter)
         bot = ctx.create_bot(base=Bot, adapter=adapter, bot_info=None)
-        event = fake_message_event_guild(message=Message("test 3"), id="123")
+        event = fake_message_event_guild(message=Message("test 3"))
         ctx.receive_event(bot, event)
         ctx.should_call_send(
             event,