Debug Smarter, Not Harder: AI Agents for Error Resolution in Computational Notebooks (2024)

Abstract: Computational notebooks became indispensable tools for research-related development, offering unprecedented interactivity and flexibility in the development process. However, these benefits come at the cost of reproducibility and an increased potential for bugs.With the rise of code-fluent Large Language Models empowered with agentic techniques, smart bug-fixing tools with a high level of autonomy have emerged.However, those tools are tuned for classical script programming and still struggle with non-linear computational notebooks.In this paper, we present an AI agent designed specifically for error resolution in a computational notebook. We have developed an agentic system capable of exploring a notebook environment by interacting with it—similar to how a user would—and integrated the system into the JetBrains service for collaborative data science called Datalore.We evaluate our approach against the pre-existing single-action solution by comparing costs and conducting a user study. Users rate the error resolution capabilities of the agentic system higher but experience difficulties with UI. We share the results of the study and consider them valuable for further improving user-agent collaboration.

## introduction computational notebooks have become a popular medium for development during the last decade, especially for data analysis, machine learning @xcite , and creating educational @xcite or scientific content @xcite . one of the main features of computational notebooks is their statefulness-thus the notebook cannot be described only by its cells, but additionally, runtime information is required. the statefulness allows to work iteratively with the runtime in an additive manner and thus to efficiently go through hypotheses @xcite . however, it causes high code entanglement @xcite and, therefore, a higher number of errors in the code. as a result, notebooks are struggling with low reproducibility rates. after a re-run, they come to the same results with a 4% probability @xcite , and 75% of them could not be executed without exceptions @xcite . the resulting debugging distracts developers from the actual task. large language models (llms), such as gpt-4 (openai, 2023), mixtral @xcite , or code llama @xcite recently demonstrated advanced capabilities in solving complex code-related problems, such as code generation @xcite , debugging @xcite , and issue solving @xcite . however, there is a lack of studies applying such models to notebooks. the difficulty lies in the stateful nature of the notebook. since the notebook requires runtime information to represent its exact current state, it is hard to gather the context for an llm, as passing the entire runtime information is impossible due to the context size limitations. ai agents allow llms to interact with such an environment iteratively. an agent can explore the environment and achieve the goal autonomously, enabling it to adjust its actions based on the received feedback. such agents have shown abilities to engage with software engineering tasks @xcite @xcite , interact with web environments @xcite , or operate embodied agents @xcite . in this work, we present an ai agent for error resolution in computational notebooks. the proposed agent was integrated into datalore, 1 a jetbrains product for collaborative data science that allows development in cloud-hosted computational notebooks. we design the agent to be capable of creating, editing, and executing cells. this approach utilizes the notebook's natural interactivity and allows gradual expansion of context. the main contributions of our paper are: ## system design in this section, we will delve into the proposed system's architecture. the system contains three parts: an agent, an environment, and a user interface. the agent is a stateful back-end service responsible for orchestrating the communication between the llm and the notebook, storing prompts, and converting llm predictions into actions in the environment. the environment is the computational notebook that-in addition to being fully functional-is responsible for executing actions provided by the agent and providing corresponding observations. the user interface defines how programmers interact with the system as a whole. the goal of the system was to conduct the necessary code changes and cell executions to solve foot_1 the given runtime exception. ## evaluation we evaluated our system from two perspectives: system performance and user experience. for the former, we compared the costs of employing such an ai agent and the single-action solution already implemented in datalore. for the latter, we conducted a user study to estimate the effect on the developers' subjective productivity and satisfaction with error resolution capabilities. ## conclusion and future work in the present work, we have presented an agentic system for error resolution in computational notebooks. our solution was integrated into jet-brains datalore. the cost of running the system tripled, yet the cost stayed within the reasonable price range. the user study revealed many directions for further user-agent interaction research, such as ensuring the user's control over the agent or better visualization of the agent's actions. utilizing smaller and cheaper models and more intelligent information retrieval holds potential for cost-efficient next generations of such systems. the context caching techniques also look promising in iterative agentic applications. to benefit the community, we publish used prompts and the answers from the user study.