Document Classification and Routing

Why Document Classification Matters

LLM-Based Classification

import openai
import os

client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

DOC_TYPES = ['invoice', 'contract', 'report', 'email', 'receipt', 'form', 'letter', 'other']

CLASSIFY_PROMPT = '''Classify this document into one of these types: {types}

Document excerpt (first 1000 characters):
{text}

Respond with ONLY the document type as a single word from the list above.'''

def classify_with_llm(text):
    response = client.chat.completions.create(
        model='gpt-4o-mini',  # cheap and fast for classification
        messages=[{
            'role': 'user',
            'content': CLASSIFY_PROMPT.format(
                types=', '.join(DOC_TYPES),
                text=text[:1000]
            )
        }],
        max_tokens=10,
        temperature=0
    )
    predicted = response.choices[0].message.content.strip().lower()
    return predicted if predicted in DOC_TYPES else 'other'

Rule-Based Classification Fallback

CLASSIFICATION_RULES = {
    'invoice': [
        'invoice', 'invoice number', 'bill to', 'amount due',
        'total amount', 'tax invoice', 'payment terms'
    ],
    'contract': [
        'agreement', 'terms and conditions', 'hereby agrees',
        'party a', 'party b', 'whereas', 'obligations'
    ],
    'report': [
        'executive summary', 'quarterly report', 'annual report',
        'findings', 'recommendations', 'methodology'
    ],
    'email': ['from:', 'to:', 'subject:', 'date:', 'dear ', 'regards,'],
    'receipt': ['receipt', 'thank you for your purchase', 'transaction id', 'cashier']
}

def classify_with_rules(text):
    text_lower = text.lower()
    scores = {}
    for doc_type, keywords in CLASSIFICATION_RULES.items():
        score = sum(1 for kw in keywords if kw in text_lower)
        if score > 0:
            scores[doc_type] = score
    if not scores:
        return None  # no match — fall through to LLM
    return max(scores, key=scores.get)

if __name__ == '__main__':
    demo_text = 'INVOICE\nBill To: Acme Corp\nAmount Due: $500\nPayment Terms: Net 30'
    print('Classified as:', classify_with_rules(demo_text))

Tiered Classification Strategy

def classify_document(text):
    # Tier 1: rule-based (free, fast)
    result = classify_with_rules(text)
    if result:
        print(f'Rule-based classification: {result}')
        return {'type': result, 'method': 'rules', 'confidence': None}

    # Tier 2: LLM (accurate, slower)
    result = classify_with_llm(text)
    print(f'LLM classification: {result}')
    return {'type': result, 'method': 'llm', 'confidence': None}

Confidence Thresholds

import json

CLASSIFY_WITH_CONFIDENCE_PROMPT = '''Classify this document. Return JSON:
{{"type": "invoice", "confidence": 0.95, "reason": "Contains invoice number and payment terms"}}

Valid types: {types}
Document excerpt: {text}

JSON:'''

def classify_with_confidence(text):
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': CLASSIFY_WITH_CONFIDENCE_PROMPT.format(
            types=', '.join(DOC_TYPES),
            text=text[:1000]
        )}],
        temperature=0
    )
    try:
        result = json.loads(response.choices[0].message.content)
        return result
    except json.JSONDecodeError:
        return {'type': 'other', 'confidence': 0.0, 'reason': 'Parse error'}

LOW_CONFIDENCE_THRESHOLD = 0.6

def classify_and_check(text):
    result = classify_with_confidence(text)
    if result['confidence'] < LOW_CONFIDENCE_THRESHOLD:
        result['needs_review'] = True
        print(f'Low confidence ({result["confidence"]}) — flagging for review')
    return result

Document Routers

def handle_invoice(text):
    # Extract: vendor, invoice number, total, due date
    extract_prompt = f'''Extract from this invoice (JSON):
{{"vendor": "", "invoice_number": "", "total": 0, "due_date": "", "line_items": []}}

{text[:2000]}\n\nJSON:'''
    return llm_call(extract_prompt)

def handle_contract(text):
    # Extract: parties, effective date, term, key obligations
    extract_prompt = f'''Extract from this contract (JSON):
{{"parties": [], "effective_date": "", "term_months": 0, "key_obligations": []}}

{text[:2000]}\n\nJSON:'''
    return llm_call(extract_prompt)

ROUTERS = {
    'invoice':  handle_invoice,
    'contract': handle_contract,
    'report':   lambda t: llm_call(f'Summarize this report in 3 bullet points:\n{t[:2000]}'),
    'email':    lambda t: llm_call(f'Extract: sender, subject, action required from this email:\n{t[:1000]}')
}

def route_document(text):
    classification = classify_document(text)
    doc_type = classification['type']
    handler = ROUTERS.get(doc_type, lambda t: llm_call(f'Describe this document:\n{t[:1000]}'))
    return handler(text)

Sub-Type Classification

CONTRACT_SUBTYPES = {
    'employment': ['employment', 'employee', 'employer', 'salary', 'compensation', 'job title'],
    'nda': ['non-disclosure', 'confidential', 'nda', 'proprietary information'],
    'service': ['service agreement', 'scope of work', 'deliverables', 'milestone'],
    'lease': ['lease', 'landlord', 'tenant', 'rent', 'premises', 'square feet']
}

