AI Data Analysis Assistant
Project 5: Build an AI Data-Analysis Assistant
This project builds an assistant that answers questions about a tabular dataset by writing and running its own analysis code, using the Code Interpreter tool from Unit 17. The core value of this approach over asking the model to reason about data in plain text is that the model computes exact answers — real aggregations, real statistics, real chart data — rather than approximating them from a text description of the data.
Scope and Design Decisions
The assistant accepts a CSV file and a natural-language question, and returns a text answer plus, when relevant, a generated chart. It is scoped to single-dataset analysis in one session rather than a multi-dataset warehouse — joining across files or connecting to a live database is a natural extension, not part of the base project.
Two decisions shape the implementation:
- The dataset is uploaded once and reused across a session's questions. Re-uploading the file on every question wastes both time and tokens; the container that Code Interpreter runs in persists for the lifetime of the conversation thread, so the file is attached once.
- Generated code and outputs are surfaced to the caller, not hidden. For a data-analysis tool, showing what computation actually produced an answer is a trust and debugging requirement, not an optional nicety — a user needs to be able to check that the model filtered the right column before trusting a number.
Uploading the Dataset and Running the First Analysis
from openai import OpenAI
client = OpenAI()
def create_analysis_session(csv_path: str) -> str:
with open(csv_path, "rb") as f:
uploaded_file = client.files.create(file=f, purpose="assistants")
response = client.responses.create(
model="gpt-5.6-terra",
tools=[{"type": "code_interpreter", "container": {"type": "auto", "file_ids": [uploaded_file.id]}}],
input=[{
"role": "user",
"content": "A CSV file has been attached. Load it and report the column "
"names, row count, and data types, without further analysis yet.",
}],
)
container_id = _extract_container_id(response)
return container_id
def _extract_container_id(response) -> str:
for item in response.output:
if item.type == "code_interpreter_call":
return item.container_id
raise RuntimeError("No code interpreter call found in the response.")
container.type: "auto" tells the API to provision a sandboxed execution container automatically and load the attached file into it; the returned container_id is what makes the container reusable across subsequent calls instead of starting a fresh, file-less container on every question. The first call intentionally asks only for a structural summary — column names, row count, dtypes — rather than jumping straight into analysis, which both confirms the file loaded correctly and gives the calling application useful metadata (for example, to populate a UI showing available columns) before any real analysis is requested.
Note: Code Interpreter container lifetime and expiration policy are managed by OpenAI and can change; a long-running analysis session should handle the case where a container has expired by re-uploading the file and creating a new one, rather than assuming a container ID remains valid indefinitely.
Asking Analytical Questions Against a Persistent Container
def ask_data_question(container_id: str, question: str) -> dict:
response = client.responses.create(
model="gpt-5.6-terra",
tools=[{"type": "code_interpreter", "container": {"type": "auto", "id": container_id}}],
input=[{"role": "user", "content": question}],
)
generated_code = []
chart_file_ids = []
for item in response.output:
if item.type == "code_interpreter_call":
generated_code.append(item.code)
for output in getattr(item, "outputs", []) or []:
if output.get("type") == "image":
chart_file_ids.append(output["file_id"])
return {
"answer": response.output_text,
"code_run": generated_code,
"chart_file_ids": chart_file_ids,
}
Reusing container.id (rather than container.type: "auto" with a fresh file_ids list) is what keeps this a persistent analysis session — the dataset and any intermediate variables the model created in earlier turns remain available, so a follow-up question like "now break that down by region" does not need to restate what "that" refers to from scratch. The function collects both the generated code and any chart file IDs the model produced, because a data-analysis assistant that hides its computation is far less trustworthy than one that shows its work — the returned code_run list lets the caller display exactly what pandas or matplotlib code the model executed to arrive at its answer.
Downloading Generated Charts
def save_charts(chart_file_ids: list[str], output_dir: str) -> list[str]:
import os
os.makedirs(output_dir, exist_ok=True)
saved_paths = []
for i, file_id in enumerate(chart_file_ids):
content = client.files.content(file_id)
path = os.path.join(output_dir, f"chart_{i}.png")
with open(path, "wb") as f:
f.write(content.read())
saved_paths.append(path)
return saved_paths
Charts the model generates inside the sandbox exist only as files within that container until explicitly retrieved through the Files API — client.files.content downloads the raw image bytes. This two-step retrieval (get the file ID from the response, then separately download its content) is a pattern worth internalizing: several OpenAI tools produce artifacts by reference rather than embedding raw bytes directly in the response payload, keeping response objects lightweight even when the underlying artifact is large.
Guarding Against Unbounded Analysis Requests
FORBIDDEN_PATTERNS = ["requests.", "urllib", "socket.", "subprocess", "os.system"]
def validate_question_scope(question: str) -> None:
lowered = question.lower()
if any(pattern in lowered for pattern in FORBIDDEN_PATTERNS):
raise ValueError(
"This question appears to request network or system access, which "
"is outside the scope of this data-analysis assistant."
)
def ask_data_question_safely(container_id: str, question: str) -> dict:
validate_question_scope(question)
return ask_data_question(container_id, question)
The Code Interpreter sandbox already isolates code execution from the host system, so this check is not a security boundary in itself — it is a scope guard that catches obviously out-of-domain requests before spending a model call on them. Real security for the execution environment comes from the sandbox itself (Unit 17 and Unit 23 cover this in more depth); validate_question_scope exists purely to keep the assistant focused on data analysis and to fail fast on requests that are clearly not what the tool is for.
Testing the Extraction Logic
class FakeCodeCall:
type = "code_interpreter_call"
def __init__(self, code, outputs):
self.code = code
self.outputs = outputs
class FakeResponse:
def __init__(self, output, text):
self.output = output
self.output_text = text
def test_extracts_code_and_chart_ids(monkeypatch_create):
fake_call = FakeCodeCall(
code="df.groupby('region')['sales'].sum().plot(kind='bar')",
outputs=[{"type": "image", "file_id": "file-abc123"}],
)
fake_response = FakeResponse(output=[fake_call], text="Sales by region are shown above.")
monkeypatch_create(lambda **kwargs: fake_response)
result = ask_data_question("container-123", "Show sales by region")
assert result["answer"] == "Sales by region are shown above."
assert result["chart_file_ids"] == ["file-abc123"]
assert "groupby" in result["code_run"][0]
print("PASS: code and chart file IDs are correctly extracted from the response")
def test_validate_question_scope_blocks_network_access():
try:
validate_question_scope("Use requests.get to fetch external data and merge it in")
assert False, "expected ValueError"
except ValueError:
print("PASS: out-of-scope network request is rejected before a model call")
def _make_monkeypatch():
original = client.responses.create
def apply(fn):
client.responses.create = fn
def restore():
client.responses.create = original
return apply, restore
monkeypatch_create, restore_create = _make_monkeypatch()
test_extracts_code_and_chart_ids(monkeypatch_create)
restore_create()
test_validate_question_scope_blocks_network_access()
The fake FakeCodeCall and FakeResponse classes mimic just enough of the real response shape — a type attribute, code, and outputs — to exercise ask_data_question's extraction logic without a real sandbox execution. This is valuable specifically because Code Interpreter calls are slow and cost real compute; a test suite that could only validate this parsing logic against live calls would be far more expensive to run on every change.
Extending This Project
Add multi-file support so the assistant can join two related CSVs (for example, orders and customers) inside the same container, and add a result-caching layer keyed on a hash of the question and container ID so repeated identical questions do not re-run the same analysis.
Common Mistakes
- Re-uploading the dataset on every question. This discards the persistent container's state and is unnecessary — reuse the container ID from the first call for the rest of the session.
- Hiding the generated code from the end user. Data-analysis answers are only as trustworthy as the computation behind them; surfacing the code the model ran is what lets a user catch a wrong filter or a misinterpreted column.
- Treating scope-guard string matching as a security control.
validate_question_scopenarrows what the assistant attempts; it is not a substitute for the sandbox's own isolation, and should never be relied on as the only defense against unsafe code execution.
Best Practices
- Reuse the sandbox container across a session's questions. It preserves loaded data and intermediate state, and avoids redundant re-uploads.
- Always retrieve and expose both the answer and the code that produced it. This is what separates a data-analysis assistant from a black box that produces numbers no one can verify.
- Download generated artifacts (charts, exported files) explicitly rather than assuming they persist indefinitely. Container-scoped files should be retrieved and stored durably as soon as they are needed beyond the current session.