Spaces:
Sleeping
Sleeping
File size: 2,016 Bytes
f65e025 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """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 [],
)
|