def classify_contract_subtype(text):
    text_lower = text.lower()
    scores = {
        subtype: sum(1 for kw in keywords if kw in text_lower)
        for subtype, keywords in CONTRACT_SUBTYPES.items()
    }
    best = max(scores, key=scores.get)
    if scores[best] == 0:
        return 'general'
    return best

def handle_contract_routed(text):
    subtype = classify_contract_subtype(text)
    print(f'Contract subtype: {subtype}')
    # Route to specialized extractor
    if subtype == 'nda':
        return extract_nda_fields(text)
    elif subtype == 'employment':
        return extract_employment_fields(text)
    else:
        return handle_contract(text)

Batch Classification Pipeline

from concurrent.futures import ThreadPoolExecutor
import time

def classify_batch(documents):
    results = []
    for doc in documents:
        text = extract_text(doc['path'])  # PDF, OCR, or plain text
        classification = classify_document(text[:1500])
        results.append({
            'doc_id':   doc['id'],
            'path':     doc['path'],
            'type':     classification['type'],
            'method':   classification['method'],
            'text':     text
        })
    return results

def process_batch(documents, max_workers=4):
    # Step 1: classify all (fast)
    classified = classify_batch(documents)

    # Step 2: group by type
    from collections import defaultdict
    by_type = defaultdict(list)
    for doc in classified:
        by_type[doc['type']].append(doc)

    # Step 3: process each group
    all_results = {}
    for doc_type, docs in by_type.items():
        handler = ROUTERS.get(doc_type)
        if handler:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = {executor.submit(handler, d['text']): d for d in docs}
                for fut, doc in futures.items():
                    all_results[doc['doc_id']] = fut.result()
    return all_results

Handling 'Other' and Unknown Types

HUMAN_REVIEW_QUEUE = []

def process_document(doc_path):
    text = extract_text(doc_path)
    classification = classify_with_confidence(text[:1500])

    # High confidence path
    if classification['confidence'] >= 0.8 and classification['type'] != 'other':
        handler = ROUTERS.get(classification['type'])
        return {
            'result': handler(text),
            'type': classification['type'],
            'auto_processed': True
        }

    # Low confidence or unknown type
    HUMAN_REVIEW_QUEUE.append({
        'path': doc_path,
        'predicted_type': classification['type'],
        'confidence': classification['confidence'],
        'reason': classification.get('reason', '')
    })

    print(f'Added to review queue: {doc_path} ({classification["confidence"]:.0%} confident)')
    return {'auto_processed': False, 'queued_for_review': True}

Classification Feedback Loop

correction_log = []

def log_correction(doc_path, predicted_type, correct_type, text_sample):
    correction_log.append({
        'doc_path': doc_path,
        'predicted': predicted_type,
        'correct': correct_type,
        'text_sample': text_sample[:200]
    })
    print(f'Logged correction: {predicted_type} -> {correct_type}')

def build_few_shot_examples(n=5):
    recent = correction_log[-n:]  # use most recent corrections
    examples = []
    for entry in recent:
        examples.append(
            f'Text: {entry["text_sample"]}\nCorrect type: {entry["correct"]}'
        )
    return '\n\n'.join(examples)

def classify_with_few_shot(text):
    few_shot = build_few_shot_examples()
    prompt = f'''Examples of correct classifications:\n{few_shot}\n\nNow classify:\n{text[:800]}\n\nType:'''
    return llm_call(prompt).strip().lower()

Structured Data Extraction After Classification

import json

class FakeMsg:
    def __init__(self, content): self.content = content

class FakeChoice:
    def __init__(self, content): self.message = FakeMsg(content)

class FakeResponse:
    def __init__(self, content): self.choices = [FakeChoice(content)]

class _Completions:
    @staticmethod
    def create(model, messages, temperature):
        return FakeResponse('{"vendor_name": "Acme", "invoice_number": "123", "total": 100}')

class _Chat:
    completions = _Completions()

class FakeClient:
    chat = _Chat()

client = FakeClient()

EXTRACTION_SCHEMAS = {
    'invoice': '{vendor_name: , invoice_number: , invoice_date: , due_date: , subtotal: 0, tax: 0, total: 0, line_items: []}',
}

def extract_structured_fields(text, doc_type):
    schema = EXTRACTION_SCHEMAS.get(doc_type)
    if not schema:
        return {'error': f'No extraction schema for type: {doc_type}'}
    prompt = (f'Extract fields from this {doc_type}. Return JSON matching this schema:\n'
              f'{schema}\n\nDocument:\n{text[:2000]}\n\nJSON:')
    response = client.chat.completions.create(model='gpt-4o-mini', messages=[{'role': 'user', 'content': prompt}], temperature=0)
    try:
        return json.loads(response.choices[0].message.content)
    except json.JSONDecodeError:
        return {'error': 'Could not parse extraction result'}

print(extract_structured_fields('Invoice #123 from Acme for $100', 'invoice'))

Knowledge Check

Recap: Document Classification and Routing