mcp-query-table 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

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/_version.py

@@ -1 +1 @@
-__version__ = "0.3.3"
+__version__ = "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/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/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/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.dist-info → mcp_query_table-0.3.4.dist-info}/METADATA

@@ -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.4.dist-info/RECORD

@@ -0,0 +1,19 @@
+mcp_query_table/__init__.py,sha256=NmnAOKcJjQoeTJMbIY79-JsNNdj-PJpnZ-x8MPjpE4o,272
+mcp_query_table/__main__.py,sha256=W3tMnZZTvkhjlLHmSYHng3CTvYX6-LjyaBDhf-ypH7w,1186
+mcp_query_table/_version.py,sha256=oYLGMpySamd16KLiaBTfRyrAS7_oyp-TOEHmzmeumwg,22
+mcp_query_table/enums.py,sha256=Hen0X1f2of69f08epun6HvYIgbpw_rf8BvoRQ184kS4,679
+mcp_query_table/server.py,sha256=oV8CSi-EW5-xlPn38R1wLu7U-pCHgAJci1W8rJFsXSk,3430
+mcp_query_table/tool.py,sha256=kCnfBNv1T88gNnFLsz-ZjXNuxbzhGvLPcPfNkkJ7Dyg,9096
+mcp_query_table/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mcp_query_table/providers/baidu.py,sha256=q0vZxl_UTO8v4LbznwT2b1hM-00jswDR8CqB3tc-X9U,3389
+mcp_query_table/providers/n.py,sha256=J0Xo-BlXveIQdDbwKkuEP4I1mfAr8i1rviD0UTztI9Y,3009
+mcp_query_table/providers/yuanbao.py,sha256=4DzoLfgg7dSP1BhiS_tKyTQ0byc1zHRuZQjyq-2pZfQ,3148
+mcp_query_table/sites/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mcp_query_table/sites/eastmoney.py,sha256=LImjpYVuM5YnXwnNzB2hkKfHofocZZScetGqMOCHZpk,4477
+mcp_query_table/sites/iwencai.py,sha256=g56pj3pbxu4mXLNnaaS3Hdx-DvEy_9OBrQJe26z4z08,5059
+mcp_query_table/sites/tdx.py,sha256=RIUQaB7Tn4AVyWaevk9SzTKIDwVO2f9erIlI-adXPLY,4126
+mcp_query_table-0.3.4.dist-info/licenses/LICENSE,sha256=rbvv_CTd7biGwT21tvhgQ2zkbPFXOoON7WFQWEdElBA,1063
+mcp_query_table-0.3.4.dist-info/METADATA,sha256=gc0VlBZxpcc_uPxkRA0WXdYNQ0cJeP3K5oTRKx5_4VU,8372
+mcp_query_table-0.3.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+mcp_query_table-0.3.4.dist-info/top_level.txt,sha256=5M_8dkO1USOX7_EWbWS6O_TEsZ5yo-AodFNKeUEgvEQ,16
+mcp_query_table-0.3.4.dist-info/RECORD,,

mcp_query_table-0.3.3.dist-info/RECORD

@@ -1,19 +0,0 @@
-mcp_query_table/__init__.py,sha256=gNXLL6MtH667HEdsUbJ4OAbGqwYO1dLdkqorV3-RtLM,238
-mcp_query_table/__main__.py,sha256=W3tMnZZTvkhjlLHmSYHng3CTvYX6-LjyaBDhf-ypH7w,1186
-mcp_query_table/_version.py,sha256=8KcCYTXH99C2-gCLuPILJvtT9YftRWJsartIx6TQ2ZY,22
-mcp_query_table/enums.py,sha256=Hen0X1f2of69f08epun6HvYIgbpw_rf8BvoRQ184kS4,679
-mcp_query_table/server.py,sha256=oV8CSi-EW5-xlPn38R1wLu7U-pCHgAJci1W8rJFsXSk,3430
-mcp_query_table/tool.py,sha256=uKqqrUrnUnvwlOipX0YODUg9WMq_Us994X8sysxe1cM,8571
-mcp_query_table/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mcp_query_table/providers/baidu.py,sha256=VkjAN6mSD4xa9kfDvbCb_Xj9Su5F2dakzYNgjoLGsR4,2738
-mcp_query_table/providers/n.py,sha256=GcqNMra1pc_DTw3QB3a2SUNzZCOrGjiuvuVDS-hTz5o,3006
-mcp_query_table/providers/yuanbao.py,sha256=BfHg3A6RGOT-GxfxkrAoiOQFFo6qydtI5tZSoRGtPiY,3244
-mcp_query_table/sites/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mcp_query_table/sites/eastmoney.py,sha256=LImjpYVuM5YnXwnNzB2hkKfHofocZZScetGqMOCHZpk,4477
-mcp_query_table/sites/iwencai.py,sha256=g56pj3pbxu4mXLNnaaS3Hdx-DvEy_9OBrQJe26z4z08,5059
-mcp_query_table/sites/tdx.py,sha256=RIUQaB7Tn4AVyWaevk9SzTKIDwVO2f9erIlI-adXPLY,4126
-mcp_query_table-0.3.3.dist-info/licenses/LICENSE,sha256=rbvv_CTd7biGwT21tvhgQ2zkbPFXOoON7WFQWEdElBA,1063
-mcp_query_table-0.3.3.dist-info/METADATA,sha256=cRXUETyJhkJ12POY1fMsZv2YHvFdBor8c_BqNOZhy5g,8372
-mcp_query_table-0.3.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-mcp_query_table-0.3.3.dist-info/top_level.txt,sha256=5M_8dkO1USOX7_EWbWS6O_TEsZ5yo-AodFNKeUEgvEQ,16
-mcp_query_table-0.3.3.dist-info/RECORD,,