Teaching Your Pipeline to Think: Agentic Record Classification

Module 5, Unit 3 | Lesson 1 of 5

By the end of this lesson, you will be able to:

• Replace the stage_classify stub with a real LLM-powered classifier that labels flagged pipeline records (K12, S10, S11)
• Design a structured classification prompt that produces consistent, parseable JSON output (K12, S11)
• Implement confidence thresholds and fallback routing for low-confidence classifications (K11, S10, S11)
• Integrate the classifier into your pipeline using the retry and parallel utilities built in Unit 2 (K12, S10, B6)

In Unit 2 you built the architecture around your classifier: retry logic, parallel processing, and the human approval checkpoint. The stage_classify function has been a stub — returning "pending" for every record. In this lesson you replace that stub with a real LLM agent that reads each flagged record and returns a structured classification.

This is the moment the pipeline starts thinking.

🔑 Key term — Structured output: A model response that conforms to a specific format — typically JSON — rather than free text. Structured output is essential for pipeline integration: your code needs to parse the classification reliably, not extract it from a sentence. You achieve this through prompt engineering (explicit JSON instructions) and response validation.

Set up the OpenAI client

Install the OpenAI SDK:

pip install openai python-dotenv

Store your API key in a .env file — never in code or version control:

OPENAI_API_KEY=sk-...

Load it at the top of every script that needs it:

from dotenv import load_dotenv
load_dotenv()

Coach Cora

Put .env in your .gitignore before your first commit — not after. Committing a key is a credential disclosure, not an untidy repository: it is a live secret, billable to your organisation, published to everyone who can read the repo. Public repositories are scraped continuously by bots hunting for exactly this, and keys have been found and used within minutes.

Here is the part people get wrong. Deleting the file and committing again does not undo it. Git keeps history, so the key still sits in the earlier commit and anyone can read it by checking out that commit — and it also survives in the reflog, in every clone and fork already taken, and in caches on the hosting platform. Even rewriting history properly with git filter-repo or BFG only fixes your copy: it changes every commit hash downstream, forces everyone to re-clone, and cannot reach the copies other people already pulled.

So treat the key as compromised the moment it is pushed. Revoke it in your API dashboard and issue a new one — that is the only step that actually stops the leak. Scrubbing the history is tidying up afterwards.

And if your organisation provides an approved API route, use that instead and follow its key management process.

Design the classification prompt

Your classifier needs to label each flagged record with one of four categories:

• compliance_risk — PII exposure, data governance violation, or regulatory concern
• operational_issue — hallucination, incorrect output, or tool failure
• cost_anomaly — usage significantly above expected cost for the task type
• data_quality — missing data, inconsistent input, or corrupted record

Create classifier.py. It loads the .env itself rather than relying on whoever imports it having done so, and it builds the OpenAI client lazily — the second of those matters more than it looks, and L3.3 explains why:

from __future__ import annotations

import json
import os
from typing import Any

from dotenv import load_dotenv

load_dotenv()  # reads .env into the environment before any client is built

_client = None


def get_client():
    """Build the OpenAI client on first use, not when this module is imported."""
    global _client
    if _client is None:
        from openai import OpenAI
        _client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    return _client

The label set, and the prompt that teaches it to the model. The four names appear twice on purpose — once as data your code can check against, once in words because the model needs to know what they mean.

CLASSIFICATION_LABELS = [
    "compliance_risk",
    "operational_issue",
    "cost_anomaly",
    "data_quality",
]

SYSTEM_PROMPT = """You are a data pipeline classifier for an AI tool usage monitoring system.

Your job is to classify flagged records from an AI usage log into exactly one category.

Categories:
- compliance_risk: PII exposure, data governance violation, or regulatory concern
- operational_issue: hallucination, incorrect output, or tool failure
- cost_anomaly: usage significantly above expected cost for the task type
- data_quality: missing data, inconsistent input, or corrupted record

Always respond with valid JSON in exactly this format:
{
  "classification": "<one of the four categories>",
  "confidence": <float between 0.0 and 1.0>,
  "reasoning": "<one sentence explaining the classification>"
}

Do not include any text outside the JSON object."""

This next part is what makes the output trustworthy rather than merely likely. Read the comment on strict.

# The schema the API is required to honour. `strict: True` makes this a
# server-side guarantee, not a request: the response will parse, the label will
# be one of the four, and confidence will be a number in range.
CLASSIFICATION_SCHEMA = {
    "name": "record_classification",
    "strict": True,
    "schema": {
        "type": "object",
        "additionalProperties": False,
        "required": ["classification", "confidence", "reasoning"],
        "properties": {
            "classification": {"type": "string", "enum": CLASSIFICATION_LABELS},
            "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0},
            "reasoning": {"type": "string"},
        },
    },
}

