NL2Pandas Agent · Sacha Tihanyi

While I’ve used pandas quite a bit, it has always been for specific projects, after which I’d not touch it for a while. The bigger the gaps in time between using it, the rustier I became with its incredibly rich API. Given my interest in finding projects that involve building useful AI agents (particularly for myself), I thought it would be an interesting exercise to build an agent that could do all the data wrangling for me, and certainly more quickly than my practice of constantly opening new browser tabs to check pandas function signatures and documentation.

The end result is an agent that can take instructions and data, be chatted with on the fly, and self-correct and recover from errors. Saving intermediate steps to files (including errors) for observability and opening results in a web browser were useful additions to what is effectively a command-line driven task-specific agent. All of this runs through a Jupyter kernel running in a Docker sandbox.

Interesting Lessons

A stateful agent needs clear observations of the environment, not just conversation history.

There’s an approach to this problem that would rely purely on conversation history, including the code produced by the agent to generate outputs at each step. This would inherently rely on the LLM, in later prompts, to dig through that growing history to identify what objects exist in state and what steps were taken earlier. However, this is a weaker architecture from the perspective of how transformer models handle long contexts, given empirical issues like the “lost in the middle” context window problem. Considering that models often demonstrate better recall of information near the beginning and end of a context window (the U-shaped performance curve), providing fresh, compact state representations near the end of the context and right before the next set of instructions to the agent should logically improve performance. This is probably even more relevant during long conversation sessions where there have been a number of errors and self-correction loops added to the message history.

Designing a complete and yet compact state probe to run after each kernel execution step was key here. Beyond mitigating the “lost in the middle” problem, this addresses the more fundamental issue of conversation history telling the agent what it attempted to execute, but not necessarily what actually happened. Code can partially execute before failing, leaving behind objects or modifying existing ones in ways that aren’t obvious from the conversation history alone. Rather than requiring the LLM to reconstruct the environment from previous code and error traces, a state probe gives it a direct observation of what actually exists in the kernel. Providing a state snapshot to the LLM’s context window, immediately before the next set of user natural language instructions, gives the agent the most up-to-date and accurate representation of state needed to generate code to complete the user’s next request. This was also crucially important when running the program in ‘compact mode’, as the exact code history of how the current state was achieved would not be present (for the sake of saving context space).
Cleaning up the mess, AKA all roads lead to Rome.

I initially set up cleanup operations to handle only user-directed process exits (!exit via the program’s command line), with the conversation history logged, but quickly had to incorporate resource cleanup after implementing the Jupyter client in a Docker sandbox, including cases where graceful exit attempts failed. Ultimately, this meant ensuring there was a common cleanup and message history saving path that all manner of system exits would flow through, be it a graceful exit initiated by the user, SIGINT, SIGTERM, startup failures, 10 consecutive model-generated code errors, or any timeouts involving Jupyter kernel hangups and LLM call issues. The key lesson here was that different modes of exit could leave resources hanging and ports stuck in use, which was a problem that had to be dealt with for future program execution. All roads needed to lead to Rome, or in other words, all manner of exits had to flow through the same ultimate cleanup process.
What’s in the model’s training distribution?

I ran into an interesting problem early on when I tried to use Polars instead of pandas as my main data manipulation library. Despite using all the same general architecture and prompts, there ended up being far more error loops that the agent had to try to recover from with Polars than with pandas. From looking at the logs of conversation history with error traces and what the agent was attempting in its effort at correcting its code, I realized that the various open-source models I was using were making far more API syntax or function-call errors with Polars than pandas. Polars, being newer and perhaps subject to greater change (and perhaps less online Q&A available for training data) may have a weaker learned representation in the models than the pandas API.

A potential fix that didn’t include just nixing Polars in favor of pandas would be to give the agent an up-to-date searchable representation of the Polars docs, or an online search capability/tool. However, this would lead to greater latency and greater token consumption than simply using pandas. This is probably going to be less of a problem going forward with newer models thanks to better Polars representation in the training data. It was still a very interesting and unexpected problem involving the model’s familiarity with different APIs to encounter and to think about the various cost-benefit tradeoffs in fixing it. I obviously went with the ‘just use pandas’ solution, although the documentation search tool call would be the more robust long-term approach, despite the increased cost and latency.

Putting It to Work

I achieved some amazing time savings using this agentic tool that made me wish I had it years ago. One of my favourite tests of the agent across different models was giving it a CSV file of over 272k lines, each with a date, exchange rate value, and exchange rate name. What was tricky was the fact that the name column contained 23 different exchange rates and would change to a new exchange rate at irregular intervals, as each exchange rate series had a different historical length. I wanted to use this data to construct MACD indicators for each currency over its entire available data history, so I submitted the following instructions:

The data contains exchange rates over time on a daily financial calendar frequency, and is listed as CSV lines of: date, FX rate, FX name. Each exchange rate is given in its full time series before the next exchange rate begins, and each time series is potentially of a different length. You will clean this data by following these steps:

Step 0

Separate these exchange rates into their own columns and ensure that there is only one date column, which is to be the index that all exchange rates share. Note that the date data is in day/month/year format to begin with.

Step 1

Select a subset of the data where there are no remaining NA or NaN values, as there are many data series in the original dataset that don’t have values until later on in the time series.

Step 2

For each exchange rate calculate the 12-day exponential moving average (EMA) and the 26-day EMA.

Step 3

For each exchange rate calculate the MACD using the standard 12-26-9 approach with the 12-day EMA and 26-day EMA, and also calculate the 9-day EMA of the MACD for the signal line.

I would normally put the agent in step mode (a command line option) for this task (though it’s not necessary) to track each step. I would thus get intermediate results written to files, and those results opened in browser windows for viewing…with the last step generating what is definitely a DataFrame larger than what should be opened in a browser. The amount of time the agent saved me versus doing it myself by hand with pandas made me shudder when I thought back to the amount of time I’ve spent on monotonous data wrangling. The productivity improvement is pretty phenomenal.