backend/app/services/classification.py

"""Document classification — zero-shot via LLM.

Returns a type + confidence + rationale, plus ranked candidates. Uses a
truncated view of the document so it stays cheap and fast.
"""
from __future__ import annotations

from app.core.logging import get_logger
from app.llm.base import LLMMessage
from app.llm.registry import get_provider
from app.schemas.documents import Classification

log = get_logger(__name__)

KNOWN_TYPES = [
    "invoice", "receipt", "contract", "purchase_order", "form",
    "resume", "employee_record", "report", "letter", "email",
    "bank_statement", "id_document", "spreadsheet", "other",
]

_SYS = (
    "You are a document classification expert. Classify the document into the single "
    f"best-fitting type from this list: {', '.join(KNOWN_TYPES)}. "
    "Pick 'employee_record' for HR records / employee rosters (even when several "
    "people are listed), 'spreadsheet' for tabular data dumps, and 'form' for "
    "fillable forms. Use 'other' only when nothing fits. "
    "Respond ONLY as JSON with keys: doc_type (string from the list), "
    "confidence (0..1), rationale (short string), "
    "candidates (array of {type, confidence} for the top 3)."
)


async def classify(markdown: str, provider: str | None = None) -> Classification:
    llm = get_provider(provider)
    excerpt = markdown[:6000] if markdown else "(empty document)"
    msgs = [
        LLMMessage(role="system", content=_SYS),
        LLMMessage(role="user", content=f"Document content:\n\n{excerpt}"),
    ]
    try:
        data = await llm.complete_json(msgs)
    except Exception as e:
        log.warning("classification failed: %s", e)
        return Classification(doc_type="other", confidence=0.0, rationale="classification error")

    return Classification(
        doc_type=str(data.get("doc_type", "other")).lower(),
        confidence=float(data.get("confidence", 0.5) or 0.5),
        rationale=str(data.get("rationale", "")),
        candidates=data.get("candidates", []) or [],
    )