DanDingNoneBot/danding_bot/plugins/chatai/__init__.py

import asyncio
import random
import os
import sys
import subprocess
from nonebot import on_message, get_plugin_config, get_driver, logger
from nonebot.adapters.onebot.v11 import MessageEvent, Bot, MessageSegment
from nonebot.plugin import PluginMetadata
from nonebot.exception import FinishedException
from openai import OpenAI
from .config import Config
from .screenshot import markdown_to_image
import pyppeteer

# 插件元信息
__plugin_meta__ = PluginMetadata(
    name="chatai",
    description="一个对接 DeepSeek 的聊天 AI 插件",
    usage="发送以 * 开头的消息，AI 会回复你，两分钟后自动撤销",
    config=Config,
)

# 获取插件配置
plugin_config = get_plugin_config(Config)

# 全局浏览器实例
_browser: pyppeteer.browser.Browser | None = None
_browser_lock = asyncio.Lock()

# OpenAI 客户端（延迟初始化）
_ai_client: OpenAI | None = None

# 撤回任务引用集合，防止被GC回收
_recall_tasks: set[asyncio.Task] = set()

# 注册消息事件处理器
message_handler = on_message(priority=50, block=True)

# 确保输出目录存在
os.makedirs("data/chatai", exist_ok=True)

# 获取 NoneBot 驱动器
driver = get_driver()


def _force_kill_chrome():
    """强制终止残留的 headless Chrome 进程（仅 pyppeteer 创建的）"""
    try:
        if sys.platform == "win32":
            # 只杀带 --headless 参数的 chrome（避免误杀用户浏览器）
            subprocess.run(
                ["taskkill", "/F", "/IM", "chrome.exe", "/FI", "MODULES eq *pyppeteer*"],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            )
        else:
            subprocess.run(
                ["pkill", "-9", "-f", "chrome.*--remote-debugging-port"],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            )
    except Exception:
        pass


@driver.on_startup
async def startup_cleanup():
    """启动时清理残留Chrome进程"""
    _force_kill_chrome()


@driver.on_shutdown
async def close_browser():
    """在 NoneBot 关闭时关闭浏览器"""
    global _browser
    async with _browser_lock:
        if _browser is not None:
            try:
                await _browser.close()
            except Exception as e:
                logger.warning(f"关闭浏览器异常: {e}")
            _browser = None
    _force_kill_chrome()


async def init_browser() -> "pyppeteer.browser.Browser":
    """初始化或复用浏览器实例"""
    global _browser
    async with _browser_lock:
        if _browser is None or not _browser.process:
            try:
                _browser = await pyppeteer.launch(
                    headless=True,
                    args=["--no-sandbox", "--disable-setuid-sandbox", "--disable-gpu"],
                )
                logger.info("chatai: 浏览器实例已创建")
            except Exception as e:
                logger.error(f"chatai: 浏览器启动失败: {e}")
                raise
    return _browser


_AI_API_TIMEOUT = 60  # seconds
_AI_SYSTEM_PROMPT = (
    "你是一个活泼可爱的群助手。请在回复中使用丰富的 Emoji 表情"
    "(如 😊 😄 🎉 ✨ 💖 等)。你的语气要俏皮活泼，像在和朋友聊天一样自然。"
    "在回答问题时要保持专业性的同时，也要让回复显得生动有趣。"
    "每条回复都必须包含至少2-3个 Emoji 表情。"
    "如果你需要展示代码片段，请确保代码中不包含任何颜文字或表情符号，"
    "保持代码的专业性和可读性。"
)


def _get_ai_client() -> OpenAI:
    """获取或创建 OpenAI 客户端（单例）"""
    global _ai_client
    if _ai_client is None:
        _ai_client = OpenAI(
            api_key=plugin_config.deepseek_token,
            base_url="https://api.siliconflow.cn/v1",
            timeout=_AI_API_TIMEOUT,
        )
    return _ai_client


async def call_ai_api(message: str) -> str:
    """调用 AI 接口"""
    client = _get_ai_client()
    response = await asyncio.get_event_loop().run_in_executor(
        None,
        lambda: client.chat.completions.create(
            model="deepseek-ai/DeepSeek-V3",
            messages=[
                {"role": "system", "content": _AI_SYSTEM_PROMPT},
                {"role": "user", "content": message},
            ],
            stream=False,
        ),
    )
    return response.choices[0].message.content or ""


@message_handler.handle()
async def handle_message(event: MessageEvent, bot: Bot):
    user_message = event.get_plaintext().strip()

    if not user_message.startswith("*"):
        return

    user_message = user_message[1:].strip()

    if not user_message:
        await asyncio.sleep(random.uniform(2, 3))
        await message_handler.finish("请输入有效内容哦~")

    try:
        browser = await init_browser()
        response = await call_ai_api(user_message)

        if response:
            await asyncio.sleep(random.uniform(2, 3))
            # 使用事件ID+时间戳避免并发路径冲突
            image_path = f"data/chatai/output_{event.message_id}.png"
            await markdown_to_image(response, image_path, browser)

            sent_message = await bot.send(
                event, MessageSegment.image(f"file:///{os.path.abspath(image_path)}")
            )

            # 清理临时图片文件
            try:
                os.remove(image_path)
            except OSError:
                pass

            # 保存task引用防止GC回收
            task = asyncio.create_task(
                _delete_message_after_delay(bot, sent_message["message_id"])
            )
            _recall_tasks.add(task)
            task.add_done_callback(_recall_tasks.discard)

    except FinishedException:
        raise
    except Exception as e:
        logger.error(f"chatai处理失败: user_id={event.user_id} error={e}")
        await asyncio.sleep(random.uniform(2, 3))
        logger.error(f"chatai详细错误: {e}")
        await message_handler.finish("出错了，请稍后再试~")


async def _delete_message_after_delay(bot: Bot, message_id: int):
    """两分钟后撤回消息"""
    await asyncio.sleep(120)
    try:
        await bot.delete_msg(message_id=message_id)
    except Exception as e:
        logger.debug(f"chatai撤回消息失败(可忽略): msg_id={message_id} error={e}")