NoteEx: Interactive Visual Context Manipulation for LLM-Assisted Exploratory Data Analysis in Computational Notebooks

TL;DR

NoteEx is a JupyterLab extension that helps data analysts manage and visualize their exploratory data analysis workflow, improving LLM assistance by enabling better context selection and mental model maintenance.

Contribution

It introduces a semantic visualization tool for EDA workflows that enhances context management and mental model externalization in computational notebooks.

Findings

Improved mental model retention with NoteEx.

More accurate and relevant LLM responses using NoteEx.

Enhanced user experience in context selection.

Abstract

Computational notebooks have become popular for Exploratory Data Analysis (EDA), augmented by LLM-based code generation and result interpretation. Effective LLM assistance hinges on selecting informative context -- the minimal set of cells whose code, data, or outputs suffice to answer a prompt. As notebooks grow long and messy, users can lose track of the mental model of their analysis. They thus fail to curate appropriate contexts for LLM tasks, causing frustration and tedious prompt engineering. We conducted a formative study (n=6) that surfaced challenges in LLM context selection and mental model maintenance. Therefore, we introduce NoteEx, a JupyterLab extension that provides a semantic visualization of the EDA workflow, allowing analysts to externalize their mental model, specify analysis dependencies, and enable interactive selection of task-relevant contexts for LLMs. A user study (n=12) against a baseline shows that NoteEx improved mental model retention and context selection, leading to more accurate and relevant LLM responses.