更新设计

Author: JohnRichard4096

nonebot_plugin_omikuji/llm_tool.py

@@ -1,4 +1,3 @@
-import random
 import typing
 
 from nonebot import get_bot, logger
@@ -12,6 +11,8 @@
     ToolFunctionSchema,
 )
 
+from nonebot_plugin_omikuji.cache import cache_omikuji, get_cached_omikuji
+
 from .config import get_config
 from .models import OMIKUJI_THEMES, OmikujiData
 from .utils import generate_omikuji
@@ -40,17 +41,16 @@ def format_omikuji(data: OmikujiData, user_name: str | None = ""):
 async def omikuji(ctx: ToolContext):
     logger.info("获取御神签")
     nb_event: MessageEvent = typing.cast(MessageEvent, ctx.event.get_nonebot_event())
+    is_group = hasattr(nb_event, "group_id")
     bot = get_bot(str(ctx.event._nbevent.self_id))
-    await bot.send(
-        ctx.event._nbevent,
-        "轻轻摇动古老的签筒，竹签哗啦作响... 心中默念所求之事... 一支签缓缓落下。",
-    )
-    level = random.choice(LEVEL)
-    data = await generate_omikuji(
-        ctx.data["theme"],
-        level=level,
-        is_group=hasattr(nb_event, "group_id"),
-    )
+
+    if (data := await get_cached_omikuji(nb_event)) is None:
+        await bot.send(
+            ctx.event._nbevent,
+            "轻轻摇动古老的签筒，竹签哗啦作响... 心中默念所求之事... 一支签缓缓落下。",
+        )
+        data = await generate_omikuji(ctx.data["theme"], is_group)
+        await cache_omikuji(nb_event, data)
     if get_config().omikuji_send_by_chat:
         return data.model_dump_json()
     msg = format_omikuji(data)

nonebot_plugin_omikuji/models.py

@@ -12,8 +12,6 @@
 from pydantic import BaseModel, Field
 
 LEVEL = ["大吉", "吉", "中吉", "小吉", "末吉", "凶", "大凶"]
-# 定义权重，让"吉"类签出现概率更高
-LEVEL_WEIGHTS = [10, 30, 20, 15, 10, 5, 2]
 OMIKUJI_THEMES = [
     "综合运势",  # 默认选项
     "恋爱姻缘",  # 情感
@@ -41,7 +39,7 @@
 
 
 def random_level() -> str:
-    return random.choices(LEVEL, weights=LEVEL_WEIGHTS)[0]
+    return random.choice(LEVEL)
 
 
 class OmikujiSections(BaseModel):
@@ -76,7 +74,7 @@ class OmikujiData(BaseModel):
                 ),
                 "sign_number": FunctionPropertySchema(
                     type="string",
-                    description="御神签编号(中文大写数字)",
+                    description="御神签编号(随机中文大写数字)",
                 ),
                 "divine_title": FunctionPropertySchema(
                     type="string",

nonebot_plugin_omikuji/utils.py

@@ -80,6 +80,8 @@ async def generate_omikuji(
     )
     assert data.tool_calls
     args = json.loads(data.tool_calls[0].function.arguments)
+    args["level"] = level
+    args["theme"] = theme
     model = OmikujiData.model_validate(args)
     if level:
         model.level = level