The schema constrains the shape of the answer. Nothing constrains whether the answer is correct, so check the values yourself. Add this to classifier.py — you will move it to a shared module in L3.3, when a second provider needs the same four guards:

def validate_result(result: dict[str, Any]) -> dict[str, Any]:
    """Check the values a model returned, not just their shape."""
    if result.get("classification") not in CLASSIFICATION_LABELS:
        raise ValueError(
            f"Unknown classification: {result.get('classification')!r}. "
            f"Expected one of: {CLASSIFICATION_LABELS}"
        )
    confidence = result.get("confidence")
    if not isinstance(confidence, (int, float)) or not 0.0 <= confidence <= 1.0:
        raise ValueError("Missing or invalid confidence score")
    if not isinstance(result.get("reasoning"), str) or not result["reasoning"].strip():
        raise ValueError("Missing or invalid reasoning")
    return result

Now the call itself. The record is flattened into the user message field by field, so the model sees named values rather than a blob of JSON. The three guards after the response handle the two failures a strict schema cannot prevent — a refusal and a truncation, both of which arrive as HTTP 200 — and then hand the parsed result to validate_result:

def classify_record(record_data: dict[str, Any]) -> dict[str, Any]:
    """
    Send a single record to the LLM for classification.
    Returns a dict with classification, confidence, and reasoning.
    Raises ValueError if the response cannot be parsed or is invalid.
    """
    user_message = f"""Classify this AI usage log record:

log_id: {record_data.get('log_id')}
tool_name: {record_data.get('tool_name')}
task_type: {record_data.get('task_type')}
flagged_issue: {record_data.get('flagged_issue')}
cost_usd: {record_data.get('cost_usd')}
time_saved_mins: {record_data.get('time_saved_mins')}
quality_score: {record_data.get('quality_score')}
human_reviewed: {record_data.get('human_reviewed')}"""

    response = get_client().chat.completions.create(
        model="gpt-4.1-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
        temperature=0.1,
        max_completion_tokens=200,
        response_format={"type": "json_schema", "json_schema": CLASSIFICATION_SCHEMA},
    )
    choice = response.choices[0]

    # With a strict schema, malformed JSON is no longer the failure mode.
    # These two are: the model can decline outright, and the response can be
    # cut off mid-object by max_completion_tokens. Both arrive as HTTP 200.
    if choice.message.refusal:
        raise ValueError(f"Model refused to classify: {choice.message.refusal}")
    if choice.finish_reason == "length":
        raise ValueError("Response truncated before the JSON was complete")

    raw = choice.message.content
    if not raw:
        raise ValueError("Model returned an empty response")

    return validate_result(json.loads(raw))

🔑 Key term — Temperature: A parameter that controls how deterministic or varied the model's output is. At temperature 0 the model always picks the highest-probability token — maximally consistent. At temperature 1 it samples more broadly — more creative but less predictable. For classification tasks, use low temperature (0.0–0.2) to get consistent results on the same input.

Curious Cat

So if strict: True guarantees the shape, why keep validating in code? Two reasons, and neither is superstition. First, the schema constrains what the answer looks like, never whether it is right — compliance_risk at 0.98 confidence on a record that is really a cost anomaly satisfies every rule in that schema and is still wrong. Second, this guarantee belongs to one provider. In L3.3 you send PII records to a local model through Ollama instead. Ollama can constrain output too, but not to the same degree — you get well-formed JSON, not a promise that classification is one of your four labels or that confidence is in range. These checks are what close that gap, which is why they live in a module both providers import. Treat model output as untrusted upstream data, exactly like the raw CSV in L1.1.

Add confidence threshold and fallback

Some records are genuinely ambiguous. When the model is uncertain, route the record for human review rather than forcing a low-confidence decision into the pipeline automatically. The threshold is a policy decision, so give it a name of its own in classifier.py:

CONFIDENCE_THRESHOLD = 0.75

Then the wrapper that applies it. Nothing else in the pipeline needs to know a threshold exists — it sees a route_override and acts on it:

def classify_with_fallback(record_data: dict[str, Any]) -> dict[str, Any]:
    """
    Classify a record and apply a confidence threshold.
    Records below the threshold are flagged for human review
    rather than routed automatically.
    """
    result = classify_record(record_data)

    if result["confidence"] < CONFIDENCE_THRESHOLD:
        result["route_override"] = "human_review"
        result["route_reason"] = (
            f"Confidence {result['confidence']:.2f} below "
            f"threshold {CONFIDENCE_THRESHOLD}"
        )
    else:
        result["route_override"] = None
        result["route_reason"] = None

    return result

