DEV Community

Ayi NEDJIMI
Ayi NEDJIMI

Posted on

Building a Document Classification System with LLMs

You have thousands of support tickets, contracts, or incident reports landing in a single queue. Someone needs to route them — to the right team, the right priority tier, or the right archive bucket. A traditional ML classifier needs labeled training data you probably don't have. A language model already understands text; the question is how to turn that understanding into a reliable, auditable pipeline.

This article walks through building a practical document classifier backed by a language model API, with structured output, uncertainty handling, and batch processing.

Why Not a Fine-Tuned Classifier?

Fine-tuning a BERT-style model is still the right answer if you have 10,000+ labeled examples and need sub-50ms latency. But for most teams the situation looks different: you have a taxonomy of 10–50 categories, a handful of examples per category, and no one available to label thousands of documents.

A language model handles this with zero-shot or few-shot classification. You describe the categories in plain text, optionally provide 1–3 examples each, and the model classifies. The tradeoff: higher per-request cost and 150–500ms latency. For batch processing and low-frequency classification pipelines, that's acceptable.

Prompt Structure and Structured Output

The single most important decision is forcing structured output. Don't ask the model to explain its classification in prose — ask it to return JSON. Most LLM APIs support constrained output (JSON mode or tool calls) that eliminates malformed responses entirely.

Here's a minimal Python implementation:

import json
import os
from openai import OpenAI

client = OpenAI()
MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")

CATEGORIES = {
    "billing": "Invoice issues, payment failures, subscription changes",
    "technical": "Bugs, crashes, integration errors, API failures",
    "security": "Suspected breaches, unusual logins, credential compromise",
    "general": "Anything that does not fit the above categories",
}

def classify_document(text: str) -> dict:
    category_block = "\n".join(
        f"- {name}: {desc}" for name, desc in CATEGORIES.items()
    )

    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a document classifier. "
                    "Return only valid JSON with keys: "
                    "category (string), confidence (float 0.0-1.0), reasoning (string)."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Categories:\n{category_block}\n\n"
                    f"Document:\n{text}\n\n"
                    "Classify this document."
                ),
            },
        ],
        response_format={"type": "json_object"},
        temperature=0,
    )

    result = json.loads(response.choices[0].message.content)

    # Normalize unexpected category values
    if result.get("category") not in CATEGORIES:
        result["category"] = "general"

    return result
Enter fullscreen mode Exit fullscreen mode

Setting temperature=0 is critical — classification is a deterministic task, not a creative one. The response_format constraint means you'll never get a JSON parse error from a malformed response.

Batch Processing with Async

One document at a time is fine for prototypes. Production workflows need throughput. The pattern below uses asyncio with a semaphore to cap concurrent API calls and avoid hitting rate limits:

import asyncio
from openai import AsyncOpenAI

async_client = AsyncOpenAI()

async def classify_batch(
    documents: list[dict],
    max_concurrent: int = 10,
) -> list[dict]:
    semaphore = asyncio.Semaphore(max_concurrent)

    async def _classify_one(doc: dict) -> dict:
        async with semaphore:
            try:
                response = await async_client.chat.completions.create(
                    model=MODEL,
                    messages=[
                        {
                            "role": "system",
                            "content": "Classify to JSON: {category, confidence, reasoning}",
                        },
                        {
                            "role": "user",
                            "content": (
                                "Categories: billing, technical, security, general\n\n"
                                f"Document:\n{doc['text']}"
                            ),
                        },
                    ],
                    response_format={"type": "json_object"},
                    temperature=0,
                )
                clf = json.loads(response.choices[0].message.content)
                return {**doc, "classification": clf, "error": None}
            except Exception as e:
                return {**doc, "classification": None, "error": str(e)}

    return await asyncio.gather(*[_classify_one(d) for d in documents])


if __name__ == "__main__":
    sample_docs = [
        {"id": 1, "text": "My card was charged twice for the same invoice."},
        {"id": 2, "text": "The /webhook endpoint returns 500 on every POST request."},
        {"id": 3, "text": "An unknown IP from Eastern Europe changed our admin credentials at 3am."},
    ]

    results = asyncio.run(classify_batch(sample_docs))
    for r in results:
        clf = r["classification"]
        if clf:
            print(f"[{r['id']}] {clf['category']} ({clf['confidence']:.2f}) — {clf['reasoning'][:60]}")
        else:
            print(f"[{r['id']}] ERROR: {r['error']}")
Enter fullscreen mode Exit fullscreen mode

With 10 concurrent requests, you can typically process 500–700 documents per minute. Errors are caught per-document so a single API failure doesn't abort the entire batch.

Routing on Confidence

The model's confidence field is not statistically calibrated, but it's a reliable signal for routing decisions:

  • ≥ 0.85: auto-classify and route immediately
  • 0.65–0.84: auto-classify, but flag for periodic spot-check
  • < 0.65: send to human review queue

For security-related documents, lower the auto-route threshold — a misclassified security incident has real consequences. If you're building a triage pipeline, validate the approach against your organization's security incident handling checklist before pushing to production.

You can also extend the schema with a secondary_category field for documents that straddle two categories. A password reset failure is both a billing and a security event; downstream routing logic can then decide which queue takes priority based on business rules.

Evaluating Before You Ship

Label 200–400 documents manually before going live. Compute precision and recall per category, not just overall accuracy — a classifier that routes 90% of documents to "general" looks decent on aggregate metrics but is useless in practice.

Common failure modes to anticipate:

Overlapping category definitions. If two categories share similar descriptions, the model will be inconsistent. Rewrite them to be mutually exclusive and add counter-examples to the prompt for ambiguous cases.

Short documents. One- or two-sentence texts give the model less signal. Concatenate subject lines, sender metadata, or document headers into the classification input before sending it to the model.

Domain drift. Customer phrasing evolves over time. A support ticket about "my plan auto-renewed" looks different from "my subscription charged me unexpectedly," even though they map to the same category. Monitor per-category error rates in production and refresh your category descriptions when drift exceeds a threshold — typically quarterly for stable domains.

Track classification latency and error rates alongside the model metrics. An API timeout that silently falls back to "general" is a reliability bug, not just a model quality problem.

The Takeaway

LLM-based document classification gets you from zero to a working pipeline in a day, without labeled training data. The decisions that matter most are: force structured output so you never parse malformed responses, set temperature to 0 for determinism, write non-overlapping category descriptions, and handle low-confidence outputs explicitly with a routing tier rather than hiding uncertainty.

The per-document cost will exceed a fine-tuned model, but the time saved on data labeling usually justifies it for teams under 10 engineers. Once production traffic generates a labeled dataset, you can use it to train a cheaper, faster model if your volume eventually demands it.


I run AYI NEDJIMI Consultants, a cybersecurity consulting firm. We publish free security hardening checklists — PDF and Excel.

Top comments (0)