Turning a Question Into a Pandas Query With an LLM

The model should never compute the answer. It should write the query that computes the answer, and something deterministic should check that query before it runs. That division is what makes this reliable enough to put in front of people.

The shape of the problem

“What was average order value by region last quarter, excluding refunds?” is a question with an exact answer that a three-line pandas expression produces. The model is good at going from the English to the expression and bad at going from the English to the number, for the reason set out in what an LLM can and cannot do with a spreadsheet: it has no accumulator, so a reduction over 40,000 rows is a plausible-looking guess.

So the architecture is fixed by that constraint. Model writes code, validator inspects the code, interpreter runs it, and the code is shown to the user next to the result so the answer is auditable. Anything that skips the middle step is a demo.

Why not just ask for the answer with the data inline

Because it costs a six-figure token count for a medium sheet, recurs on every follow-up question, and produces an answer that is confidently wrong in the third significant figure with no exception raised. The generated-code path sends a schema of a few hundred tokens regardless of whether the table has a thousand rows or fifty million.

What the model needs to see

Not the data. A schema, and specifically a schema with the values in it, because the single most common failure is a category the model guessed the spelling of.

def describe(df, max_cats: int = 12) -> str:
    lines = []
    for col in df.columns:
        s = df[col]
        line = f"- {col}: {s.dtype}"
        if s.dtype == "object" or str(s.dtype) == "category":
            n = s.nunique(dropna=True)
            if n <= max_cats:
                vals = sorted(map(str, s.dropna().unique()))
                line += f" — one of: {vals}"
            else:
                line += f" — free text, {n} distinct"
        elif pd.api.types.is_numeric_dtype(s):
            line += f" — range {s.min()} to {s.max()}"
        elif pd.api.types.is_datetime64_any_dtype(s):
            line += f" — {s.min().date()} to {s.max().date()}"
        if s.isna().any():
            line += f", {s.isna().mean():.1%} null"
        lines.append(line)
    return "\n".join(lines)

Enumerating the categories is what turns df[df.status == "refunded"] into df[df.status == "REFUND"] when that is what the column actually holds. Including the date range does the same job for “last quarter”, which is otherwise resolved against whatever the model believes today’s date to be. Where the types themselves are uncertain, settle them first with the approach in LLM schema inference from a CSV.

Validating before executing

Generated code is untrusted input. It arrives from a probabilistic process, it may have been influenced by text in the data itself, and eval on it gives whatever produced it the privileges of your process. Parse it into an abstract syntax tree and walk the tree against an allowlist — string matching on forbidden substrings is not a control, since getattr(__builtins__, "ev" + "al") defeats it trivially.

import ast

ALLOWED_CALLS = {
    "sum","mean","median","min","max","count","nunique","std","var",
    "groupby","agg","sort_values","head","tail","reset_index","round",
    "value_counts","isin","between","str","dt","fillna","dropna",
    "astype","abs","query","loc","iloc","merge","pivot_table","size",
}
FORBIDDEN_NODES = (ast.Import, ast.ImportFrom, ast.FunctionDef,
                   ast.ClassDef, ast.Lambda, ast.Delete, ast.Global,
                   ast.While, ast.For, ast.Try, ast.With)

def validate(code: str, columns: set[str]) -> str | None:
    try:
        tree = ast.parse(code, mode="eval")
    except SyntaxError as e:
        return f"does not parse: {e}"
    for node in ast.walk(tree):
        if isinstance(node, FORBIDDEN_NODES):
            return f"disallowed construct: {type(node).__name__}"
        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
            return f"dunder access: {node.attr}"
        if isinstance(node, ast.Call):
            fn = node.func
            name = fn.attr if isinstance(fn, ast.Attribute) else getattr(fn, "id", None)
            if name not in ALLOWED_CALLS:
                return f"call not allowed: {name}"
        if isinstance(node, ast.Name) and node.id not in {"df", "pd"}:
            return f"unknown name: {node.id}"
    return None

Three properties make this worth the twenty lines. Parsing in mode="eval" rejects statements outright, so assignments and semicolon-chained payloads never reach the walk. Blocking attributes beginning with an underscore closes the __class__ / __subclasses__ escape that is the standard route out of a restricted Python evaluation. And restricting bare names to df and pd means no other object in the process is reachable by name.

Executing safely

Validation is necessary and not sufficient. Add three more guards. Run against a copy of the frame, so a mutation cannot corrupt the real one. Impose a wall-clock timeout, because an accidental cross-join on two large frames is allowed by every rule above and will exhaust memory. And cap the returned size, since a query returning a million rows is not an answer to a question anybody asked.

The function

1. Build the schema string with describe(df). Cache it; it changes only when the frame does.
2. Send one request with a system prompt that states the rules: return a single pandas expression over a frame named df, no imports, no assignments, no printing, no explanation, and use only column names from the schema. Set temperature to 0 — there is one right query and no reason to sample.
3. Strip any markdown fence from the response and take the expression. Models return fenced code even when told not to; handle it rather than fighting it.
4. Run validate(code, set(df.columns)). On failure, send the error text back for one retry — a column-name error is corrected reliably on the second attempt — and fail closed after that. Do not loop indefinitely.
5. Execute with eval(code, {"__builtins__": {}}, {"df": df.copy(), "pd": pd}) under a timeout, and truncate the result to a row cap.
6. Return the result and the code. Displaying the query beside the answer is the only mechanism by which a user can catch a query that ran cleanly and answered the wrong question.

The failures you will actually hit

• Ambiguous questions produce confident wrong queries. “Top customers” is by revenue, by order count or by margin, and the model picks one silently. Have it emit the interpretation in a comment above the expression, and show that too.
• Relative dates. “Last quarter” needs today’s date, which the model does not reliably know. Put the current date in the prompt explicitly and require an absolute range in the generated code.
• Nulls change every aggregate. Pandas skips NaN in mean() by default, so an average over a column that is 30% null is an average of the other 70% and nobody said so. Surface the null rate next to the answer.
• The query is valid, runs, and answers a different question. This is the dominant failure and no validator catches it, because nothing is wrong with the code. The only defence is showing the query. Treat that as a requirement, not a debug feature.
• Prompt injection through the data. If the schema includes enumerated category values, those values came from the table and can contain instructions. The allowlist is the control that makes this survivable: a hostile string can change what the model writes, but it cannot make the validator accept an import.