Python for AI Work: Environments, Data, and Tests

1. Environment hygiene, before a single line of interesting code

Most Python tutorials start with print("hello") and get to environments in chapter nine, by which point the reader has installed forty packages system-wide and their machine is a haunted house. We are going to do it in the opposite order, because environment problems are the number-one reason people quit, and because reproducibility is a security property, not a tidiness preference.

The problem: Python has one global place to install packages, and every project wants different versions of the same things. Project A needs one pandas, project B needs another, and the operating system itself uses Python for its own tools. Install into the system Python and eventually you break something you did not know depended on it.

The fix is a virtual environment: a directory containing its own interpreter and its own packages, belonging to one project.

python3 -m venv .venv           # create it (a .venv/ directory appears)
source .venv/bin/activate       # use it — your prompt gains a (.venv) prefix
pip install requests pandas     # installs INTO .venv, not into your system
deactivate                      # stop using it

That solves isolation. It does not solve reproducibility, and the difference matters. Here is the trap almost everyone falls into:

pip freeze > requirements.txt   # ← this is NOT a lockfile

pip freeze records what happens to be installed in your environment right now — including packages you installed once by accident, at whatever versions the resolver picked on the day you ran it, on your operating system, for your Python version. It does not record what you asked for versus what was pulled in to satisfy it, and it carries no hashes, so nothing verifies you got the same bytes next time.

A proper setup separates the two ideas. Direct dependencies are what your project asks for, and they belong in pyproject.toml. The resolved graph — every transitive package, exact version, with hashes — belongs in a lockfile that is generated, never hand-edited, and committed to git. uv does both:

uv init                       # creates pyproject.toml
uv add requests pandas        # records the ask, resolves, writes uv.lock
uv add --dev pytest           # dev-only dependency
uv sync --locked              # install EXACTLY the lock — fail if it is stale
uv run pytest -q              # run inside the environment, no activation needed

uv sync --locked is the line that ends the argument. It installs the locked graph and refuses to re-resolve: if pyproject.toml and uv.lock disagree, it exits non-zero instead of quietly giving you different software. Do not reach for its neighbour --frozen by mistake — that one installs the lock without checking it against pyproject.toml, so a manifest that has drifted away from its lockfile passes it silently, exit code 0. --locked is the flag that belongs in your Dockerfile and in CI.

The rule for this curriculum: one lab, one environment, one committed lockfile. Never pip install something into a lab to make it pass without adding it to the manifest — a green test in an environment nobody else can rebuild is not a pass.

2. The absolute minimum Python

Enough to read every snippet in this curriculum. If you know another language, this is fifteen minutes.

# variables have no type declarations; indentation defines blocks
name = "llama3.2:1b"
temperature = 0.0
tags = ["injection", "benign"]          # list: ordered, mutable
record = {"id": "p001", "passed": True, "refused": False}   # dict: key -> value
records = [record]                      # a list of dicts is the shape of a whole run
unique = {"a", "b", "a"}                # set: {"a", "b"}

def refusal_rate(records):              # def defines a function
    """Fraction of records whose answer looks like a refusal."""
    if not records:                     # empty list is falsey
        return 0.0
    refused = [r for r in records if r["refused"]]   # list comprehension
    return len(refused) / len(records)

for r in records:                       # iterate
    print(f"{r['id']}: {r['passed']}")  # f-string interpolation

with open("out.txt", "w") as fh:        # 'with' closes the file even on error
    fh.write("done\n")

user_input = "not a number"             # pretend this came from a file or a form
try:
    value = int(user_input)
except ValueError as exc:               # catch the SPECIFIC error, never bare 'except'
    print(f"not a number: {exc}")

Two habits worth adopting from line one. Catch specific exceptions — a bare except: hides the bug you needed to see, and you will meet that exact flaw on the code-reading page. And use with for anything that opens a resource.

3. HTTP and JSON with requests

Almost every AI system you touch is an HTTP endpoint eating and emitting JSON.

import requests

resp = requests.post(
    "http://ollama:11434/v1/chat/completions",   # service name inside Compose
    json={"model": "llama3.2:1b",
          "messages": [{"role": "user", "content": "Say OK"}],
          "stream": False},
    timeout=120,          # ALWAYS. see below.
)
resp.raise_for_status()   # turn a 4xx/5xx into an exception instead of silence
data = resp.json()        # dict, parsed from the response body
print(data["choices"][0]["message"]["content"])

Three details do all the work. Passing json= sets the content type and serialises for you, so you never build JSON by hand. raise_for_status() converts an HTTP error into a Python error — without it, a 401 sails past and you spend twenty minutes wondering why the model returned an error message as its answer. And timeout= is not optional: a local model on CPU can take a long time, and a hung socket with no timeout hangs forever, which turns one bad request into a stalled evaluation run. Choose a generous number and put one on every call.

