Sort incoming documents into folders with an LLM

import json
import os
import shutil
import time
from pathlib import Path

from openai import OpenAI

INBOX = Path(os.environ.get("INBOX_DIR", "./inbox"))
OUTBOX = Path(os.environ.get("OUTBOX_DIR", "./sorted"))
CATEGORIES = ["invoice", "contract", "support", "other"]
SUFFIXES = {".txt", ".md", ".eml"}
MIN_CONFIDENCE = 0.6
client = OpenAI()  # reads OPENAI_API_KEY

SYSTEM = (
    f"Classify the document into exactly one of: {', '.join(CATEGORIES)}. "
    'Reply with JSON: {"category": str, "confidence": number 0-1, "summary": str}.'
)


def classify(text: str) -> dict:
    resp = client.chat.completions.create(
        model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
        response_format={"type": "json_object"},
        temperature=0,
        messages=[
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": text[:12_000]},  # bound the cost per file
        ],
    )
    data = json.loads(resp.choices[0].message.content or "{}")
    try:
        confident = float(data.get("confidence", 0)) >= MIN_CONFIDENCE
    except (TypeError, ValueError):
        confident = False
    if data.get("category") not in CATEGORIES or not confident:
        data["category"] = "review"  # a human decides
    return data


def move(path: Path, folder: str) -> None:
    dest = OUTBOX / folder
    dest.mkdir(parents=True, exist_ok=True)
    shutil.move(path, dest / path.name)


def process(path: Path) -> None:
    result = classify(path.read_text(encoding="utf-8", errors="replace"))
    move(path, result["category"])
    with (OUTBOX / "log.jsonl").open("a", encoding="utf-8") as log:
        log.write(json.dumps({"file": path.name, **result}) + "\n")
    print(f"{path.name} -> {result['category']}")


def main(poll_seconds: int = 10) -> None:
    INBOX.mkdir(parents=True, exist_ok=True)
    while True:
        for path in sorted(INBOX.iterdir()):
            # Skip folders, other types and files modified in the last few seconds
            if path.suffix not in SUFFIXES or time.time() - path.stat().st_mtime < 5:
                continue
            try:
                process(path)
            except Exception as exc:
                print(f"{path.name} failed: {exc}")
                move(path, "failed")
        time.sleep(poll_seconds)


if __name__ == "__main__":
    main()