Build a data analysis agent

pip install boxlite openai

import asyncio
import re
from pathlib import Path

from boxlite import CodeBox
from openai import OpenAI

client = OpenAI()  # reads OPENAI_API_KEY (and OPENAI_BASE_URL if set)


def extract_code(text: str) -> str:
    """Pull the code out of a fenced block in the model's reply."""
    match = re.search(r"```(?:python)?\n(.*?)```", text, re.S)
    return (match.group(1) if match else text).strip()


async def analyze(csv_path: str, question: str) -> str:
    # 1) Send only the header and three sample rows as a schema hint
    lines = Path(csv_path).read_text().splitlines()
    schema_preview = "\n".join(lines[:4])

    response = client.chat.completions.create(
        model="<YOUR_MODEL>",  # e.g. "gpt-4o-mini", or any model id your endpoint serves
        messages=[{"role": "user", "content": (
            "You write Python that reads /data/sales.csv with pandas and prints answers. "
            "No explanation, code only.\n"
            f"CSV header and sample rows:\n{schema_preview}\n"
            f"Task: {question}"
        )}],
    )
    code = extract_code(response.choices[0].message.content)

    # 2) disk_size_gb=4 leaves room for pandas and numpy
    async with CodeBox(disk_size_gb=4) as box:
        # 3) Copy the real data into a persistent path — not /tmp, which is tmpfs
        await box.copy_in(csv_path, "/data/sales.csv")
        # 4) Install dependencies inside the box
        await box.install_packages("pandas")
        # 5) Run the generated code; run() returns stdout as a string
        return await box.run(code)


async def main() -> None:
    # Sample data; in production this is the file your user uploaded
    Path("sales.csv").write_text(
        "date,region,product,units,revenue\n"
        "2026-01-05,North,Widget,120,2400\n"
        "2026-01-05,South,Widget,90,1800\n"
        "2026-01-06,North,Gadget,40,2000\n"
        "2026-01-06,South,Gadget,75,3750\n"
        "2026-01-07,North,Widget,150,3000\n"
        "2026-01-07,South,Gadget,60,3000\n"
        "2026-01-08,North,Gadget,55,2750\n"
        "2026-01-08,South,Widget,110,2200\n"
    )

    try:
        answer = await analyze(
            "sales.csv",
            "Each region's total revenue and the single best-selling product by units.",
        )
        print(answer)
    except RuntimeError as exc:
        print(f"sandbox failed to start: {exc}")


asyncio.run(main())

import asyncio

from boxlite import CodeBox

CHART_CODE = """
import pandas as pd
import matplotlib
matplotlib.use("Agg")            # required: there is no display inside the box
import matplotlib.pyplot as plt

df = pd.read_csv("/data/sales.csv")
df.groupby("region")["revenue"].sum().plot(kind="bar")
plt.title("Revenue by Region")
plt.tight_layout()
plt.savefig("/data/revenue.png")  # persistent path, so copy_out can reach it
print("chart written")
"""


async def main() -> None:
    try:
        async with CodeBox(disk_size_gb=4) as box:
            await box.copy_in("sales.csv", "/data/sales.csv")
            await box.install_packages("pandas", "matplotlib")
            print(await box.run(CHART_CODE))            # -> "chart written"
            await box.copy_out("/data/revenue.png", "revenue.png")
        print("chart saved to ./revenue.png")
    except RuntimeError as exc:
        print(f"sandbox failed to start: {exc}")


asyncio.run(main())

export OPENAI_API_KEY="<YOUR_API_KEY>"
python data_analysis.py

Region Revenue:
North: 10150
South: 10750

Best-selling Products:
North: Widget
South: Widget