Smarter Jupyter Notebook Setup in CodeGrade: Using AI to Draft Your Notebook Checks

Jupyter Notebooks put explanation, data and code in one place, and CodeGrade grades them cell by cell. Writing the check for every cell is the part that still takes time.

AI can draft those checks. Every notebook cell carries its own code and output, so you can give the AI exactly the context it needs and get a targeted check back in seconds. You decide what counts and what it is worth.

Three ways to use AI when setting up notebook assignments:

• Assertion cells that check a student's variables and results
• The JSON cell trick for perfectly targeted checks
• Hardening your checks so they are fair across all submissions

This builds on our earlier post on AI prompts for Semgrep rules and unit tests.

How notebook grading works in CodeGrade

Students work on a template notebook in the browser. Grading happens with Run In Notebook steps: each one adds a checking cell to the student's notebook and runs it. Set the step to Run after notebook, and the check runs after all of the student's cells, so every variable and DataFrame the student created is there. A few plain assert statements decide pass or fail. Each check connects to a rubric row, and students see the result the moment they submit.

Two prerequisites:

• The template names the variables. Checks look for avg_price, so the task must tell students to use that name. A template cell with avg_price = ... does this.
• The Setup block installs the packages. Everything the notebook imports must be installed in the AutoTest Setup with python3 -m pip install --break-system-packages pandas numpy (adjust the list). AI drafts that install script too.

New to notebook grading? Start with the Jupyter Notebook autograding guide or the blog post How to auto-grade Jupyter Notebooks.

A notebook check is a short Python cell with assertions. That is exactly what AI drafts well, given the right context.

Prompt 1: Assertion cells from the task and your solution

Use this when you have the exercise text and your solution cell.

Write a single Python cell that checks a student's work in a Jupyter
notebook. The cell runs AFTER all of the student's cells, so the
student's variables already exist.

The task the student had to do:
[paste the exercise text]

My reference solution cell:
[paste your solution code]

Requirements for the checking cell:
- use plain assert statements, no test framework
- every assert must have a clear, student-friendly message that says
 WHAT is wrong, not what the correct answer is
- check the results (variables, DataFrame shapes and values), not the
 student's code style
Output only the code cell.

Two lines in that prompt matter most. Checking results keeps grading fair: students can reach the same DataFrame in different ways. Messages that do not reveal the answer mean you can show the assertion output to students as feedback.

Prompt 2: The JSON cell trick

A .ipynb file is JSON. Every cell is a small JSON object with the exact source code and, if the notebook has been run, the exact output. Instead of describing a cell to the AI, paste the cell itself.

In JupyterLab or VS Code, select the cell in your solution notebook, copy it, and paste it into the AI chat. You get this:

{
 "cell_type": "code",
 "source": [
   "sales_2024 = df[df['year'] == 2024]\n",
   "avg_price = sales_2024['price'].mean().round(2)\n",
   "avg_price"
 ],
 "outputs": [
   { "data": { "text/plain": ["117.36"] } }
 ]
}


Then:

Below is one cell from my solution Jupyter notebook, in raw .ipynb JSON,
including its output.

[paste the cell JSON here]

Write a single Python checking cell for the student version of this
notebook. It runs after all of the student's cells. It must:
- assert that the variables this cell creates exist
- assert their values match the output shown in the JSON (use round()
 or a tolerance for floats)
- use clear assert messages that do not reveal the expected values
Output only the code cell.


The AI returns:

assert 'sales_2024' in dir(), "Create a variable called sales_2024 with the 2024 rows."
assert 'avg_price' in dir(), "Create a variable called avg_price with the average price."
assert hasattr(sales_2024, 'shape'), "sales_2024 should be a DataFrame."
assert (sales_2024['year'] == 2024).all(), "sales_2024 still contains rows from other years."
assert len(sales_2024) == 1842, "sales_2024 does not have the expected number of rows."
assert abs(float(avg_price) - 117.36) < 0.01, "avg_price is not the average price of the 2024 sales."


The AI saw the real variable names and the real output, so nothing is guessed. Drop the cell into a Run after notebook step, connect the rubric row, done. The same works for DataFrames, plots stored in variables and trained models.

One caveat: a check at the end of the notebook sees the final value of each variable. If a later exercise reuses avg_price, set that check to run after the specific cell instead.

Optional: shorter, consistent feedback with cg-feedback-helpers

A failed assert shows students a Python traceback with the message at the bottom. CodeGrade's own package, cg-feedback-helpers, turns that into one line. It has ready-made assertions such as variable_exists, equals_float, has_length and is_of_type, each with a clear default message, so your checks read the same across the whole notebook and you write fewer messages yourself. Add this line to the Setup block:

python3 -m pip install --break-system-packages cg-feedback-helpers


Then add one line to either prompt:

Use the pytest_asserter from cg_feedback_helpers for all assertions,
pass globals() to variable_exists, and round floats before equals_float.


The same check then comes back as:

import warnings; warnings.filterwarnings("ignore")
from cg_feedback_helpers import pytest_asserter as a

a.variable_exists("sales_2024", globals())
a.variable_exists("avg_price", globals())
a.has_attr("shape", sales_2024, negative_feedback="sales_2024 should be a DataFrame.")
a.is_true((sales_2024["year"] == 2024).all(), negative_feedback="sales_2024 still contains rows from other years.")
a.has_length(sales_2024, 1842)
a.equals_float(round(float(avg_price), 2), 117.36, negative_feedback="avg_price is not the average price of the 2024 sales.")


On a failure the student sees a single line, for example FeedbackAssertionError: Expected variable 'avg_price' to exist, or your own text when you pass negative_feedback. When a check passes, any positive_feedback you added is printed. The warnings line on top hides an IPython notice that would otherwise appear under the message.

We tested this in a notebook step, and three details matter. Use pytest_asserter, not the package's default asserter: the default one reports through a channel the notebook step does not display, so a failing check only shows FeedbackAssertionError: 1. Do not call emit_success() in a notebook cell: it ends the cell with SystemExit: 0 and the step counts as failed even though every check passed. And equals_float uses a very tight tolerance, so round the value first, as in the example.

Prompt 3: Harden the checks before real submissions arrive

First drafts are too strict in some places (exact float equality, column order) and overly tolerant in others (no length check). Once you have a working set:

Here are the checking cells for my notebook assignment:
[paste your checking cells]

Improve them for grading real student work:
- floats compared with round() or a tolerance, never exact equality
- DataFrame checks tolerant to row order (sort before comparing)
- accept equivalent correct answers where they exist
- add a length/shape check wherever values are checked
- keep every assert message student-friendly and answer-free
Output the complete revised cells only.


Then the one habit that catches most mistakes: run your solution and your empty template against the finished checks. The solution must pass everything; the template must fail everything. If the template passes a check, that check is not checking the student's work. Build a snapshot to verify both. If a check misbehaves, paste the AutoTest output back into the AI and ask for a fix.

If the notebook uses random data, seed it in the template (np.random.seed(...)) so the expected values hold for every student.

Putting it together

1. Template and solution first. The template names the variables students must use; the Setup block installs the packages.
2. One check per graded cell. Prompt 1 or the JSON cell trick per exercise, each into a Run In Notebook step set to run after the notebook, connected to its rubric row.
3. Harden and verify. Prompt 3 over the full set, then build a snapshot: solution passes, template fails.

For a notebook with ten exercises this takes about half an hour, most of it reading what the AI drafted. The pedagogy stays yours: what each exercise teaches, what it is worth, how strict the checks are. The AI writes the assertions.

Questions about notebook grading or AI in your workflow? Our team is here to help at support@codegrade.com.