mcp-query-table 0.3.3__tar.gz → 0.3.4__tar.gz

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mcp_query_table
-Version: 0.3.3
+Version: 0.3.4
 Summary: query table from website, support MCP
 Author-email: wukan <wu-kan@163.com>
 License: MIT License

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table/__init__.py

@@ -3,4 +3,5 @@ from ._version import __version__
 from .enums import QueryType, Site, Provider
 from .tool import BrowserManager, query, chat
 
-TIMEOUT = 1000 * 60 * 2  # 2分钟，在抓取EventStream数据时等待数据返回，防止外层30秒超时
+TIMEOUT = 1000 * 60 * 3  # 3分钟，在抓取EventStream数据时等待数据返回，防止外层30秒超时
+TIMEOUT_60 = 1000 * 60  # 1分钟

mcp_query_table-0.3.4/mcp_query_table/_version.py

@@ -0,0 +1 @@
+__version__ = "0.3.4"

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table/providers/baidu.py

@@ -5,14 +5,14 @@
 """
 import json
 
-from loguru import logger
 from playwright.async_api import Page
 
 import mcp_query_table
-from mcp_query_table.tool import GlobalVars
+from mcp_query_table.tool import GlobalVars, split_images
 
 _PAGE0_ = "https://chat.baidu.com/search"
 _PAGE1_ = "https://chat.baidu.com/aichat/api/conversation"
+_PAGE2_ = "https://chat.baidu.com/aichat/api/file/upload"
 
 G = GlobalVars()
 
@@ -69,16 +69,30 @@ async def chat(page: Page,
                create: bool,
                files: list[str],
                ) -> str:
+    async def on_file_chooser(file_chooser):
+        # 文件选择对话框
+        await file_chooser.set_files(files)
+
     if not page.url.startswith(_PAGE0_):
         create = True
 
     if create:
         await page.goto(_PAGE0_)
 
+    # 文件上传
     if len(files) > 0:
-        # TODO
-        logger.warning("抱歉，百度AI搜索的上传文件功能未突破")
-
+        imgs, docs = split_images(files)
+        assert len(imgs) == 0 or len(docs) == 0, "不能同时包含图片和文档"
+
+        page.on("filechooser", on_file_chooser)
+        async with page.expect_response(f"{_PAGE2_}*", timeout=mcp_query_table.TIMEOUT_60) as response_info:
+            if len(imgs) > 0:
+                await page.locator(".cs-input-upload-icon").last.click()
+            else:
+                await page.locator(".cs-input-upload-icon").first.click()
+        page.remove_listener("filechooser", on_file_chooser)
+
+    # 提交问题
     await page.route(_PAGE1_, on_route)
     async with page.expect_response(_PAGE1_, timeout=mcp_query_table.TIMEOUT) as response_info:
         await page.locator("#chat-input-box").fill(prompt)

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table/providers/n.py

@@ -93,7 +93,7 @@ async def chat(page: Page,
         await page.goto(_PAGE0_)
         if len(files) > 0:
             # 只能在新会话中上传文件
-            async with page.expect_response(_PAGE3_, timeout=mcp_query_table.TIMEOUT) as response_info:
+            async with page.expect_response(_PAGE3_, timeout=mcp_query_table.TIMEOUT_60) as response_info:
                 await page.locator("input[type=\"file\"]").set_input_files(files)
     else:
         name = "提出后续问题，Enter发送，Shift+Enter 换行"

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table/providers/yuanbao.py

@@ -7,7 +7,7 @@ import re
 from playwright.async_api import Page
 
 import mcp_query_table
-from mcp_query_table.tool import GlobalVars, is_image
+from mcp_query_table.tool import GlobalVars, split_images
 
 _PAGE0_ = "https://yuanbao.tencent.com/"
 _PAGE1_ = "https://yuanbao.tencent.com/api/chat"
@@ -76,20 +76,15 @@ async def chat(page: Page,
         await page.goto(_PAGE0_)
 
     if len(files) > 0:
-        is_img, is_doc = False, False
-        for f in files:
-            if is_image(f):
-                is_img = True
-            else:
-                is_doc = True
-        assert is_img ^ is_doc, "不能同时包含图片和文档"
+        imgs, docs = split_images(files)
+        assert len(imgs) == 0 or len(docs) == 0, "不能同时包含图片和文档"
 
         # 点击上传文件按钮，才会出现上传文件的input
         await page.get_by_role("button").filter(has_text=re.compile(r"^$")).last.click()
 
         # 上传文件
-        async with page.expect_response(_PAGE2_, timeout=mcp_query_table.TIMEOUT) as response_info:
-            if is_img:
+        async with page.expect_response(_PAGE2_, timeout=mcp_query_table.TIMEOUT_60) as response_info:
+            if len(imgs) > 0:
                 await page.locator("input[type=\"file\"]").nth(-2).set_input_files(files)
             else:
                 await page.locator("input[type=\"file\"]").last.set_input_files(files)

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table/tool.py

@@ -2,7 +2,7 @@ import subprocess
 import sys
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Optional, List, Tuple
 
 import pandas as pd
 from loguru import logger
@@ -142,6 +142,12 @@ class BrowserManager:
             # 防止开发者工具被使用
             if page.url.startswith("devtools://"):
                 continue
+            # 防止chrome扩展被使用
+            if page.url.startswith("chrome-extension://"):
+                continue
+            # 防止edge扩展被使用
+            if page.url.startswith("extension://"):
+                continue
             self.pages.append(page)
 
     async def _try_launch(self) -> None:
@@ -275,3 +281,15 @@ def is_image(path: str) -> bool:
     img_ext = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']
     ext = Path(path).suffix.lower()
     return ext in img_ext
+
+
+def split_images(files: List[str]) -> Tuple[List[str], List[str]]:
+    """图片列表分成两部分"""
+    imgs = []
+    docs = []
+    for f in files:
+        if is_image(f):
+            imgs.append(f)
+        else:
+            docs.append(f)
+    return imgs, docs

{mcp_query_table-0.3.3 → mcp_query_table-0.3.4}/mcp_query_table.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mcp_query_table
-Version: 0.3.3
+Version: 0.3.4
 Summary: query table from website, support MCP
 Author-email: wukan <wu-kan@163.com>
 License: MIT License

mcp_query_table-0.3.3/mcp_query_table/_version.py

@@ -1 +0,0 @@
-__version__ = "0.3.3"