DS-STAR: Data Science Agent via Iterative Planning and Verification

TL;DR

DS-STAR introduces a data science agent capable of handling heterogeneous data formats by first performing automatic per-file analysis to generate rich contextual descriptions, then solving tasks through an iterative loop of planning, implementation, execution, and LLM-based verification. The core contributions are the data-file analysis module, an LLM-based verifier that judges plan sufficiency, and a sequential planning mechanism that refines plans until they are verified as sufficient. Empirical results on DABStep, KramaBench, and DA-Code show state-of-the-art performance, especially on hard tasks requiring multi-file, multi-format data processing. The approach advances practical data science automation and lays groundwork for future human-in-the-loop enhancements to further boost performance and applicability.

Abstract

Data science, which transforms raw data into actionable insights, is critical for data-driven decision-making. However, these tasks are often complex, involving steps for exploring multiple data sources and synthesizing findings to deliver insightful answers. While large language models (LLMs) show significant promise in automating this process, they often struggle with heterogeneous data formats and generate sub-optimal analysis plans, as verifying plan sufficiency is inherently difficult without ground-truth labels for such open-ended tasks. To overcome these limitations, we introduce DS-STAR, a novel data science agent. Specifically, DS-STAR makes three key contributions: (1) a data file analysis module that automatically explores and extracts context from diverse data formats, including unstructured types; (2) a verification step where an LLM-based judge evaluates the sufficiency of the analysis plan at each stage; and (3) a sequential planning mechanism that starts with a simple, executable plan and iteratively refines it based on the DS-STAR's feedback until its sufficiency is verified. This iterative refinement allows DS-STAR to reliably navigate complex analyses involving diverse data sources. Our experiments show that DS-STAR achieves state-of-the-art performance across three challenging benchmarks: DABStep, KramaBench, and DA-Code. Moreover, DS-STAR particularly outperforms baselines on hard tasks that require processing multiple data files with heterogeneous formats.

Figures (6)

• Figure 1: (a) Data science tasks require processing data files in various formats (e.g., csv, txt, xlsx, md) to answer user queries, such as data analysis for extracting useful insights or predictive tasks using machine learning. Data science agents are designed to accomplish this by writing code scripts (e.g., Python) and answering based on the output from their execution. In addition to the solution code, the output may include a trained model for prediction tasks, processed databases, text-formatted answers, visual charts, and more. (b) We report normalized accuracy (%) for both easy (answer can be found within a single file) and hard (requires processing multiple files) tasks in the DABStep, KramaBench, and DA-Code benchmarks. DS-STAR consistently and significantly outperforms competitive baselines, particularly in challenging hard tasks, showing DS-STAR's superiority in handling multiple data sources in heterogeneous formats.
• Figure 2: Overview of DS-STAR. (a) DS-STAR generates a Python script that extracts essential information for analyzing heterogeneous input data files. (b) Starting with an initial plan, i.e., a simple single step, DS-STAR implements the plan and executes it to obtain intermediate results. Next, using a verifer agent, DS-STAR determines whether the current plan is sufficient to resolve the user's query. If it is sufficient, the currently implemented code becomes the final solution. If not, DS-STAR's router agent decides whether to add the next step or remove the incorrect step from the current plan. Finally, DS-STAR adds the next step and implements it again. This iterative process of planning, implementation, execution, and verification is repeated until the verifier agent determines that the plan is sufficient to answer the user's query of until the maximum number of iterations is reached.
• Figure 3: Qualitative result. We observed that while ReAct yao2023react failed to filter out the 'NextPay' value, leading to an incorrect answer, DS-STAR successfully filtered out due to the use of data file analysis agent. See Appendix \ref{['app:qual']} for more qualitative results of DS-STAR.
• Figure 4: Sensitivity analysis on the number of refinement steps using DABStep. (a) Difficult tasks require more iterations to generate sufficient plans. Specifically, hard-level tasks require an average 5.6 iterations, while easy-level tasks require 3.0 iterations. Also, more than 50% of easy-level tasks are done only with a single round. (b) Performing more iterations allows the agent to generate sufficient plans, resulting in better performance for both easy and hard level tasks.
• Figure 5: Generated bar chart.