移除冗余代码，添加command

Author: JohnRichard4096

nonebot_plugin_omikuji/__init__.py

@@ -10,7 +10,7 @@
 require("nonebot_plugin_orm")
 from importlib import metadata
 
-from nonebot_plugin_suggarchat.API import ToolsManager
+from nonebot_plugin_suggarchat.API import Menu, ToolsManager
 
 from . import llm_tool, sql_models
 from .cache import OmikujiCache
@@ -40,6 +40,7 @@ async def init():
     conf = get_config()
     if conf.enable_omikuji:
         ToolsManager().register_tool(TOOL_DATA)
+        Menu().reg_menu("omikuji", "抽御神签", "[可选]主题")
     logger.info("正在初始化缓存数据......")
     os.makedirs(get_cache_dir(), exist_ok=True)
     for cache in get_cache_dir().glob("*.json"):

nonebot_plugin_omikuji/commands.py

@@ -0,0 +1,38 @@
+import random
+import typing
+
+from nonebot import on_command
+from nonebot.adapters.onebot.v11 import Bot, GroupMessageEvent, Message, MessageEvent
+from nonebot.params import CommandArg
+
+from .cache import cache_omikuji, get_cached_omikuji
+from .config import get_config
+from .models import OMIKUJI_THEMES, THEME_TYPE
+from .utils import format_omikuji, generate_omikuji
+
+omikuji = on_command(
+    "omikuji",
+    aliases={"御神签", "抽签"},
+    priority=10,
+    block=True,
+    rule=lambda: get_config().enable_omikuji,
+)
+
+
+@omikuji.handle()
+async def _(bot: Bot, event: MessageEvent, args: Message = CommandArg()):
+    if args:
+        theme = args.extract_plain_text()
+        if theme not in OMIKUJI_THEMES:
+            await omikuji.finish(
+                f"当前可用御神签主题：{''.join(i + ',' for i in OMIKUJI_THEMES)}"
+            )
+    else:
+        theme = random.choice(list(OMIKUJI_THEMES))
+    theme = typing.cast(THEME_TYPE, theme)
+    is_group = isinstance(event, GroupMessageEvent)
+    if (data := await get_cached_omikuji(event)) is None:
+        data = await generate_omikuji(theme, is_group)
+        await cache_omikuji(event, data)
+    msg = format_omikuji(data)
+    await omikuji.finish(msg)

nonebot_plugin_omikuji/llm_tool.py

@@ -1,16 +1,40 @@
+import random
 import typing
 
 from nonebot import get_bot, logger
 from nonebot.adapters.onebot.v11 import MessageEvent
 from nonebot_plugin_suggarchat.API import (
+    FunctionDefinitionSchema,
+    FunctionParametersSchema,
+    FunctionPropertySchema,
     ToolContext,
     ToolData,
+    ToolFunctionSchema,
 )
 
-from .cache import get_cached_omikuji
 from .config import get_config
-from .models import FUNC_META
-from .utils import format_omikuji, get_omikuji
+from .models import OMIKUJI_THEMES, OmikujiData
+from .utils import generate_omikuji
+
+LEVEL = ["大吉", "吉", "中吉", "小吉", "末吉", "凶", "大凶"]
+
+
+def format_omikuji(data: OmikujiData, user_name: str | None = ""):
+    ln = "\n"
+    msg = f"""{data.intro}
+{(user_name + "，" if user_name else "")}你的签上刻了什么？
+
+＝＝＝ 御神签 第{data.sign_number} ＝＝＝
+✨ 天启：{data.divine_title}
+🌸 运势：{data.level} - {data.theme}
+
+{"".join(f"▫ {section.name}{ln}{section.content}{ln}" for section in data.sections)}
+
+⚖ 真言偈：{data.maxim}
+
+{data.end}
+"""
+    return msg
 
 
 async def omikuji(ctx: ToolContext):
@@ -21,16 +45,36 @@ async def omikuji(ctx: ToolContext):
         ctx.event._nbevent,
         "轻轻摇动古老的签筒，竹签哗啦作响... 心中默念所求之事... 一支签缓缓落下。",
     )
-    if (data := await get_cached_omikuji(nb_event)) is None:
-        data = await get_omikuji(
-            ctx.data["theme"],
-            is_group=hasattr(nb_event, "group_id"),
-        )
+    level = random.choice(LEVEL)
+    data = await generate_omikuji(
+        ctx.data["theme"],
+        level=level,
+        is_group=hasattr(nb_event, "group_id"),
+    )
     if get_config().omikuji_send_by_chat:
         return data.model_dump_json()
     msg = format_omikuji(data)
     await bot.send(nb_event, msg)
     ctx.matcher.cancel_nonebot_process()
 
 
+FUNC_META = ToolFunctionSchema(
+    strict=True,
+    function=FunctionDefinitionSchema(
+        name="御神签(omikuji)",
+        description="抽取一个御神签",
+        parameters=FunctionParametersSchema(
+            type="object",
+            properties={
+                "theme": FunctionPropertySchema(
+                    type="string",
+                    description="御神签主题（如果包含不良内容则随机选择）",
+                    enum=OMIKUJI_THEMES,
+                )
+            },
+            required=["theme"],
+        ),
+    ),
+)
+
 TOOL_DATA = ToolData(data=FUNC_META, func=omikuji, custom_run=True)

nonebot_plugin_omikuji/sql_models.py

@@ -3,7 +3,6 @@
 import asyncio
 from datetime import datetime
 from functools import lru_cache
-from typing import Literal
 
 from nonebot_plugin_orm import Model
 from sqlalchemy import JSON, Index, String, UniqueConstraint
@@ -21,7 +20,7 @@ class OmikujiCache(Model):
     __tablename__ = "omikuji_cache"
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     level: Mapped[str] = mapped_column(String(64), nullable=False)
-    theme: Mapped[Literal[THEME_TYPE]] = mapped_column(String(64), nullable=False)
+    theme: Mapped[THEME_TYPE] = mapped_column(String(64), nullable=False)
     sections: Mapped[dict[str, list[str]]] = mapped_column(
         JSON,
         default={},

nonebot_plugin_omikuji/utils.py

@@ -53,7 +53,9 @@ async def _hit_cache_omikuji(
             end=random.choice([i.content for i in cache.end]),
         )
         return model
-async def get_omikuji(
+
+
+async def generate_omikuji(
     theme: THEME_TYPE,
     is_group: bool = False,
     level: str = "",