Replace the stage_classify stub

Open stages.py and replace the entire stage_classify function (the stub from L2.1) with the implementation below. Do not add this alongside the old version — delete the stub first.

Three functions, added to stages.py. The first is only a shape change — the flat payload of the fields the prompt asks for:

from classifier import classify_with_fallback


def record_payload(record: ToolUsageRecord) -> dict[str, Any]:
    """The fields the classifier prompt asks for, flattened."""
    return {
        "log_id": record.log_id,
        "tool_name": record.tool_name.value,
        "task_type": record.task_type,
        "flagged_issue": record.flagged_issue.value,
        "cost_usd": record.cost_usd,
        "time_saved_mins": record.time_saved_mins,
        "quality_score": record.quality_score,
        "human_reviewed": record.human_reviewed,
    }

The second classifies a single record. Clean records (flagged_issue == "none") return immediately without touching the LLM — they need no classification, and the cheapest API call is the one you do not make. Anything actually flagged goes through the retry wrapper from L2.2, so a transient API failure does not lose the record:

def classify_item(item: dict[str, Any]) -> dict[str, Any]:
    """Classify one record. Clean records skip the LLM entirely."""
    record = item["record"]

    if record.flagged_issue == FlaggedIssue.NONE:
        return {
            **item,
            "classification": "none",
            "confidence": 1.0,
            "reasoning": "No issue flagged - classification not required.",
            "route_override": None,
            "model_used": "none",
            "llm_called": False,
        }

    result = with_retry(
        lambda: classify_with_fallback(record_payload(record)),
        max_attempts=3,
        base_delay=1.0,
    )

    return {
        **item,
        "classification": result["classification"],
        "confidence": result["confidence"],
        "reasoning": result["reasoning"],
        "route_override": result.get("route_override"),
        "route_reason": result.get("route_reason"),
        "llm_called": True,
        "model_used": "openai/gpt-4.1-mini",
    }

The third is the stage. It runs the batch and records what happened. llm_calls is counted here rather than estimated later, and L3.2 reuses it alongside the per-record token totals that its cost arithmetic actually runs on:

def stage_classify(state: PipelineState) -> PipelineState:
    """Stage 3: Classify flagged records using the OpenAI classifier."""
    started = time.time()

    items = [{"record": r} for r in state.valid_records]
    results, errs = run_in_parallel(classify_item, items, max_workers=3)
    state.classified_records = results
    llm_calls = sum(1 for r in results if r.get("llm_called"))

    state.stage_results.append(StageResult(
        stage_name="classify",
        records_in=len(state.valid_records),
        records_out=len(results),
        records_failed=len(errs),
        duration_seconds=round(time.time() - started, 3),
        # run_in_parallel keeps the whole failed item, record object and all.
        # The output payload is JSON, so reduce each failure to the two facts a
        # recipient needs: which record, and what went wrong.
        errors=[
            {"log_id": e["item"]["record"].log_id, "error": e["error"]}
            for e in errs
        ],
        metadata={
            "llm_calls": llm_calls,
            "skipped_clean": len(results) - llm_calls,
            "model": "gpt-4.1-mini",
        },
    ))

    print(
        f"[classify] {len(results)} classified "
        f"({llm_calls} LLM calls, "
        f"{len(results) - llm_calls} skipped clean, "
        f"{len(errs)} errors)"
    )
    return state

Promoting classify_item to its own function rather than nesting it inside the stage is not cosmetic. L3.3 and L4.2 both rewrite it, and a function you can point at is a function you can replace wholesale.

The stage_route function from L2.2 checks route_override before applying the normal deterministic route rules. That means a low-confidence classification is not silently treated as a normal ops or standard record — it joins the human approval path for review.

Coach Cora

Trace that precedence carefully with a PII record. If the classifier returns low confidence on a pii_detected row, route_override wins, the record lands in human_review_queue — and it never enters the compliance queue at all. Both routes reach the approval checkpoint, so a human still sees it and nothing is lost today. But your compliance queue is no longer a complete list of PII records, and if anyone ever reports from that queue, the count will be quietly wrong. This is what it looks like when a model's output is allowed to override a deterministic governance rule. In L2.1 you argued that compliance routing must never depend on an LLM; here is the line of code where that argument gets tested. Two defensible fixes: check pii_detected before route_override so the compliance route always wins, or keep the current order and record the record in both queues. Pick one, and write down why — this is exactly the kind of decision the L4.3 narrative is asking you to justify.

