chatgpt-mirai-qq-bot-web-search 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

{chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info → chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: chatgpt-mirai-qq-bot-web-search
-Version: 0.2.2
+Version: 0.2.4
 Summary: WebSearch adapter for lss233/chatgpt-mirai-qq-bot
 Home-page: https://github.com/chuanSir123/web_search
 Author: chuanSir
@@ -18,7 +18,7 @@ Requires-Dist: playwright
 Requires-Dist: trafilatura
 Requires-Dist: lxml-html-clean
 
-# OneBot-adapter for ChatGPT-Mirai-QQ-Bot
+# web_search for ChatGPT-Mirai-QQ-Bot
 
 本项目是 [ChatGPT-Mirai-QQ-Bot](https://github.com/lss233/chatgpt-mirai-qq-bot) 的一个插件，用于将OneBot协议的消息转换为ChatGPT-Mirai-QQ-Bot的消息格式。
 

chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/RECORD

@@ -0,0 +1,11 @@
+web_search/__init__.py,sha256=bpu1gpf9tq3mOAoaGMM-8S6GBw2GzRAG6ClHGlp-CVw,4607
+web_search/blocks.py,sha256=NJqWOAp7X-WMM3c3Xn3qLb9IVV090W3MqELHZTAw8tk,8959
+web_search/config.py,sha256=DhLiERBJR2V5Boglf7Aq9Rbc4vsvLIh67CrLDIPeqA0,398
+web_search/web_searcher.py,sha256=760WTBGusxYntigGi5Wsbm2rRHEWq3sTwfNlf3ew3b0,20919
+web_search/example/roleplayWithWebSearch.yaml,sha256=C-dGy3z8gcRcmxzurssP-kPRLqMf1TYR-nnNUaJjISE,7468
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/LICENSE,sha256=ILBn-G3jdarm2w8oOrLmXeJNU3czuJvVhDLBASWdhM8,34522
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/METADATA,sha256=RaMgn6mSy-K2tEAvQMzmwDwO1rttpBHfrWqGP5VS9Yo,1734
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/entry_points.txt,sha256=o3kRDSdSmSdnCKlK6qS57aN0WpI4ab-Nxub2NwUrjf0,64
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/top_level.txt,sha256=PoNm8MJYw_y8RTMaNlY0ePLoNHxVUAE2IHDuL5fFubI,11
+chatgpt_mirai_qq_bot_web_search-0.2.4.dist-info/RECORD,,

web_search/__init__.py

@@ -6,7 +6,7 @@ from .config import WebSearchConfig
 from .web_searcher import WebSearcher
 from dataclasses import dataclass
 from kirara_ai.workflow.core.block import BlockRegistry
-from .blocks import WebSearchBlock,WebSearchByKeywordBlock
+from .blocks import WebSearchBlock,WebSearchByKeywordBlock, DouyinVideoSearchBlock
 from .blocks import AppendSystemPromptBlock
 from kirara_ai.ioc.inject import Inject
 from kirara_ai.ioc.container import DependencyContainer
@@ -31,6 +31,7 @@ class WebSearchPlugin(Plugin):
         try:
             self.block_registry.register("web_search", "search", WebSearchBlock)
             self.block_registry.register("web_search_by_keyword", "search", WebSearchByKeywordBlock)
+            self.block_registry.register("douyin_video_search", "search", DouyinVideoSearchBlock)
         except Exception as e:
             logger.warning(f"WebSearchPlugin failed: {e}")
         try:

web_search/blocks.py

@@ -6,6 +6,8 @@ from .config import WebSearchConfig
 from kirara_ai.llm.format.message import LLMChatMessage
 from kirara_ai.llm.format.response import LLMChatResponse
 from kirara_ai.ioc.container import DependencyContainer
+import re
+from kirara_ai.im.message import IMMessage
 
 def get_options_provider(container: DependencyContainer, block: Block) -> List[str]:
     return ["bing", "google", "baidu"]
@@ -166,3 +168,64 @@ class AppendSystemPromptBlock(Block):
 
         return {"messages": messages}
 
+class DouyinVideoSearchBlock(Block):
+    """抖音视频搜索Block"""
+    name = "douyin_video_search"
+    description = "通过关键词搜索抖音视频"
+    container: DependencyContainer
+    inputs = {
+        "keyword": Input(name="keyword", label="搜索关键字", data_type=str, description="搜索关键词"),
+        "count": Input(name="count", label="视频数量", data_type=int, description="需要获取的视频数量")
+    }
+
+    outputs = {
+        "results": Output(name="results", label="搜索结果", data_type=str, description="视频链接列表")
+    }
+
+    def __init__(self, name: str = None, timeout: Optional[int] = 10, proxy: str = None):
+        super().__init__(name)
+        self.searcher = None
+        self.config = WebSearchConfig()
+        self.timeout = timeout
+        self.proxy = proxy
+
+    def _ensure_searcher(self):
+        """同步方式初始化searcher"""
+        if not self.searcher:
+            try:
+                loop = asyncio.get_event_loop()
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            self.searcher = loop.run_until_complete(WebSearcher.create())
+
+    def execute(self, **kwargs) -> Dict[str, Any]:
+        keyword = kwargs["keyword"]
+        count = kwargs["count"]
+
+        if not keyword:
+            return {"results": ""}
+
+        self._ensure_searcher()
+
+        try:
+            try:
+                loop = asyncio.get_event_loop()
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+
+            results = loop.run_until_complete(
+                self.searcher.search_douyin_videos(
+                    keyword=keyword,
+                    count=count,
+                    timeout=self.timeout,
+                    proxy=self.proxy,
+                    sender =  self.container.resolve(IMMessage).sender.user_id
+                )
+            )
+            return {"results": f"\n以下是抖音视频搜索结果:\n{results}"}
+        except Exception as e:
+            print(e)
+            return {"results": f"搜索失败: {str(e)}"}
+

web_search/web_searcher.py

@@ -8,14 +8,27 @@ import subprocess
 import sys
 from kirara_ai.logger import get_logger
 import os
+import re
+import requests
+import json
+from kirara_ai.im.message import IMMessage
+from kirara_ai.im.sender import ChatSender
+import yaml
+from datetime import datetime, date
 
 logger = get_logger("WebSearchPlugin")
-
+user_videoIds = {}
 class WebSearcher:
+
     def __init__(self):
+
         self.playwright = None
         self.browser = None
         self.context = None
+        # 获取当前文件所在目录
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        self.video_ids_file = os.path.join(current_dir, "douyin_video_ids.yaml")
+        self.video_ids = self._load_video_ids()
         self.search_engines = {
             'bing': {
                 'url': 'https://www.bing.com/search?q={}',
@@ -49,10 +62,12 @@ class WebSearcher:
     async def _ensure_initialized(self,proxy):
         """确保浏览器已初始化"""
         try:
+            if self.context:
+                return self.context
             self.playwright = await async_playwright().start()
 
             # 创建用户数据目录路径
-            user_data_dir = os.path.join(os.path.expanduser("~"), ".playwright_user_data")
+            user_data_dir = os.path.join(os.path.expanduser("~"), ".playwright_user_data")+f'{random.randint(1, 1000000)}'
             os.makedirs(user_data_dir, exist_ok=True)
 
             # 合并所有选项到一个字典
@@ -323,3 +338,143 @@ class WebSearcher:
             await self.context.close()
         if self.playwright:
             await self.playwright.stop()
+
+    def _load_video_ids(self):
+        """从YAML文件加载视频ID记录"""
+        try:
+            today = str(date.today())
+            if os.path.exists(self.video_ids_file):
+                with open(self.video_ids_file, 'r', encoding='utf-8') as f:
+                    data = yaml.safe_load(f) or {}
+                    # 检查是否是今天的数据
+                    if data.get('date') == today:
+                        return data.get('video_ids', {})
+
+            # 如果文件不存在、数据为空或日期不是今天，创建新的空记录
+            empty_data = {
+                'date': today,
+                'video_ids': {}
+            }
+            with open(self.video_ids_file, 'w', encoding='utf-8') as f:
+                yaml.safe_dump(empty_data, f, allow_unicode=True)
+            return empty_data['video_ids']
+        except Exception as e:
+            logger.error(f"Failed to load video IDs: {e}")
+            return {}
+
+    def _save_video_ids(self):
+        """保存视频ID记录到YAML文件"""
+        try:
+            data = {
+                'date': str(date.today()),
+                'video_ids': self.video_ids
+            }
+            # 确保目录存在
+            os.makedirs(os.path.dirname(self.video_ids_file), exist_ok=True)
+            # 使用 'w' 模式覆盖写入文件
+            with open(self.video_ids_file, 'w', encoding='utf-8') as f:
+                yaml.safe_dump(data, f, allow_unicode=True)
+        except Exception as e:
+            logger.error(f"Failed to save video IDs: {e}")
+
+    async def search_douyin_videos(self, keyword: str, count: int = 5, timeout: int = 10, proxy: str = None, sender: str = None) -> str:
+        """搜索抖音视频"""
+        context = await self._ensure_initialized(proxy)
+        page = None
+
+        if not sender:
+            return "需要提供发送者ID"
+
+        if sender not in self.video_ids:
+            self.video_ids[sender] = []
+
+        try:
+            # 构建搜索URL
+            search_url = f'https://www.douyin.com/search/{urllib.parse.quote(keyword)}'
+            page = await context.new_page()
+
+            # 设置User-Agent
+            await page.set_extra_http_headers({
+                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+            })
+
+            # 访问搜索页面并等待包含waterfall_item_的内容出现
+            await page.goto(search_url, wait_until='domcontentloaded')
+
+            # 等待页面中出现waterfall_item_字符串
+            await page.wait_for_function('''
+                () => document.documentElement.innerHTML.includes('waterfall_item_')
+            ''', timeout=timeout * 1000)
+
+            # 获取ttwid
+            url = "https://ttwid.bytedance.com/ttwid/union/register/"
+            ttjson = {"region": "cn", "aid": 1768, "needFid": "false", "service": "www.ixigua.com",
+                        "migrate_info": {"ticket": "", "source": "node"}, "cbUrlProtocol": "https", "union": "true"}
+            ttresponse = requests.post(url, json=ttjson)
+            tt = ttresponse.cookies.get_dict()['ttwid']
+
+            video_links = []
+            videoCount = 0
+            max_scroll_attempts = 10  # 最大滚动尝试次数
+            scroll_attempt = 0
+            while videoCount < count and scroll_attempt < max_scroll_attempts:
+                # 提取视频ID
+                video_elements = await page.query_selector_all('div[id^="waterfall_item_"]')
+
+                # 处理当前页面上的视频
+                for i, element in enumerate(video_elements):
+                    if videoCount >= count:
+                        break
+
+                    # 获取元素的id属性
+                    div_id = await element.get_attribute('id')
+                    video_id = div_id.replace('waterfall_item_', '')
+                    if video_id not in self.video_ids[sender]:
+                        if video_id.isdigit():  # 确保是数字ID
+                            url = f'https://www.douyin.com/video/{video_id}'
+                            try:
+                                headers = {
+                                    "referer": "https://www.douyin.com/",
+                                    "user-agent": "Mozilla/5.0 (Linux; Android 12; 2210132C Build/SP1A.210812.016) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.104 Mobile Safari/537.36",
+                                    "cookie": "ttwid=" + tt
+                                }
+                                video_url = url
+                                aweme_id = re.findall('video/(\d+)', url)[0]
+                                url1 = f"https://www.iesdouyin.com/share/video/{aweme_id}"
+                                resp1 = requests.get(url1, headers=headers).text.encode('gbk', errors='ignore').decode('gbk')
+                                json_data = resp1.split("window._ROUTER_DATA = ")[1].split("</script>")[0]
+                                resp1 = json.loads(json_data.encode('gbk', errors='ignore').decode('gbk'))
+                                video_url = resp1["loaderData"]["video_(id)/page"]["videoInfoRes"]["item_list"][0]["video"]["play_addr"]["url_list"][0]
+                                video_url = video_url.replace("playwm", "play").replace("720p", "1080p")
+                                imgresponse = requests.get(video_url, allow_redirects=False)
+                                new_url = imgresponse.headers.get('Location')
+                                if new_url:  # 只有获取到重定向地址才添加
+                                    logger.debug(new_url)
+                                    video_links.append(f'[{videoCount+1}] {new_url}')
+                                    videoCount += 1
+                                    self.video_ids[sender].append(video_id)
+                                    self._save_video_ids()  # 每添加一个视频就保存一次
+                            except Exception as e:
+                                continue
+
+                # 如果还没有获取够视频，继续滚动加载
+                if videoCount < count:
+                    # 滚动到页面底部
+                    await page.evaluate('window.scrollTo(0, document.body.scrollHeight)')
+                    await asyncio.sleep(2)  # 等待新内容加载
+                    scroll_attempt += 1
+                if scroll_attempt == 10:
+                    self.video_ids[sender] = []
+                    self._save_video_ids()
+
+            return ("视频url地址:\n"+"\n".join(video_links)) if video_links else "未找到视频"
+
+        except Exception as e:
+            logger.error(f"抖音视频搜索失败 - 关键词: {keyword} - 错误: {e}", exc_info=True)
+            return f"搜索失败: {str(e)}"
+        finally:
+            if page:
+                try:
+                    await page.close()
+                except Exception as e:
+                    logger.error(f"关闭页面错误: {e}")

chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-web_search/__init__.py,sha256=cDyiTyQ4i-7brk3qfQrPBF5XQNpDacWmauzJwTlK6VI,4485
-web_search/blocks.py,sha256=S3RsV9CCTKAsKUNhewg__ejEpJRDz7DTawtH05WRgE8,6732
-web_search/config.py,sha256=DhLiERBJR2V5Boglf7Aq9Rbc4vsvLIh67CrLDIPeqA0,398
-web_search/web_searcher.py,sha256=0zLgMsWCK71gStyWpFjup5WfxHx3tBTf3rGwM7Ae7Zs,13332
-web_search/example/roleplayWithWebSearch.yaml,sha256=C-dGy3z8gcRcmxzurssP-kPRLqMf1TYR-nnNUaJjISE,7468
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/LICENSE,sha256=ILBn-G3jdarm2w8oOrLmXeJNU3czuJvVhDLBASWdhM8,34522
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/METADATA,sha256=Skc-om3-q2TH-RUte06VDlqyOD4QC17hpqaS2UanRCM,1738
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/entry_points.txt,sha256=o3kRDSdSmSdnCKlK6qS57aN0WpI4ab-Nxub2NwUrjf0,64
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/top_level.txt,sha256=PoNm8MJYw_y8RTMaNlY0ePLoNHxVUAE2IHDuL5fFubI,11
-chatgpt_mirai_qq_bot_web_search-0.2.2.dist-info/RECORD,,