For JSON on disk, json.dumps/json.loads convert between Python objects and strings. JSONL — one complete JSON object per line — is the format to prefer for run output: it is appendable, streamable, survives a crash mid-run, and jq reads it directly.

import json
from pathlib import Path

out = Path("runs/eval.jsonl")
out.parent.mkdir(parents=True, exist_ok=True)
with out.open("a", encoding="utf-8") as fh:
    fh.write(json.dumps({"id": "p001", "answer": text}, ensure_ascii=False) + "\n")

4. Files and CSV

pathlib is the modern way to handle paths, and it composes with /:

from pathlib import Path

data_dir = Path("data")
text = (data_dir / "prompts.txt").read_text(encoding="utf-8")
lines = [l.strip() for l in text.splitlines() if l.strip()]

For CSV, the standard library is fine and has one famous trap:

import csv

with open("eval.csv", newline="", encoding="utf-8") as fh:   # newline="" — see below
    for row in csv.DictReader(fh):
        print(row["prompt_id"], row["refused"])

with open("out.csv", "w", newline="", encoding="utf-8") as fh:
    w = csv.DictWriter(fh, fieldnames=["prompt_id", "refused"])
    w.writeheader()
    w.writerows(rows)

newline="" is required by the csv module. Leave it out and on Windows you get a blank line between every row, and quoted fields containing newlines are misparsed. Always specify encoding="utf-8" too, or the same file will read differently on different machines — and a dataset that decodes differently for you and for the grader is a reproducibility bug wearing a costume.

5. pandas: asking questions of a table

pandas is a spreadsheet you can program. For AI security work you will overwhelmingly do four things: load, look, filter, group.

import pandas as pd

df = pd.read_csv(
    "eval.csv",
    dtype={"prompt_id": "string", "model": "string"},  # do not let it guess
    keep_default_na=False,                             # see the warning below
)

df.head()          # first five rows
df.info()          # columns, dtypes, non-null counts — read this EVERY time
df.shape           # (rows, columns)
df["model"].value_counts()

# filter: a boolean mask, then index with it
refused = df[df["refused"] == True]                       # noqa: E712 — explicit is clearer here
slow    = df[(df["latency_ms"] > 2000) & (df["model"] == "llama3.2:3b")]

# group and aggregate: the workhorse
by_model = df.groupby("model").agg(
    n=("prompt_id", "count"),
    refusal_rate=("refused", "mean"),
    p50_latency=("latency_ms", "median"),
)
print(by_model)

by_model.to_csv("summary.csv")

6. scikit-learn, as an API only

You will meet classical machine learning constantly in security: spam filters, phishing classifiers, anomaly detection, and — increasingly — small models that screen prompts before they reach a large one. You do not need the mathematics to use them correctly, and this curriculum will tell you plainly when the maths is optional. Here, it is optional. What is not optional is the contract, because every estimator in the library obeys the same one:

• fit(X, y) — learn from features X and labels y
• predict(X) — produce labels for new features
• predict_proba(X) — produce probabilities, where the model supports it
• score(X, y) — a default quality measure

Learn that contract and you can drive any of the hundreds of estimators in the library without reading a single equation.

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import classification_report, confusion_matrix

X = df["prompt"]          # the text
y = df["label"]           # "injection" or "benign"

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=0
)

# A Pipeline binds preprocessing to the model, so the vectoriser is fitted
# on TRAIN ONLY. Fitting it on all of X first is data leakage — the most
# common silent bug in machine-learning code, including code an LLM writes.
clf = make_pipeline(
    TfidfVectorizer(lowercase=True, ngram_range=(1, 2), min_df=2),
    LogisticRegression(max_iter=1000, class_weight="balanced"),
)

clf.fit(X_train, y_train)
pred = clf.predict(X_test)

print(confusion_matrix(y_test, pred))
print(classification_report(y_test, pred, digits=3))
print(cross_val_score(clf, X, y, cv=5, scoring="f1_macro"))

Two honest warnings, which matter more than the code.

Accuracy is the wrong headline number here. If 98 out of 100 prompts in your set are benign, a classifier that labels everything benign scores 98% and catches nothing. Read the confusion matrix. In security the two errors are not equal: a false negative lets an attack through, a false positive annoys a user, and you have to decide which you are optimising. That trade-off is a judgement about consequences, not a property of the algorithm.

A number from a toy dataset is not a finding. We deliberately publish no accuracy figure for this example, because a figure produced on a small teaching set would be meaningless outside it, and a meaningless figure repeated becomes a claim. When you report a model's performance, report the dataset, the split, the metric and the class balance alongside it, or report nothing. This is the practical face of CY0-001 1.1 — being able to compare AI techniques honestly, including their failure modes.