Test the classifier in isolation

⚠️ Read this before you run anything. As written, stage_classify sends every flagged record to OpenAI — including the two pii_detected rows. That is a real governance problem, and it is deliberate: this is the naive baseline, and L3.3 is where you fix it by routing PII to a local model instead. You are seeing the wrong version first on purpose, because "send everything to the cloud API" is what most people build by default and it is worth recognising as a decision rather than an accident.

For the same reason, the isolation test below uses LOG-008 — a hallucination-flagged record — rather than one of the PII rows. Get into the habit now: when you hand-test against a live API, pick a record you would be comfortable seeing in a third party's logs.

Before running the full pipeline, test classify_record directly:

# quick_test.py
from dotenv import load_dotenv
load_dotenv()

from classifier import classify_record

result = classify_record({
    "log_id": "LOG-008",          # A real hallucination-flagged row
    "tool_name": "GPT-4",
    "task_type": "summarisation",
    "flagged_issue": "hallucination",
    "cost_usd": 0.062,
    "time_saved_mins": -5,
    "quality_score": 3.1,
    "human_reviewed": True,
})
print(result)
# Expected: {"classification": "operational_issue", "confidence": ~0.9, "reasoning": "..."}

If you get a valid classification, run the full pipeline:

python pipeline.py

The [classify] line will now show real LLM call counts. Check pipeline_output.json — the classify stage metadata should include llm_calls, skipped_clean, and the model name.

Build activity

1. Create classifier.py with the system prompt, classify_record, and classify_with_fallback.
2. Test classify_record in isolation on LOG-008 (hallucination-flagged) and LOG-011 (cost_overrun-flagged). Record the classifications and confidence scores you receive. Deliberately do not hand-test a pii_detected record — see the note above.
3. Update stage_classify in stages.py with the real implementation.
4. Run python pipeline.py end-to-end and verify the classify stage metadata in pipeline_output.json.
5. Set CONFIDENCE_THRESHOLD to 0.99 and re-run. How many records now receive route_override: "human_review" and pause at the approval checkpoint? Reset to 0.75 after testing.

Challenge Chase

The classifier currently sends each flagged record as a separate API call. For large datasets this is expensive and slow. Design a batch classification approach: group records by flagged_issue type and send them in a single prompt that asks the model to return a JSON array of classifications. What is the maximum safe batch size before you risk exceeding gpt-4.1-mini's context window? And note that batching changes the failure mode: a strict schema will still guarantee a well-formed array, so the partial failure you have to design for is not malformed JSON — it is an array that comes back with 8 entries when you sent 10, or one truncated by max_completion_tokens before it closed. How would you detect each, and how would you decide between retrying the whole batch and falling back to single calls?

The code from this lesson

The pipeline with a real agentic classifier in it. You will need your own OpenAI key: copy the .env.example in the archive to .env and put your key there. The code files are assembled from the code blocks of this lesson and the ones before it.

Take it if you would rather read and run the code than type it out. Typing it is still the better way to learn it — this is here so a typo cannot cost you the lesson.

Checklist

• classifier.py is created with the system prompt, classify_record, and classify_with_fallback
• classify_record returns a valid object with classification, confidence, and reasoning
• Low-confidence records receive route_override: "human_review" rather than a forced classification
• stage_classify uses the real classifier with retry and parallel processing from Unit 2
• Clean records skip the LLM call entirely
• python pipeline.py shows real LLM call counts in the [classify] output
• pipeline_output.json stage metadata includes llm_calls, skipped_clean, and model

Quick Check

KSB evidence focus

• K12 — Understands how to set up, interact with and generate APIs, databases and spreadsheets. You have completed a full LLM API integration: client setup, structured prompting, response parsing, output validation, error handling, and retry logic. This is a reusable pattern for any OpenAI-compatible API your pipeline needs to call.

• S11 — Can design, implement and test an AI/ML model or system. The classifier is a designed system with deliberate choices at every layer: the prompt specifies the task, temperature controls variance, Pydantic-style validation enforces the output contract, and a confidence threshold handles uncertainty. Each decision is documented and testable.

• S10 — Can carry out data processing tasks, including processing/cleaning, data transformations, and feature engineering. Routing clean records away from the LLM and only classifying flagged ones is a data processing optimisation that reduces API cost and latency without changing the pipeline output for clean records.

Up next: Every one of those API calls costs money, and right now you have no idea how much. Lesson 2 shows you how to read token usage straight out of the API response, turn it into a real cost figure, and put a guardrail on it before a production data volume turns a rounding error into a budget.