275 lines
9.3 KiB
Python
275 lines
9.3 KiB
Python
"""
|
||
Excel 表格分析：下载、schema 探测、pandas 沙箱执行。
|
||
|
||
设计原则：
|
||
- MCP 容器内执行，与 backend 进程隔离（建议 Docker 部署）
|
||
- 仅允许对已加载的 ``df`` 做 pandas 运算，禁止 import / 文件 / 网络
|
||
- 代码必须将最终答案赋给 ``result``
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import ast
|
||
import json
|
||
import tempfile
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
import pandas as pd
|
||
|
||
MAX_DOWNLOAD_BYTES = 20 * 1024 * 1024
|
||
DOWNLOAD_TIMEOUT = 120.0
|
||
EXEC_TIMEOUT_SEC = 15
|
||
MAX_RESULT_ROWS = 200
|
||
|
||
FORBIDDEN_CALL_NAMES = frozenset(
|
||
{
|
||
"open",
|
||
"exec",
|
||
"eval",
|
||
"compile",
|
||
"__import__",
|
||
"getattr",
|
||
"setattr",
|
||
"delattr",
|
||
"globals",
|
||
"locals",
|
||
"vars",
|
||
"dir",
|
||
"help",
|
||
"input",
|
||
"breakpoint",
|
||
"exit",
|
||
"quit",
|
||
}
|
||
)
|
||
|
||
FORBIDDEN_ATTR_NAMES = frozenset({"__import__", "__subclasses__", "__globals__", "__code__"})
|
||
|
||
SAFE_BUILTINS = {
|
||
"len": len,
|
||
"str": str,
|
||
"int": int,
|
||
"float": float,
|
||
"list": list,
|
||
"dict": dict,
|
||
"tuple": tuple,
|
||
"set": set,
|
||
"sum": sum,
|
||
"min": min,
|
||
"max": max,
|
||
"range": range,
|
||
"round": round,
|
||
"abs": abs,
|
||
"sorted": sorted,
|
||
"enumerate": enumerate,
|
||
"zip": zip,
|
||
"map": map,
|
||
"filter": filter,
|
||
"any": any,
|
||
"all": all,
|
||
"bool": bool,
|
||
"True": True,
|
||
"False": False,
|
||
"None": None,
|
||
}
|
||
|
||
|
||
def _validate_pandas_code(code: str) -> None:
|
||
if not code or not code.strip():
|
||
raise ValueError("pandas_code 不能为空")
|
||
if len(code) > 8000:
|
||
raise ValueError("pandas_code 过长（上限 8000 字符）")
|
||
|
||
tree = ast.parse(code, mode="exec")
|
||
|
||
for node in ast.walk(tree):
|
||
if isinstance(node, (ast.Import, ast.ImportFrom)):
|
||
raise ValueError("禁止 import，请仅使用已提供的 pd 与 df")
|
||
if isinstance(node, ast.Call):
|
||
func = node.func
|
||
if isinstance(func, ast.Name) and func.id in FORBIDDEN_CALL_NAMES:
|
||
raise ValueError(f"禁止调用: {func.id}")
|
||
if isinstance(func, ast.Attribute) and func.attr in FORBIDDEN_ATTR_NAMES:
|
||
raise ValueError(f"禁止访问: {func.attr}")
|
||
if isinstance(node, ast.Attribute) and node.attr in FORBIDDEN_ATTR_NAMES:
|
||
raise ValueError(f"禁止访问属性: {node.attr}")
|
||
|
||
|
||
def _download_excel(file_url: str) -> Path:
|
||
url = (file_url or "").strip()
|
||
if not url.startswith(("http://", "https://")):
|
||
raise ValueError("file_url 必须是 http(s) 可下载地址")
|
||
|
||
suffix = ".xlsx"
|
||
lower = url.split("?", 1)[0].lower()
|
||
if lower.endswith(".xls"):
|
||
suffix = ".xls"
|
||
elif lower.endswith(".csv"):
|
||
suffix = ".csv"
|
||
|
||
with httpx.Client(timeout=DOWNLOAD_TIMEOUT, follow_redirects=True) as client:
|
||
with client.stream("GET", url) as resp:
|
||
resp.raise_for_status()
|
||
cl = resp.headers.get("content-length")
|
||
if cl and int(cl) > MAX_DOWNLOAD_BYTES:
|
||
raise ValueError(f"文件过大（>{MAX_DOWNLOAD_BYTES // (1024 * 1024)}MB）")
|
||
|
||
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
|
||
try:
|
||
total = 0
|
||
for chunk in resp.iter_bytes():
|
||
total += len(chunk)
|
||
if total > MAX_DOWNLOAD_BYTES:
|
||
raise ValueError(f"文件过大（>{MAX_DOWNLOAD_BYTES // (1024 * 1024)}MB）")
|
||
tmp.write(chunk)
|
||
tmp.flush()
|
||
return Path(tmp.name)
|
||
except Exception:
|
||
tmp.close()
|
||
Path(tmp.name).unlink(missing_ok=True)
|
||
raise
|
||
finally:
|
||
tmp.close()
|
||
|
||
|
||
def _resolve_sheet_name(sheet_name: str, excel_path: Path) -> str | int:
|
||
raw = (sheet_name or "").strip()
|
||
xl = pd.ExcelFile(excel_path)
|
||
if not raw or raw.lower() in ("0", "first", "default"):
|
||
return xl.sheet_names[0] if xl.sheet_names else 0
|
||
if raw.isdigit():
|
||
idx = int(raw)
|
||
if idx < 0 or idx >= len(xl.sheet_names):
|
||
raise ValueError(f"sheet 索引越界: {idx}，共 {len(xl.sheet_names)} 个 sheet")
|
||
return xl.sheet_names[idx]
|
||
if raw not in xl.sheet_names:
|
||
raise ValueError(f"sheet 不存在: {raw}，可选: {xl.sheet_names}")
|
||
return raw
|
||
|
||
|
||
def _load_dataframe(excel_path: Path, sheet_name: str) -> pd.DataFrame:
|
||
if excel_path.suffix.lower() == ".csv":
|
||
return pd.read_csv(excel_path)
|
||
resolved = _resolve_sheet_name(sheet_name, excel_path)
|
||
return pd.read_excel(excel_path, sheet_name=resolved)
|
||
|
||
|
||
def _format_result(result: Any) -> str:
|
||
if isinstance(result, pd.DataFrame):
|
||
df = result.head(MAX_RESULT_ROWS)
|
||
payload = {
|
||
"type": "dataframe",
|
||
"rows": len(result),
|
||
"truncated": len(result) > MAX_RESULT_ROWS,
|
||
"data": json.loads(df.to_json(orient="records", force_ascii=False)),
|
||
}
|
||
return json.dumps(payload, ensure_ascii=False)
|
||
if isinstance(result, pd.Series):
|
||
s = result.head(MAX_RESULT_ROWS)
|
||
payload = {
|
||
"type": "series",
|
||
"length": len(result),
|
||
"truncated": len(result) > MAX_RESULT_ROWS,
|
||
"data": json.loads(s.to_json(force_ascii=False)),
|
||
}
|
||
return json.dumps(payload, ensure_ascii=False)
|
||
if isinstance(result, (int, float, bool)) or result is None:
|
||
return json.dumps({"type": "scalar", "value": result}, ensure_ascii=False)
|
||
if isinstance(result, str):
|
||
return json.dumps({"type": "text", "value": result[:5000]}, ensure_ascii=False)
|
||
return json.dumps({"type": "other", "value": str(result)[:5000]}, ensure_ascii=False)
|
||
|
||
|
||
def _sandbox_execute(code: str, df: pd.DataFrame) -> Any:
|
||
_validate_pandas_code(code)
|
||
namespace: dict[str, Any] = {"pd": pd, "df": df, "result": None}
|
||
compiled = compile(code, "<pandas_sandbox>", "exec")
|
||
|
||
def _run() -> None:
|
||
exec(compiled, {"__builtins__": SAFE_BUILTINS}, namespace) # noqa: S102
|
||
|
||
import concurrent.futures
|
||
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||
fut = pool.submit(_run)
|
||
try:
|
||
fut.result(timeout=EXEC_TIMEOUT_SEC)
|
||
except concurrent.futures.TimeoutError as e:
|
||
raise TimeoutError(f"执行超时（>{EXEC_TIMEOUT_SEC}s）") from e
|
||
|
||
if namespace.get("result") is None:
|
||
raise ValueError("请在 pandas_code 末尾将答案赋给变量 result，例如: result = df['列'].sum()")
|
||
return namespace["result"]
|
||
|
||
|
||
def get_excel_schema(file_url: str, sheet_name: str = "") -> str:
|
||
"""
|
||
下载 Excel/CSV 并返回 sheet 列表、列名、类型与前几行预览（JSON 字符串）。
|
||
"""
|
||
path = _download_excel(file_url)
|
||
try:
|
||
if path.suffix.lower() == ".csv":
|
||
df = pd.read_csv(path, nrows=5)
|
||
preview_rows = json.loads(
|
||
df.head(5).to_json(orient="records", force_ascii=False)
|
||
)
|
||
payload = {
|
||
"format": "csv",
|
||
"sheets": ["default"],
|
||
"default_sheet": "default",
|
||
"schemas": {
|
||
"default": {
|
||
"columns": [str(c) for c in df.columns],
|
||
"dtypes": {str(c): str(df[c].dtype) for c in df.columns},
|
||
"preview_rows": preview_rows,
|
||
}
|
||
},
|
||
}
|
||
return json.dumps(payload, ensure_ascii=False)
|
||
|
||
xl = pd.ExcelFile(path)
|
||
schemas: dict[str, Any] = {}
|
||
for name in xl.sheet_names:
|
||
df = pd.read_excel(path, sheet_name=name, nrows=5)
|
||
schemas[name] = {
|
||
"columns": [str(c) for c in df.columns],
|
||
"dtypes": {str(c): str(df[c].dtype) for c in df.columns},
|
||
"preview_rows": json.loads(
|
||
df.head(5).to_json(orient="records", force_ascii=False)
|
||
),
|
||
}
|
||
payload = {
|
||
"format": "excel",
|
||
"sheets": xl.sheet_names,
|
||
"default_sheet": xl.sheet_names[0] if xl.sheet_names else "",
|
||
"schemas": schemas,
|
||
"hint": "分析时先选 sheet_name，再调用 run_pandas_on_excel；代码中只用 df 与 pd，结果赋给 result",
|
||
}
|
||
if sheet_name.strip():
|
||
resolved = _resolve_sheet_name(sheet_name, path)
|
||
key = resolved if isinstance(resolved, str) else xl.sheet_names[int(resolved)]
|
||
payload["selected_sheet"] = key
|
||
return json.dumps(payload, ensure_ascii=False)
|
||
finally:
|
||
path.unlink(missing_ok=True)
|
||
|
||
|
||
def run_pandas_on_excel(file_url: str, pandas_code: str, sheet_name: str = "0") -> str:
|
||
"""
|
||
在沙箱中对 Excel/CSV 执行 pandas 代码。
|
||
|
||
pandas_code 示例::
|
||
result = df[df['区域'] == '华东']['销售额'].sum()
|
||
result = df.groupby('部门')['金额'].mean().sort_values(ascending=False).head(10)
|
||
"""
|
||
path = _download_excel(file_url)
|
||
try:
|
||
df = _load_dataframe(path, sheet_name)
|
||
if len(df) > 100_000:
|
||
raise ValueError("表格行数过多（>10万），请让用户拆分文件或缩小范围")
|
||
result = _sandbox_execute(pandas_code, df)
|
||
return _format_result(result)
|
||
finally:
|
||
path.unlink(missing_ok=True)
|