"""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 [], )