AI-Powered Document Processing Pipeline

"""Document preprocessing: PDF text extraction and OCR."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import fitz  # PyMuPDF


@dataclass
class PageResult:
    page_number: int
    text: str
    is_scanned: bool
    confidence: float  # OCR confidence, 1.0 for native text


def extract_text(pdf_path: Path) -> list[PageResult]:
    """Extract text from a PDF, using OCR for scanned pages.

    Strategy:
    1. Try native text extraction first (PyMuPDF).
    2. If a page has fewer than 50 characters of native text,
       treat it as scanned and apply OCR.
    """
    doc = fitz.open(str(pdf_path))
    results = []

    for page_num, page in enumerate(doc):
        native_text = page.get_text("text").strip()

        if len(native_text) > 50:
            # Digital page — use native text
            results.append(PageResult(
                page_number=page_num + 1,
                text=native_text,
                is_scanned=False,
                confidence=1.0,
            ))
        else:
            # Scanned page — apply OCR
            ocr_text, confidence = _ocr_page(page)
            results.append(PageResult(
                page_number=page_num + 1,
                text=ocr_text,
                is_scanned=True,
                confidence=confidence,
            ))

    doc.close()
    return results


def _ocr_page(page) -> tuple[str, float]:
    """OCR a single page using Tesseract via PyMuPDF."""
    # PyMuPDF can perform OCR directly with Tesseract
    tp = page.get_textpage_ocr(flags=fitz.TEXT_PRESERVE_WHITESPACE)
    text = page.get_text("text", textpage=tp).strip()
    # Estimate confidence from character count vs page area
    confidence = min(len(text) / 500, 1.0)  # Rough heuristic
    return text, confidence

"""Document type classification using LLM."""

from __future__ import annotations

from enum import Enum

from anthropic import Anthropic

client = Anthropic()


class DocumentType(str, Enum):
    INVOICE = "invoice"
    CONTRACT = "contract"
    REPORT = "report"
    FORM = "form"
    LETTER = "letter"
    UNKNOWN = "unknown"


async def classify_document(text_preview: str) -> tuple[DocumentType, float]:
    """Classify a document by type using the first 500 tokens.

    Returns:
        Tuple of (document_type, confidence_score).
    """
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=128,
        messages=[{
            "role": "user",
            "content": (
                "Classify this document. Return JSON with 'type' "
                "(invoice, contract, report, form, letter, unknown) "
                "and 'confidence' (0.0-1.0).\n\n"
                f"Document preview:\n{text_preview[:2000]}"
            ),
        }],
    )
    import json
    data = json.loads(response.content[0].text)
    return DocumentType(data["type"]), float(data["confidence"])

"""Structured extraction using LLM with Pydantic schemas."""

from __future__ import annotations

from datetime import date
from typing import Optional

from anthropic import Anthropic
from pydantic import BaseModel, Field

client = Anthropic()


class ExtractedField(BaseModel):
    """A single extracted field with confidence."""
    value: str | float | None
    confidence: float = Field(ge=0, le=1)


class InvoiceExtraction(BaseModel):
    """Extraction schema for invoices."""
    vendor_name: ExtractedField
    invoice_number: ExtractedField
    invoice_date: ExtractedField
    due_date: ExtractedField
    total_amount: ExtractedField
    currency: ExtractedField
    line_items: list[dict] = Field(default_factory=list)
    tax_amount: Optional[ExtractedField] = None
    payment_terms: Optional[ExtractedField] = None


EXTRACTION_PROMPT = """Extract structured data from this {doc_type}.
Return JSON matching this exact schema. For each field, provide
the extracted value and a confidence score (0.0-1.0).

If a field is not present in the document, set value to null
and confidence to 0.0.

Schema fields: {schema_fields}

Document text:
{document_text}"""


async def extract_fields(
    document_text: str,
    doc_type: str,
    schema_class: type[BaseModel],
) -> BaseModel:
    """Extract structured fields from a document.

    Args:
        document_text: Full text of the document.
        doc_type: The classified document type.
        schema_class: Pydantic model defining the extraction schema.

    Returns:
        Populated Pydantic model with extracted fields.
    """
    schema_fields = ", ".join(schema_class.model_fields.keys())

    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": EXTRACTION_PROMPT.format(
                doc_type=doc_type,
                schema_fields=schema_fields,
                document_text=document_text[:8000],
            ),
        }],
    )

    import json
    data = json.loads(response.content[0].text)
    return schema_class(**data)