Spaces:
Running
Running
File size: 6,694 Bytes
64a008c faefb1f 64a008c 722c296 64a008c faefb1f 64a008c e463de2 64a008c e463de2 64a008c e463de2 64a008c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | from __future__ import annotations
import asyncio
import json
import logging
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
from app.services.dataset_metadata_service import extract_metadata
from app.utils.http_utils import download_url
from app.utils.subprocess_utils import run_subprocess
logger = logging.getLogger(__name__)
_PYTHON = getattr(sys, "executable", None) or "python3"
_MAX_OUTPUT_BYTES = 65536
_MAX_CONCURRENT = 8
_DOWNLOAD_TIMEOUT = 120
_semaphore = asyncio.Semaphore(_MAX_CONCURRENT)
class CSVAnalysisError(Exception):
pass
async def _resolve_source(source: Union[str, bytes]) -> Tuple[bytes, Optional[str]]:
if isinstance(source, str) and source.lower().startswith(("http://", "https://")):
return await download_url(source, timeout_seconds=_DOWNLOAD_TIMEOUT)
elif isinstance(source, bytes):
return source, None
else:
raise TypeError(f"Unsupported source type: {type(source)}")
def _run_subprocess(cmd: List[str], timeout: float, max_output: int) -> Dict[str, Any]:
return run_subprocess(cmd, timeout, max_output)
_CHAT_SCRIPT = """\
import json, sys, io, base64, traceback
import pandas as pd, numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
df = pd.read_csv(r"{csv_path}")
with open(r"{blocks_path}", "r") as _f:
_data = json.load(_f)
_results = {{"analyze": [], "visualization": []}}
_ns = {{"df": df, "pd": pd, "np": np, "plt": plt, "sns": sns}}
_EXCLUDED = {{"pd", "np", "plt", "sns", "df", "json", "sys", "io", "base64", "traceback", "matplotlib", "seaborn"}}
for _b in _data.get("analyze", []):
_code = (_b.get("python_code") or "").strip()
if not _code:
_results["analyze"].append({{"success": True, "output": "", "error": None}})
continue
_old = sys.stdout
sys.stdout = io.StringIO()
_pre = {{k for k in _ns if not k.startswith("_")}} - _EXCLUDED
try:
exec(_code, _ns)
_output = sys.stdout.getvalue()
if not _output:
_post = {{k for k in _ns if not k.startswith("_")}} - _EXCLUDED
_new = sorted(_post - _pre)
if "final_result" in _ns:
_output += f"final_result = {{repr(_ns['final_result'])}}\\n"
for _k in _new:
if _k == "final_result":
continue
_v = _ns[_k]
try:
_s = repr(_v)
except Exception:
_s = "<unrepresentable>"
_output += f"{{_k}} = {{_s}}\\n"
_results["analyze"].append({{"success": True, "output": _output, "error": None}})
except Exception:
_results["analyze"].append({{"success": False, "output": sys.stdout.getvalue(), "error": traceback.format_exc()}})
finally:
sys.stdout = _old
for _b in _data.get("visualization", []):
_code = (_b.get("python_code") or "").strip()
if not _code:
_results["visualization"].append({{"success": True, "image_base64": "", "error": None}})
continue
_full = _code + (
"\\nfrom io import BytesIO\\nimport base64\\n"
"_buf = BytesIO()\\nplt.savefig(_buf, format='png', bbox_inches='tight', dpi=150)\\n"
"_buf.seek(0)\\nprint(base64.b64encode(_buf.read()).decode(), end='')\\n"
"plt.close('all')\\n"
)
_old = sys.stdout
sys.stdout = io.StringIO()
try:
exec(_full, _ns)
_results["visualization"].append({{"success": True, "image_base64": sys.stdout.getvalue().strip(), "error": None}})
except Exception:
_results["visualization"].append({{"success": False, "image_base64": None, "error": traceback.format_exc()}})
finally:
sys.stdout = _old
plt.close("all")
print(json.dumps(_results))
"""
async def execute_csv_chat_blocks(
source: Union[str, bytes],
analyze_blocks: List[Dict[str, Any]],
viz_blocks: List[Dict[str, Any]],
timeout: int = 60,
) -> Dict[str, Any]:
data, _ = await _resolve_source(source)
if not data:
return {"success": False, "results": None, "error": "No data provided"}
async with _semaphore:
run_dir = None
start = time.monotonic()
try:
run_dir = Path(tempfile.mkdtemp())
csv_path = run_dir / "data.csv"
csv_path.write_bytes(data)
blocks_path = run_dir / "blocks.json"
blocks_path.write_text(
json.dumps({"analyze": analyze_blocks, "visualization": viz_blocks}),
encoding="utf-8",
)
script = _CHAT_SCRIPT.format(
csv_path=csv_path.as_posix(),
blocks_path=blocks_path.as_posix(),
)
script_path = run_dir / "chat_exec.py"
script_path.write_text(script, encoding="utf-8")
cmd = [_PYTHON, str(script_path)]
result = await asyncio.to_thread(_run_subprocess, cmd, timeout, _MAX_OUTPUT_BYTES)
elapsed_ms = (time.monotonic() - start) * 1000
if result["exit_code"] != 0:
return {
"success": False,
"results": None,
"error": result["stderr"] or "Subprocess failed",
"execution_time_ms": round(elapsed_ms, 2),
}
parsed = json.loads(result["stdout"])
return {
"success": True,
"results": parsed,
"error": None,
"execution_time_ms": round(elapsed_ms, 2),
}
except json.JSONDecodeError as exc:
elapsed_ms = (time.monotonic() - start) * 1000
return {"success": False, "results": None, "error": f"Failed to parse output: {exc}", "execution_time_ms": round(elapsed_ms, 2)}
except FileNotFoundError:
elapsed_ms = (time.monotonic() - start) * 1000
return {"success": False, "results": None, "error": f"Python runtime ({_PYTHON}) not found", "execution_time_ms": round(elapsed_ms, 2)}
except Exception as exc:
elapsed_ms = (time.monotonic() - start) * 1000
logger.exception("CSV chat execution error")
return {"success": False, "results": None, "error": f"Execution error: {exc}", "execution_time_ms": round(elapsed_ms, 2)}
finally:
if run_dir and run_dir.exists():
shutil.rmtree(run_dir, ignore_errors=True)
async def get_dataset_info(source: Union[str, bytes, Any]) -> Dict[str, Any]:
return await extract_metadata(source)
|