feat:初版

plugins/PlayletSearch/__init__.py

@@ -0,0 +1 @@
+# 短剧搜索插件

plugins/PlayletSearch/main.py

@@ -0,0 +1,471 @@
+"""
+短剧搜索插件
+
+用户发送 /搜索短剧 xxx 来搜索短剧并获取视频链接
+"""
+
+import tomllib
+import asyncio
+import aiohttp
+from pathlib import Path
+from loguru import logger
+from typing import List, Dict, Optional
+from utils.plugin_base import PluginBase
+from utils.decorators import on_text_message
+
+
+class PlayletSearch(PluginBase):
+    """短剧搜索插件"""
+
+    description = "搜索短剧并获取视频链接"
+    author = "Assistant"
+    version = "1.0.0"
+
+    def __init__(self):
+        super().__init__()
+        self.config = None
+
+    async def async_init(self):
+        """异步初始化"""
+        config_path = Path(__file__).parent / "config.toml"
+        with open(config_path, "rb") as f:
+            self.config = tomllib.load(f)
+        logger.success("短剧搜索插件已加载")
+
+    @on_text_message(priority=70)
+    async def handle_message(self, bot, message: dict):
+        """处理文本消息"""
+        content = message.get("Content", "").strip()
+        from_wxid = message.get("FromWxid", "")
+        is_group = message.get("IsGroup", False)
+
+        # 精确匹配指令
+        if not content.startswith("/搜索短剧 "):
+            return True
+
+        # 检查是否启用
+        if not self.config["behavior"]["enabled"]:
+            return True
+
+        # 检查群聊过滤
+        if is_group:
+            enabled_groups = self.config["behavior"]["enabled_groups"]
+            disabled_groups = self.config["behavior"]["disabled_groups"]
+
+            if from_wxid in disabled_groups:
+                return True
+            if enabled_groups and from_wxid not in enabled_groups:
+                return True
+
+        # 提取短剧名称
+        keyword = content[6:].strip()  # 去掉 "/搜索短剧 "
+        if not keyword:
+            await bot.send_text(from_wxid, "❌ 请输入短剧名称\n格式：/搜索短剧 短剧名称")
+            return False
+
+        logger.info(f"搜索短剧: {keyword}")
+        await bot.send_text(from_wxid, f"🔍 正在搜索短剧：{keyword}\n请稍候...")
+
+        try:
+            # 第一步：搜索短剧
+            search_result = await self._search_playlet(keyword)
+            if not search_result:
+                await bot.send_text(from_wxid, f"❌ 未找到短剧：{keyword}")
+                return False
+            
+            book_id, cover_url = search_result
+
+            # 第二步：获取剧集列表
+            episode_result = await self._get_episode_list(book_id, keyword)
+            if episode_result is None:
+                await bot.send_text(from_wxid, "❌ 获取剧集列表失败")
+                return False
+            
+            video_list, detail_cover = episode_result
+            if not video_list:
+                await bot.send_text(from_wxid, "❌ 获取剧集列表失败")
+                return False
+            
+            # 优先使用搜索结果的cover，因为detail_cover可能是错误的
+            final_cover_url = cover_url if cover_url else detail_cover
+
+            # 限制集数
+            max_episodes = self.config["behavior"]["max_episodes"]
+            if len(video_list) > max_episodes:
+                video_list = video_list[:max_episodes]
+
+            # 第三步：并发获取所有视频URL
+            video_urls = await self._get_video_urls(video_list)
+
+            # 第四步：构造并发送聊天记录
+            await self._send_chat_records(bot, from_wxid, keyword, video_urls, final_cover_url)
+
+            logger.success(f"短剧搜索完成: {keyword}, {len(video_urls)} 集")
+
+        except Exception as e:
+            logger.error(f"短剧搜索失败: {e}")
+            await bot.send_text(from_wxid, f"❌ 搜索失败: {str(e)}")
+
+        return False
+
+    async def _search_playlet(self, keyword: str) -> Optional[tuple]:
+        """搜索短剧，返回 (book_id, cover_url)"""
+        url = self.config["api"]["base_url"]
+        params = {
+            "key": self.config["api"]["api_key"],
+            "keyword": keyword
+        }
+
+        timeout = aiohttp.ClientTimeout(total=self.config["api"]["timeout"])
+        max_retries = self.config["behavior"].get("max_retries", 15)
+
+        for attempt in range(max_retries):
+            try:
+                async with aiohttp.ClientSession(timeout=timeout) as session:
+                    async with session.post(url, params=params) as resp:
+                        if resp.status != 200:
+                            logger.warning(f"搜索短剧失败 (尝试{attempt+1}/{max_retries}): HTTP {resp.status}")
+                            if attempt < max_retries - 1:
+                                await asyncio.sleep(1)
+                                continue
+                            return None
+
+                        result = await resp.json()
+
+                        if result.get("code") != 0:
+                            logger.warning(f"搜索短剧失败 (尝试{attempt+1}/{max_retries}): {result.get('msg')}")
+                            if attempt < max_retries - 1:
+                                await asyncio.sleep(1)
+                                continue
+                            return None
+
+                        data = result.get("data", [])
+                        if not data:
+                            logger.warning(f"搜索短剧无结果 (尝试{attempt+1}/{max_retries})")
+                            if attempt < max_retries - 1:
+                                await asyncio.sleep(1)
+                                continue
+                            return None
+
+                        # 返回第一个结果的 book_id 和 cover
+                        first_result = data[0]
+                        book_id = first_result.get("book_id")
+                        cover_url = first_result.get("cover", "")
+                        
+                        # URL解码
+                        import urllib.parse
+                        if cover_url:
+                            cover_url = urllib.parse.unquote(cover_url)
+                        
+                        logger.info(f"找到短剧: {first_result.get('title')}, book_id={book_id}, cover_url={cover_url}")
+                        return (book_id, cover_url)
+            except Exception as e:
+                logger.error(f"搜索短剧异常 (尝试{attempt+1}/{max_retries}): {e}")
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(1)
+                else:
+                    return None
+
+        return None
+
+    async def _get_episode_list(self, book_id: str, keyword: str) -> tuple:
+        """获取剧集列表（带重试），返回 (video_list, detail_cover)"""
+        url = self.config["api"]["base_url"]
+        params = {
+            "key": self.config["api"]["api_key"],
+            "book_id": book_id,
+            "keyword": keyword
+        }
+
+        timeout = aiohttp.ClientTimeout(total=self.config["api"]["timeout"])
+        max_retries = self.config["behavior"].get("max_retries", 15)
+
+        for attempt in range(max_retries):
+            try:
+                async with aiohttp.ClientSession(timeout=timeout) as session:
+                    async with session.post(url, params=params) as resp:
+                        if resp.status != 200:
+                            logger.error(f"获取剧集列表失败 (尝试{attempt+1}/{max_retries}): HTTP {resp.status}")
+                            if attempt < max_retries - 1:
+                                await asyncio.sleep(1)
+                                continue
+                            return ([], "")
+
+                        result = await resp.json()
+
+                        if result.get("code") != 0:
+                            logger.error(f"获取剧集列表失败: code={result.get('code')}, msg={result.get('msg')}")
+                            return ([], "")
+
+                        data = result.get("data", {})
+                        video_list = data.get("video_list") or []
+                        detail = data.get("detail", {})
+                        detail_cover = detail.get("cover", "")
+                        
+                        # URL解码
+                        import urllib.parse
+                        if detail_cover:
+                            detail_cover = urllib.parse.unquote(detail_cover)
+
+                        if not video_list:
+                            logger.warning(f"剧集列表为空 (尝试{attempt+1}/{max_retries})")
+                            if attempt < max_retries - 1:
+                                await asyncio.sleep(1)
+                                continue
+                            return ([], "")
+
+                        logger.info(f"获取到 {len(video_list)} 集, detail_cover={detail_cover}")
+                        return (video_list, detail_cover)
+            except Exception as e:
+                logger.error(f"获取剧集列表异常 (尝试{attempt+1}/{max_retries}): {e}")
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(1)
+                else:
+                    return ([], "")
+
+        return ([], "")
+
+    async def _get_video_url(self, session: aiohttp.ClientSession, video_id: str) -> Optional[str]:
+        """获取单个视频URL（带重试）"""
+        url = self.config["api"]["base_url"]
+        params = {
+            "key": self.config["api"]["api_key"],
+            "video_id": video_id
+        }
+
+        max_retries = self.config["behavior"].get("max_retries", 15)
+        for attempt in range(max_retries):
+            try:
+                async with session.post(url, params=params) as resp:
+                    if resp.status != 200:
+                        logger.warning(f"获取视频URL失败 (video_id={video_id}, 尝试{attempt+1}/{max_retries}): HTTP {resp.status}")
+                        if attempt < max_retries - 1:
+                            await asyncio.sleep(1)
+                            continue
+                        return None
+
+                    result = await resp.json()
+
+                    if result.get("code") != 0:
+                        logger.warning(f"获取视频URL失败 (video_id={video_id}, 尝试{attempt+1}/{max_retries}): code={result.get('code')}, msg={result.get('msg')}")
+                        if attempt < max_retries - 1:
+                            await asyncio.sleep(1)
+                            continue
+                        return None
+
+                    data = result.get("data", {})
+                    video = data.get("video", {})
+                    video_url = video.get("url")
+
+                    if not video_url:
+                        logger.warning(f"获取视频URL失败 (video_id={video_id}, 尝试{attempt+1}/{max_retries}): 返回数据中没有url字段")
+                        if attempt < max_retries - 1:
+                            await asyncio.sleep(1)
+                            continue
+                        return None
+
+                    return video_url
+
+            except Exception as e:
+                logger.error(f"获取视频URL异常 (video_id={video_id}, 尝试{attempt+1}/{max_retries}): {e}")
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(1)
+                else:
+                    return None
+
+        return None
+
+    async def _get_video_urls(self, video_list: List[Dict]) -> List[Dict]:
+        """并发获取所有视频URL"""
+        timeout = aiohttp.ClientTimeout(total=self.config["api"]["timeout"])
+        max_concurrent = self.config["behavior"].get("max_concurrent_videos", 10)
+        
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            semaphore = asyncio.Semaphore(max_concurrent)
+            
+            async def get_single_video(video):
+                async with semaphore:
+                    video_id = video.get("video_id")
+                    title = video.get("title")
+                    
+                    url = await self._get_video_url(session, video_id)
+                    if url:
+                        return {"title": title, "url": url}
+                    else:
+                        logger.warning(f"获取视频URL失败: {title} (video_id={video_id})")
+                        return None
+            
+            # 并发执行所有任务
+            tasks = [get_single_video(video) for video in video_list]
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # 过滤掉失败的结果
+            valid_results = []
+            for result in results:
+                if isinstance(result, dict) and result:
+                    valid_results.append(result)
+                elif isinstance(result, Exception):
+                    logger.error(f"获取视频URL异常: {result}")
+            
+            return valid_results
+
+    async def _send_chat_records(self, bot, from_wxid: str, playlet_name: str, video_urls: List[Dict], cover_url: str = ""):
+        if not video_urls:
+            await bot.send_text(from_wxid, "❌ 未获取到任何视频链接")
+            return
+
+        import uuid
+        import time
+        import hashlib
+        import xml.etree.ElementTree as ET
+
+        is_group = from_wxid.endswith("@chatroom")
+        recordinfo = ET.Element("recordinfo")
+        info_el = ET.SubElement(recordinfo, "info")
+        info_el.text = f"{playlet_name} 链接合集"
+        is_group_el = ET.SubElement(recordinfo, "isChatRoom")
+        is_group_el.text = "1" if is_group else "0"
+        datalist = ET.SubElement(recordinfo, "datalist")
+        datalist.set("count", str(len(video_urls)))
+        desc_el = ET.SubElement(recordinfo, "desc")
+        desc_el.text = f"{playlet_name} 链接合集"
+        fromscene_el = ET.SubElement(recordinfo, "fromscene")
+        fromscene_el.text = "3"
+
+        for item in video_urls:
+            di = ET.SubElement(datalist, "dataitem")
+            di.set("datatype", "5")
+            di.set("dataid", uuid.uuid4().hex)
+            src_local_id = str((int(time.time() * 1000) % 90000) + 10000)
+            new_msg_id = str(int(time.time() * 1000))
+            create_time = str(int(time.time()))
+            ET.SubElement(di, "srcMsgLocalid").text = src_local_id
+            ET.SubElement(di, "sourcetime").text = time.strftime("%Y-%m-%d %H:%M", time.localtime(int(create_time)))
+            ET.SubElement(di, "fromnewmsgid").text = new_msg_id
+            ET.SubElement(di, "srcMsgCreateTime").text = create_time
+            ET.SubElement(di, "sourcename").text = playlet_name
+            ET.SubElement(di, "sourceheadurl").text = cover_url or ""
+            ET.SubElement(di, "datatitle").text = item.get("title") or ""
+            ET.SubElement(di, "datadesc").text = "点击观看"
+            ET.SubElement(di, "datafmt").text = "url"
+            ET.SubElement(di, "link").text = item.get("url") or ""
+            ET.SubElement(di, "ischatroom").text = "1" if is_group else "0"
+            weburlitem = ET.SubElement(di, "weburlitem")
+            ET.SubElement(weburlitem, "thumburl").text = cover_url or ""
+            ET.SubElement(di, "thumbwidth").text = "200"
+            ET.SubElement(di, "thumbheight").text = "200"
+            ET.SubElement(weburlitem, "title").text = item.get("title") or ""
+            ET.SubElement(weburlitem, "link").text = item.get("url") or ""
+            ET.SubElement(weburlitem, "desc").text = "点击观看"
+            appmsgshareitem = ET.SubElement(weburlitem, "appmsgshareitem")
+            ET.SubElement(appmsgshareitem, "itemshowtype").text = "-1"
+            dataitemsource = ET.SubElement(di, "dataitemsource")
+            ET.SubElement(dataitemsource, "hashusername").text = hashlib.sha256(from_wxid.encode("utf-8")).hexdigest()
+
+        record_xml = ET.tostring(recordinfo, encoding="unicode")
+
+        appmsg_parts = [
+            "<appmsg appid=\"\" sdkver=\"0\">",
+            f"<title>{playlet_name} 链接合集</title>",
+            f"<des>{playlet_name}</des>",
+            "<type>19</type>",
+            "<url>https://support.weixin.qq.com/cgi-bin/mmsupport-bin/readtemplate?t=page/favorite_record__w_unsupport</url>",
+            "<appattach><cdnthumbaeskey></cdnthumbaeskey><aeskey></aeskey></appattach>",
+            f"<recorditem><![CDATA[{record_xml}]]></recorditem>",
+            "<percent>0</percent>",
+            "</appmsg>"
+        ]
+        appmsg_xml = "".join(appmsg_parts)
+
+        await bot._send_data_async(11214, {"to_wxid": from_wxid, "content": appmsg_xml})
+        logger.success(f"已发送聊天记录，包含 {len(video_urls)} 集视频链接")
+
+    def get_llm_tools(self) -> List[dict]:
+        """返回LLM工具定义，供AIChat插件调用"""
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search_playlet",
+                    "description": "搜索短剧并获取视频链接",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "keyword": {
+                                "type": "string",
+                                "description": "短剧名称或关键词"
+                            }
+                        },
+                        "required": ["keyword"]
+                    }
+                }
+            }
+        ]
+
+    async def execute_llm_tool(self, tool_name: str, arguments: dict, bot, from_wxid: str) -> dict:
+        """执行LLM工具调用，供AIChat插件调用"""
+        try:
+            if not self.config["behavior"]["enabled"]:
+                return {"success": False, "message": "短剧搜索插件未启用"}
+
+            # 检查群聊过滤
+            is_group = from_wxid.endswith("@chatroom")
+            if is_group:
+                enabled_groups = self.config["behavior"]["enabled_groups"]
+                disabled_groups = self.config["behavior"]["disabled_groups"]
+
+                if from_wxid in disabled_groups:
+                    return {"success": False, "message": "此群聊未启用短剧搜索功能"}
+                if enabled_groups and from_wxid not in enabled_groups:
+                    return {"success": False, "message": "此群聊未启用短剧搜索功能"}
+
+            if tool_name == "search_playlet":
+                keyword = arguments.get("keyword")
+                if not keyword:
+                    return {"success": False, "message": "缺少短剧名称参数"}
+
+                logger.info(f"LLM工具调用搜索短剧: {keyword}")
+                await bot.send_text(from_wxid, f"🔍 正在搜索短剧：{keyword}\n请稍候...")
+
+                # 第一步：搜索短剧
+                search_result = await self._search_playlet(keyword)
+                if not search_result:
+                    await bot.send_text(from_wxid, f"❌ 未找到短剧：{keyword}")
+                    return {"success": False, "message": f"未找到短剧：{keyword}"}
+                
+                book_id, cover_url = search_result
+
+                # 第二步：获取剧集列表
+                episode_result = await self._get_episode_list(book_id, keyword)
+                if episode_result is None:
+                    await bot.send_text(from_wxid, "❌ 获取剧集列表失败")
+                    return {"success": False, "message": "获取剧集列表失败"}
+                
+                video_list, detail_cover = episode_result
+                if not video_list:
+                    await bot.send_text(from_wxid, "❌ 获取剧集列表失败")
+                    return {"success": False, "message": "获取剧集列表失败"}
+                
+                # 优先使用搜索结果的cover
+                final_cover_url = cover_url if cover_url else detail_cover
+
+                # 限制集数
+                max_episodes = self.config["behavior"]["max_episodes"]
+                if len(video_list) > max_episodes:
+                    video_list = video_list[:max_episodes]
+
+                # 第三步：并发获取所有视频URL
+                video_urls = await self._get_video_urls(video_list)
+
+                # 第四步：构造并发送聊天记录
+                await self._send_chat_records(bot, from_wxid, keyword, video_urls, final_cover_url)
+
+                logger.success(f"短剧搜索完成: {keyword}, {len(video_urls)} 集")
+                return {"success": True, "message": f"短剧搜索完成：{keyword}，共{len(video_urls)}集"}
+            else:
+                return None  # 不是本插件的工具，返回None让其他插件处理
+                
+        except Exception as e:
+            logger.error(f"LLM工具执行失败: {e}")
+            await bot.send_text(from_wxid, f"❌ 搜索失败: {str(e)}")
+            return {"success": False, "message": f"执行失败: {str(e)}"}