Skip to content
All work

An agent for data science, not just the SQL

65% execution accuracy on BIRD, 8 points above GPT-4o

At UCLA Trustworthy AI Lab I built the tool-calling layer, the ingestion pipeline, and the data model behind an agent that plans and runs a whole analysis from a plain-language prompt. It beat GPT-4o on BIRD, and the leverage was upstream of the model.

65%
Execution accuracy on the BIRD benchmark
18% → 5%
Agent failure rate, after re-architecting context
80
Distinct user schemas served, no preprocessing
100+
Concurrent sessions, scaling from zero
A bar chart of execution accuracy on the BIRD benchmark. This agent scores 65 percent; GPT-4o scores 57 percent on the same tasks, a gap of eight percentage points. execution accuracy · bird benchmark 0% 25% 50% 75% 100% This agent 65% GPT-4o 57% +8 pts same tasks · same schemas · n = 400+ sessions
Fig. 1Execution accuracy on BIRD against GPT-4o on the same tasks and schemas. Execution accuracy is BIRD's own measure: the generated code has to run and return the right rows, not merely look correct.

Context

Most “talk to your data” systems are text-to-SQL: one prompt in, one query out. Real analysis isn’t shaped like that. It’s plan, query, look at what came back, revise, query again, and the interesting failures happen between the steps, not inside them.

Two things have to be true before any of that reasoning matters. The agent needs tools it can compose, and it needs the user’s data in a shape it can actually query. My work was mostly the second one, and that’s the half that usually gets skipped.

BIRD is a good measuring stick because its questions require reasoning over messy, realistic schemas rather than clean toy tables.

The problem

Nobody is going to clean their data first. A researcher exports a CSV from somewhere and expects to ask questions about it. No declared schema, no types, no promise about size. A platform that asks them to preprocess has already lost the user it was built for.

Planning. Decompose a vague request into ordered steps, pick the right tool for each, and recognize when a result means the plan was wrong.

State, which people underestimate. A multi-step agent on serverless has no natural place to keep its own state: steps land on different Lambda invocations, run concurrently, and can fail partway. Losing state mid-analysis doesn’t just fail the task. It fails it silently, leaving a plausible partial answer.

What I built

The tool-calling layer, the surface the agent plans against. A prompt and a file go in; runnable Python, charts, and exportable tables come out. Most of the work was granularity: how much a single tool should do, and how much of its result should travel back to the planner rather than stay in storage. 65% execution accuracy on BIRD, 8 points above the GPT-4o baseline (Fig. 1).

An ingestion pipeline that infers rather than asks (Fig. 2). Uploads go straight to S3 through presigned URLs, so a 300 MB Parquet never passes through the application server at all. The pipeline then samples the file, infers the schema and each column’s type, and registers a typed table the tools can query. It has served 80 distinct user schemas with no manual preprocessing, which is the difference between a demo and something someone opens a second time.

A user's CSV, JSON or Parquet file uploads straight to S3 through a presigned URL, never passing through the application server. The pipeline reads a sample, infers the schema and column types, and registers a typed table the agent's tools can query. 80 distinct user schemas, tables up to 300 megabytes, and no manual preprocessing. raw file csv · json · parquet up to 300 MB no app server in the upload path Presigned S3 direct to bucket Infer schema + types Typed table 80 schemas Agent tools sql · pandas types are inferred from a sample, never declared by the user zero manual preprocessing · the user brings a file, not a schema
Fig. 2The ingestion path. A user brings a file, not a schema. The pipeline samples it, infers the column types, and registers a table the agent's tools can query, so nothing downstream depends on the user preparing their data first.

An event-sourced data model in PostgreSQL (Fig. 3), on Supabase, with larger artifacts in S3. Every step appends an immutable event: the conversation turn, the model’s response, each tool call, and the failure if there was one. Current state is a fold over that log. That solved distributed state outright. Concurrent steps can’t clobber each other, a failed step leaves a record rather than corruption, and any session replays exactly.

Then I queried it. The log was built for durability, but it doubled as instrumentation, and that turned out to be the more valuable half. Grouping failures by tool path showed a small set of paths spending most of the context window before the agent ever reached an answer. That pointed at context management rather than prompt wording, and redesigning it cut the failure rate from 18% to 5%.

Requests enter through API Gateway to a Dockerized Lambda. Each step appends an immutable event to a log stored in PostgreSQL and S3. Current state is a fold over that log, which makes any session replayable. API Gateway concurrent requests Lambda Docker · CDK append-only event log plan query observe revise answer append fold PostgreSQL events S3 artifacts state = fold(events) · any session replays exactly 400+ sessions · zero data loss · concurrent steps cannot clobber
Fig. 3State as a fold over an append-only log: concurrent steps can't clobber each other, and a failure leaves a record rather than a corrupted document. That record is also the dataset I later mined to find where the agent was spending its context window.

A serverless backend sized for concurrent inference. Container-image Lambda behind API Gateway, provisioned with CDK. Container images because the analysis stack, pandas and everything underneath it, runs far past Lambda’s 250 MB zipped package limit. CDK because the infrastructure had to be reproducible for anyone rerunning the experiments. It scales from zero to 100+ concurrent sessions with no idle compute cost, which is the right shape for a research platform that sits unused between sessions.

A data-centric evaluation harness. Rather than end-to-end pass/fail, it scores across all 80 table schemas, 120+ question types, and 20+ tool-calling functions, so a failure attributes to a specific tool, schema shape, or question class instead of to the agent as a whole.

The hard part

Resisting the urge to fix the prompt. Every failure looks like a prompting problem until something can attribute it. Once the harness and the event log together could say “this question class fails on this tool, and here is how much context it burned getting there,” the fixes stopped being verbal and became structural: tool granularity, and the shape of what came back to the planner.

Outcome

65% execution accuracy on BIRD, 8 points above GPT-4o on the same tasks (Fig. 1), with the failure rate down from 18% to 5%. Beating a frontier model here wasn’t a modeling win. The leverage sat upstream of the model: getting arbitrary user data into a queryable shape without asking the user for help, and instrumenting the system well enough to know exactly where it lost.

What I’d do differently

Build the instrumentation first. I built the event log for durability and the harness third, so every improvement before them was guesswork I couldn’t measure. On an agent project the log and the harness aren’t test infrastructure. They’re the instrument you’re doing the work with.