Skip to content

DatawiseAgent: A Notebook-Centric LLM Agent Framework for Adaptive and Robust Data Science Automation

🕒 Published (v1): 2025-03-10 08:32 UTC · Source: Arxiv · Venue: EMNLP 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

DatawiseAgent is a notebook-centric LLM agent framework that unifies all agent–user–environment communication as interleaved markdown and code cells, governed by a finite-state transducer (FST) across four stages: DFS-like planning, incremental execution, self-debugging, and post-filtering. It achieves state-of-the-art performance on data analysis, scientific visualization, and predictive modeling benchmarks while degrading gracefully on smaller open-source models.

Problem

Existing data science agents are siloed to pipeline phases (feature engineering, model selection, etc.), struggle to generalize across task types and LLM capabilities, and assume access to SOTA proprietary models (e.g., GPT-4o), making them brittle in resource-constrained or privacy-sensitive settings. No prior framework provides a unified, long-horizon, self-repairing control structure that works across diverse LLMs.

Method

DatawiseAgent has two core components:

Unified interaction representation: All agent–user–environment communication (instructions, tool calls, environment info, observations, debug traces) is encoded as sequences of interleaved markdown and executable code cells within a computational notebook. This single format applies at every stage, reducing format-switching overhead for weaker models.

FST-based multi-stage architecture: The agent is formalized as a nondeterministic finite-state transducer (NFST) with state space \(Q = \{q_{plan}, q_{inc}, q_{debug}, q_{filter}, q_0\}\). State transitions are governed by \(\delta(q, \sigma, f)\) where \(\sigma\) is the agent-generated action signal and \(f \in \{\text{error}, \text{no\_error}\}\) is execution feedback. The four stages are: 1. DFS-like planning: non-linear, tree-structured subgoal selection — advance, backtrack, or terminate. 2. Incremental execution: step-by-step cell generation with fine-grained feedback per cell rather than one-shot code generation. 3. Self-debugging: iterative faulty-code repair using execution traces. 4. Post-filtering: if debugging succeeds, extracts clean code; if not, emits a concise diagnostic markdown report to prevent context pollution, then resumes.

Key Contributions

  • Notebook-centric unified interaction representation that collapses all communication modalities into markdown+code cell sequences.
  • FST formalism for agent lifecycle management enabling modular stage extension and fine-grained ablation.
  • DFS-like non-linear planning paired with incremental per-cell execution for long-horizon task solving without requiring strong LLMs.
  • Code repair module (self-debugging + post-filtering) that prevents error trace accumulation from polluting the context window.
  • Demonstrated model-agnostic robustness: graceful degradation from GPT-4o to Qwen2.5-7B.

Results

  • Data analysis (InfiAgent-DABench, 257 challenges): DatawiseAgent achieves 82.88% ABQ (GPT-4o mini), 85.99% (GPT-4o, matching Taskweaver), and 81.71% (Qwen2.5-72B) — best or tied-best across all model configurations vs. ReAct, AutoGen, Taskweaver, and Data Interpreter.
  • Scientific visualization (MatplotBench, 100 cases): Average scores of 58.60 (GPT-4o mini, +20.51 over direct decoding), 64.33 (GPT-4o, new SOTA), and 61.88 (Qwen2.5-72B) with visual tool — best across all three model configurations.
  • Predictive modeling (DSBench data modeling, 74 Kaggle tasks): 98.64% task success rate and RPG of 53.18 (GPT-4o), 98.64%/46.61 (GPT-4o mini), 91.89%/42.90 (Qwen2.5-72B) — surpassing AutoGen (71.62%/34.74 with GPT-4o) and Code Interpreter (54.05%/26.14 with GPT-4) by large margins.
  • Robustness across model sizes (InfiAgent-DABench, Qwen2.5 7B–72B): DatawiseAgent maintains top ABQ across all sizes; performance gap over competing methods widens on smaller models.

Limitations

  • Evaluation covers only three benchmarks; broader task diversity (e.g., multi-modal data, NLP pipelines) is not assessed.
  • Self-debugging is triggered reactively on execution errors; proactive error anticipation is not addressed.
  • The 3600-second task time limit may under-represent agent capability on very long-horizon real-world tasks.
  • Context window growth from long cell sequences in extended tasks is not explicitly managed beyond post-filtering.
  • The FST transition function is hand-designed; learned or adaptive control flow is not explored.

Relevance to Harnesses / Meta-Harnesses

DatawiseAgent is a concrete example of a single-agent meta-harness: the FST-based architecture functions as an explicit control harness that schedules, reroutes, and repairs sub-processes (planning, execution, debugging, filtering) without relying on a separate orchestrator agent. The post-filtering stage — which discards polluted traces and synthesizes a diagnostic stub — mirrors the "clean context handoff" pattern central to robust multi-stage harness design. The unified cell-sequence interaction representation is directly analogous to a harness's shared state/blackboard, making it a concrete implementation blueprint. The modular stage design and explicit ablation support harness researchers seeking to understand which orchestration components contribute most to robustness under varying model capacity.