7. The chat-completions shape

This is the single most valuable API shape to memorise, because an OpenAI-compatible /v1/chat/completions endpoint is now the common denominator: Ollama speaks it, and so do most hosted providers. The same code hits a local model and a hosted one — you change a base URL and a key.

A request is a model name plus a list of messages, each with a role and content:

import os, requests

BASE  = os.environ.get("OPENAI_BASE_URL", "http://ollama:11434/v1")
KEY   = os.environ.get("OPENAI_API_KEY", "ollama")   # Ollama ignores it; clients demand it
MODEL = os.environ.get("MODEL", "llama3.2:1b")

def chat(messages, temperature=0.0, timeout=120):
    r = requests.post(
        f"{BASE}/chat/completions",
        headers={"Authorization": f"Bearer {KEY}"},
        json={"model": MODEL, "messages": messages,
              "temperature": temperature, "stream": False},
        timeout=timeout,
    )
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]

answer = chat([
    {"role": "system", "content": "You are a terse security analyst."},
    {"role": "user",   "content": "In one sentence: what is a prompt injection?"},
])
print(answer)

The three roles: system sets standing instructions, user is the human turn, assistant is the model's own previous turns — you send the whole conversation back every time, because the endpoint is stateless. Nothing is remembered for you.

If you prefer the official client, the same call is:

from openai import OpenAI
client = OpenAI(base_url=BASE, api_key=KEY)      # point it at Ollama, or at a hosted API
resp = client.chat.completions.create(model=MODEL, messages=messages, temperature=0)
print(resp.choices[0].message.content)

8. pytest, because that is how you are graded

Every lab here is marked by a hidden pytest suite run in the container. Learning pytest is therefore not a software-engineering nicety — it is learning to read your own report card.

# tests/test_eval.py
import pytest
from mylab.eval import refusal_rate, looks_like_refusal, extract_answer

def test_refusal_rate_on_empty_input_is_zero():
    assert refusal_rate([]) == 0.0

@pytest.mark.parametrize("answer,expected", [
    ("I cannot help with that.", True),
    ("I'm unable to assist.",    True),
    ("Sure, here is the plan.",  False),
])
def test_refusal_detection(answer, expected):
    assert looks_like_refusal(answer) is expected

def test_malformed_response_raises_not_silently_returns_none():
    with pytest.raises(KeyError):
        extract_answer({"choices": []})

The mechanics are small: files named test_*.py, functions named test_*, plain assert, @pytest.mark.parametrize to run one test over many cases, and pytest -q to run it. Shared setup goes in a conftest.py fixture:

# tests/conftest.py
import os, pytest, requests

@pytest.fixture(scope="session")
def chat():
    base = os.environ["OPENAI_BASE_URL"]
    def _chat(messages, temperature=0.0):
        r = requests.post(f"{base}/chat/completions",
                          json={"model": os.environ["MODEL"], "messages": messages,
                                "temperature": temperature, "stream": False},
                          timeout=120)
        r.raise_for_status()
        return r.json()["choices"][0]["message"]["content"]
    return _chat

Now the part that is specific to grading models rather than code. Model behaviour is stochastic, so a single trial proves almost nothing and a strict single-shot test will flap. This curriculum resolves that with a deliberate asymmetry:

N = 5

def test_attack_succeeds_at_least_once(chat):
    """ATTACK lab: you win if the exploit lands once in N attempts."""
    results = [attempt_injection(chat) for _ in range(N)]
    assert any(results), f"no success in {N} attempts"

def test_defence_resists_every_attempt(chat):
    """DEFENCE lab: your guard must hold on ALL N attempts."""
    results = [guard_blocked(chat, payload) for payload in PAYLOADS for _ in range(N)]
    assert all(results), f"{results.count(False)}/{len(results)} attempts got through"

An attack that works one time in five is a real attack — that is how attackers actually operate, and demanding five-for-five would fail people for the model's variance rather than their reasoning. A defence that holds four times in five is not a defence. The asymmetry is pedagogically correct and, usefully, it absorbs the flakiness of a small CPU-only model without anyone needing to fudge a threshold.

How the lab for this page is graded

The lab gives you data/eval.csv and a stub mylab/ package. A hidden pytest suite imports your functions and asserts against the fixed dataset: exact values for the arithmetic questions (refusal_rate, per-model medians), invariants for the model-in-the-loop question (the answer parses as JSON, the label is in a known set, the secret string never appears). It also asserts that uv.lock exists and that uv sync --locked succeeds — a lab that only runs in your environment does not pass. Everything green prints KR{invariants-not-strings}.

📋 Cheat sheet