diff --git a/app/api/__init__.py b/app/api/__init__.py index ae98cbd7..95d7d268 100644 --- a/app/api/__init__.py +++ b/app/api/__init__.py @@ -11,6 +11,7 @@ from app.api.api_tokens import router as api_tokens_router from app.api.azure import router as azure_router from app.api.backup import router as backup_router from app.api.billing import router as billing_router +from app.api.classification_rules import router as classification_rules_router from app.api.database import router as database_router from app.api.diagnostic import router as diagnostic_router from app.api.dropbox import router as dropbox_router @@ -82,3 +83,4 @@ router.include_router(imap_accounts_router) router.include_router(integrations_router) router.include_router(notifications_router) router.include_router(scheduled_jobs_router) +router.include_router(classification_rules_router) diff --git a/app/api/classification_rules.py b/app/api/classification_rules.py new file mode 100644 index 00000000..594ade24 --- /dev/null +++ b/app/api/classification_rules.py @@ -0,0 +1,325 @@ +"""Classification Rules API endpoints. + +Provides CRUD operations for managing custom document classification rules. +System-wide rules (``owner_id IS NULL``) can only be managed by admins. +""" + +from __future__ import annotations + +import logging +from typing import Annotated, Any + +from fastapi import APIRouter, Depends, HTTPException, Request, status +from pydantic import BaseModel, Field +from sqlalchemy.orm import Session + +from app.auth import require_login +from app.database import get_db +from app.models import ClassificationRuleModel +from app.utils.classification_rules import ( + BUILTIN_CATEGORIES, + RULE_TYPE_CONTENT, + RULE_TYPE_FILENAME, + RULE_TYPE_METADATA, +) + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/classification-rules", tags=["classification"]) + +DbSession = Annotated[Session, Depends(get_db)] + +_VALID_RULE_TYPES = {RULE_TYPE_FILENAME, RULE_TYPE_CONTENT, RULE_TYPE_METADATA} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _get_user_id(request: Request) -> str: + """Extract the user identifier from the request session.""" + user = getattr(request.state, "user", None) + if user and hasattr(user, "get"): + return user.get("sub") or user.get("email") or "anonymous" + return "anonymous" + + +def _is_admin(request: Request) -> bool: + """Check whether the current user is an admin.""" + user = getattr(request.state, "user", None) + if user and hasattr(user, "get"): + groups = user.get("groups", []) + return "admin" in groups or "Admin" in groups + return False + + +# --------------------------------------------------------------------------- +# Pydantic schemas +# --------------------------------------------------------------------------- + + +class RuleCreate(BaseModel): + """Schema for creating a classification rule.""" + + name: str = Field(..., min_length=1, max_length=255) + category: str = Field(..., min_length=1, max_length=100) + rule_type: str = Field(..., description="One of: filename_pattern, content_keyword, metadata_match") + pattern: str = Field(..., min_length=1, max_length=1000) + priority: int = Field(default=0, ge=0, le=1000) + case_sensitive: bool = False + enabled: bool = True + + +class RuleUpdate(BaseModel): + """Schema for updating a classification rule.""" + + name: str | None = Field(default=None, min_length=1, max_length=255) + category: str | None = Field(default=None, min_length=1, max_length=100) + rule_type: str | None = Field(default=None) + pattern: str | None = Field(default=None, min_length=1, max_length=1000) + priority: int | None = Field(default=None, ge=0, le=1000) + case_sensitive: bool | None = None + enabled: bool | None = None + + +class RuleResponse(BaseModel): + """Schema for a classification rule response.""" + + id: int + owner_id: str | None + name: str + category: str + rule_type: str + pattern: str + priority: int + case_sensitive: bool + enabled: bool + + model_config = {"from_attributes": True} + + +# --------------------------------------------------------------------------- +# Endpoints +# --------------------------------------------------------------------------- + + +@router.get("/categories") +@require_login +async def list_categories(request: Request) -> dict[str, str]: + """Return all built-in classification categories. + + Custom categories created via rules are not included here; they are + discovered dynamically when rules are evaluated. + """ + return BUILTIN_CATEGORIES + + +@router.get("/rule-types") +@require_login +async def list_rule_types(request: Request) -> list[dict[str, str]]: + """Return the supported rule types with descriptions.""" + return [ + { + "type": RULE_TYPE_FILENAME, + "label": "Filename Pattern", + "description": "Regex pattern matched against the original filename.", + }, + { + "type": RULE_TYPE_CONTENT, + "label": "Content Keyword", + "description": "Pipe-separated keywords matched against the OCR text.", + }, + { + "type": RULE_TYPE_METADATA, + "label": "Metadata Match", + "description": "field=value pattern matched against existing AI metadata.", + }, + ] + + +@router.get("/") +@require_login +async def list_rules(request: Request, db: DbSession) -> list[dict[str, Any]]: + """List classification rules visible to the current user. + + Returns both system rules (``owner_id IS NULL``) and the user's own rules. + """ + user_id = _get_user_id(request) + rules = ( + db.query(ClassificationRuleModel) + .filter((ClassificationRuleModel.owner_id.is_(None)) | (ClassificationRuleModel.owner_id == user_id)) + .order_by(ClassificationRuleModel.priority.desc(), ClassificationRuleModel.id) + .all() + ) + return [ + { + "id": r.id, + "owner_id": r.owner_id, + "name": r.name, + "category": r.category, + "rule_type": r.rule_type, + "pattern": r.pattern, + "priority": r.priority, + "case_sensitive": r.case_sensitive, + "enabled": r.enabled, + } + for r in rules + ] + + +@router.post("/", status_code=status.HTTP_201_CREATED) +@require_login +async def create_rule(request: Request, body: RuleCreate, db: DbSession) -> dict[str, Any]: + """Create a new custom classification rule. + + The rule is owned by the current user. Admins may create system-wide + rules by setting ``owner_id`` to ``null`` (not yet exposed). + """ + if body.rule_type not in _VALID_RULE_TYPES: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid rule_type. Must be one of: {', '.join(sorted(_VALID_RULE_TYPES))}", + ) + + user_id = _get_user_id(request) + + # Check for duplicate name within the user's scope + existing = ( + db.query(ClassificationRuleModel) + .filter(ClassificationRuleModel.owner_id == user_id, ClassificationRuleModel.name == body.name) + .first() + ) + if existing: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"A rule named '{body.name}' already exists.", + ) + + rule = ClassificationRuleModel( + owner_id=user_id, + name=body.name, + category=body.category, + rule_type=body.rule_type, + pattern=body.pattern, + priority=body.priority, + case_sensitive=body.case_sensitive, + enabled=body.enabled, + ) + try: + db.add(rule) + db.commit() + db.refresh(rule) + except Exception: + db.rollback() + raise + + logger.info("Classification rule created: id=%s, user=%s", rule.id, user_id) + return { + "id": rule.id, + "owner_id": rule.owner_id, + "name": rule.name, + "category": rule.category, + "rule_type": rule.rule_type, + "pattern": rule.pattern, + "priority": rule.priority, + "case_sensitive": rule.case_sensitive, + "enabled": rule.enabled, + } + + +@router.get("/{rule_id}") +@require_login +async def get_rule(request: Request, rule_id: int, db: DbSession) -> dict[str, Any]: + """Get a single classification rule by ID.""" + user_id = _get_user_id(request) + rule = db.query(ClassificationRuleModel).filter(ClassificationRuleModel.id == rule_id).first() + if rule is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Rule not found") + + # Users can see system rules and their own rules + if rule.owner_id is not None and rule.owner_id != user_id and not _is_admin(request): + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Rule not found") + + return { + "id": rule.id, + "owner_id": rule.owner_id, + "name": rule.name, + "category": rule.category, + "rule_type": rule.rule_type, + "pattern": rule.pattern, + "priority": rule.priority, + "case_sensitive": rule.case_sensitive, + "enabled": rule.enabled, + } + + +@router.put("/{rule_id}") +@require_login +async def update_rule(request: Request, rule_id: int, body: RuleUpdate, db: DbSession) -> dict[str, Any]: + """Update an existing classification rule. + + Users can only update their own rules. Admins can update any rule. + """ + user_id = _get_user_id(request) + rule = db.query(ClassificationRuleModel).filter(ClassificationRuleModel.id == rule_id).first() + if rule is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Rule not found") + + if rule.owner_id != user_id and not _is_admin(request): + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Cannot modify this rule") + + if body.rule_type is not None and body.rule_type not in _VALID_RULE_TYPES: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid rule_type. Must be one of: {', '.join(sorted(_VALID_RULE_TYPES))}", + ) + + update_data = body.model_dump(exclude_unset=True) + for field_name, value in update_data.items(): + setattr(rule, field_name, value) + + try: + db.commit() + db.refresh(rule) + except Exception: + db.rollback() + raise + + logger.info("Classification rule updated: id=%s, user=%s", rule.id, user_id) + return { + "id": rule.id, + "owner_id": rule.owner_id, + "name": rule.name, + "category": rule.category, + "rule_type": rule.rule_type, + "pattern": rule.pattern, + "priority": rule.priority, + "case_sensitive": rule.case_sensitive, + "enabled": rule.enabled, + } + + +@router.delete("/{rule_id}", status_code=status.HTTP_204_NO_CONTENT) +@require_login +async def delete_rule(request: Request, rule_id: int, db: DbSession) -> None: + """Delete a classification rule. + + Users can only delete their own rules. Admins can delete any rule. + """ + user_id = _get_user_id(request) + rule = db.query(ClassificationRuleModel).filter(ClassificationRuleModel.id == rule_id).first() + if rule is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Rule not found") + + if rule.owner_id != user_id and not _is_admin(request): + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Cannot delete this rule") + + try: + db.delete(rule) + db.commit() + except Exception: + db.rollback() + raise + + logger.info("Classification rule deleted: id=%s, user=%s", rule_id, user_id) diff --git a/app/api/pipelines.py b/app/api/pipelines.py index 8175832f..6acf273f 100644 --- a/app/api/pipelines.py +++ b/app/api/pipelines.py @@ -117,8 +117,14 @@ PIPELINE_STEP_TYPES: dict[str, dict[str, Any]] = { }, "classify": { "label": "Document Classification", - "description": "Classify the document type using AI without full metadata extraction.", - "config_schema": {}, + "description": "Classify the document type using built-in and custom rules (filename patterns, content keywords, metadata matching).", + "config_schema": { + "use_builtin_rules": { + "type": "boolean", + "default": True, + "description": "Include the pre-built classification rules (invoice, contract, receipt, etc.).", + }, + }, }, } diff --git a/app/celery_worker.py b/app/celery_worker.py index 4881f8a9..b1cf054f 100644 --- a/app/celery_worker.py +++ b/app/celery_worker.py @@ -22,6 +22,7 @@ from app.tasks.batch_tasks import ( # noqa: F401 sync_search_index, ) from app.tasks.check_credentials import check_credentials +from app.tasks.classify_document import classify_document_task # noqa: F401 from app.tasks.compute_embedding import backfill_missing_embeddings, compute_document_embedding # noqa: F401 from app.tasks.convert_to_pdf import convert_to_pdf # noqa: F401 from app.tasks.convert_to_pdfa import convert_to_pdfa # noqa: F401 diff --git a/app/models.py b/app/models.py index 0cc7b53a..e48a49cf 100644 --- a/app/models.py +++ b/app/models.py @@ -833,3 +833,48 @@ class ScheduledJob(Base): created_at = Column(DateTime(timezone=True), server_default=func.now()) updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now()) + + +class ClassificationRuleModel(Base): + """Custom document classification rule. + + Rules are evaluated during the ``classify`` pipeline step to assign a + category to a document. System-wide rules have ``owner_id IS NULL``; + user-specific rules belong to a single owner. + """ + + __tablename__ = "classification_rules" + + id = Column(Integer, primary_key=True, index=True) + + # NULL = system-wide rule visible to all users. + owner_id = Column(String, nullable=True, index=True) + + # Human-readable rule name (unique per owner). + name = Column(String(255), nullable=False) + + # Target category (e.g. "invoice", "contract", "receipt"). + category = Column(String(100), nullable=False, index=True) + + # Rule type: "filename_pattern", "content_keyword", or "metadata_match". + rule_type = Column(String(50), nullable=False) + + # The matching pattern: + # - filename_pattern: a regex + # - content_keyword: pipe-separated keywords + # - metadata_match: "field=value" + pattern = Column(String(1000), nullable=False) + + # Higher priority rules are evaluated first (default 0). + priority = Column(Integer, nullable=False, default=0) + + # Whether pattern matching is case-sensitive. + case_sensitive = Column(Boolean, nullable=False, default=False) + + # Disabled rules are skipped during classification. + enabled = Column(Boolean, nullable=False, default=True) + + created_at = Column(DateTime(timezone=True), server_default=func.now()) + updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now()) + + __table_args__ = (UniqueConstraint("owner_id", "name", name="uq_classification_rules_owner_name"),) diff --git a/app/tasks/classify_document.py b/app/tasks/classify_document.py new file mode 100644 index 00000000..846270a1 --- /dev/null +++ b/app/tasks/classify_document.py @@ -0,0 +1,173 @@ +"""Celery task for rule-based document classification. + +This task is executed as a pipeline step (``step_type="classify"``). It +applies built-in and user-defined classification rules against the document's +filename, OCR text, and existing AI metadata to assign a ``document_type`` +category. + +The result is stored in the ``ai_metadata`` JSON blob on the +:class:`~app.models.FileRecord` (field ``classification``). +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +from app.celery_app import celery +from app.database import SessionLocal +from app.models import ClassificationRuleModel, FileRecord +from app.tasks.retry_config import BaseTaskWithRetry +from app.utils import log_task_progress +from app.utils.classification_rules import ( + ClassificationResult, + classify_document, + db_rule_to_engine_rule, +) + +logger = logging.getLogger(__name__) + +STEP_NAME = "classify_document" + + +def _load_custom_rules(owner_id: str | None) -> list[Any]: + """Load enabled custom classification rules from the database. + + Returns engine-level :class:`ClassificationRule` dataclass instances. + Rules are loaded in priority-descending order. System rules + (``owner_id IS NULL``) and the user's own rules are both included. + """ + with SessionLocal() as db: + query = db.query(ClassificationRuleModel).filter(ClassificationRuleModel.enabled.is_(True)) + if owner_id: + query = query.filter( + (ClassificationRuleModel.owner_id.is_(None)) | (ClassificationRuleModel.owner_id == owner_id) + ) + else: + query = query.filter(ClassificationRuleModel.owner_id.is_(None)) + rules = query.order_by(ClassificationRuleModel.priority.desc()).all() + return [db_rule_to_engine_rule(r) for r in rules] + + +@celery.task(base=BaseTaskWithRetry, bind=True) +def classify_document_task( + self: Any, + file_id: int, + owner_id: str | None = None, +) -> dict[str, Any]: + """Classify a document using rule-based matching. + + This task: + 1. Loads the :class:`FileRecord` from the database. + 2. Gathers filename, OCR text, and existing AI metadata. + 3. Loads built-in + user-defined classification rules. + 4. Runs the classification engine. + 5. Persists the result into ``ai_metadata.classification``. + + Args: + file_id: Primary key of the :class:`FileRecord` to classify. + owner_id: Owner identifier for loading user-specific rules. + + Returns: + Dict with ``category``, ``confidence``, and ``matched_rules``. + """ + task_id = self.request.id + + log_task_progress( + task_id, + STEP_NAME, + "in_progress", + f"Starting classification for file {file_id}", + file_id=file_id, + ) + + try: + with SessionLocal() as db: + file_record: FileRecord | None = db.query(FileRecord).filter(FileRecord.id == file_id).first() + if file_record is None: + log_task_progress( + task_id, + STEP_NAME, + "failure", + f"FileRecord {file_id} not found", + file_id=file_id, + ) + return {"status": "error", "detail": "File not found"} + + # Gather inputs + filename = file_record.original_filename or "" + text = file_record.ocr_text or "" + existing_metadata: dict[str, Any] = {} + if file_record.ai_metadata: + try: + existing_metadata = json.loads(file_record.ai_metadata) + except (json.JSONDecodeError, TypeError): + existing_metadata = {} + + # Load custom rules + effective_owner = owner_id or file_record.owner_id + custom_rules = _load_custom_rules(effective_owner) + + # Run classification engine + result: ClassificationResult = classify_document( + filename=filename, + text=text, + metadata=existing_metadata, + custom_rules=custom_rules, + ) + + # Persist result into ai_metadata + classification_data = { + "category": result.category, + "confidence": result.confidence, + "matched_rules": [ + { + "rule_name": m.rule_name, + "rule_type": m.rule_type, + "category": m.category, + "confidence": m.confidence, + } + for m in result.matched_rules + ], + } + + existing_metadata["classification"] = classification_data + + # If no document_type was set yet, populate it from the classification + if not existing_metadata.get("document_type"): + from app.utils.classification_rules import BUILTIN_CATEGORIES + + existing_metadata["document_type"] = BUILTIN_CATEGORIES.get( + result.category, result.category.replace("_", " ").title() + ) + + file_record.ai_metadata = json.dumps(existing_metadata, ensure_ascii=False) + db.commit() + + log_task_progress( + task_id, + STEP_NAME, + "success", + f"Classified as '{result.category}' with confidence {result.confidence}", + file_id=file_id, + detail=f"Matched {len(result.matched_rules)} rule(s)", + ) + + return { + "status": "success", + "category": result.category, + "confidence": result.confidence, + "matched_rules": len(result.matched_rules), + } + + except Exception as e: + logger.exception("Classification failed for file %s: %s", file_id, e) + log_task_progress( + task_id, + STEP_NAME, + "failure", + f"Classification failed: {e}", + file_id=file_id, + ) + raise diff --git a/app/utils/classification_rules.py b/app/utils/classification_rules.py new file mode 100644 index 00000000..4c8f61c2 --- /dev/null +++ b/app/utils/classification_rules.py @@ -0,0 +1,384 @@ +""" +Rule-based document classification engine. + +Provides pre-built categories and a rule matcher that classifies documents +using filename patterns, content keywords, and metadata fields. Custom +rules stored in the database are evaluated alongside the built-in defaults. + +Usage:: + + from app.utils.classification_rules import classify_document + + result = classify_document( + filename="2024-03-01_Invoice_Acme.pdf", + text="Invoice total: $1,234.56", + metadata={"absender": "Acme Corp"}, + custom_rules=custom_rules_from_db, + ) + # result -> ClassificationResult(category="invoice", confidence=85, matched_rules=[...]) +""" + +from __future__ import annotations + +import logging +import re +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Pre-built categories +# --------------------------------------------------------------------------- + +#: Canonical category names recognised by the system. Users may also define +#: their own categories via custom rules. +BUILTIN_CATEGORIES: dict[str, str] = { + "invoice": "Invoice", + "contract": "Contract", + "receipt": "Receipt", + "letter": "Letter", + "report": "Report", + "bank_statement": "Bank Statement", + "tax_document": "Tax Document", + "insurance": "Insurance Document", + "payslip": "Payslip", + "unknown": "Unknown", +} + +# --------------------------------------------------------------------------- +# Rule type constants +# --------------------------------------------------------------------------- + +RULE_TYPE_FILENAME = "filename_pattern" +RULE_TYPE_CONTENT = "content_keyword" +RULE_TYPE_METADATA = "metadata_match" + + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + + +@dataclass +class ClassificationRule: + """A single classification rule.""" + + name: str + category: str + rule_type: str # filename_pattern | content_keyword | metadata_match + pattern: str # regex for filename, keyword(s) for content, "field=value" for metadata + priority: int = 0 # higher = evaluated first + case_sensitive: bool = False + + def __post_init__(self) -> None: + if self.rule_type not in (RULE_TYPE_FILENAME, RULE_TYPE_CONTENT, RULE_TYPE_METADATA): + raise ValueError(f"Invalid rule_type: {self.rule_type!r}") + + +@dataclass +class MatchedRule: + """Records which rule matched and why.""" + + rule_name: str + rule_type: str + category: str + confidence: int + + +@dataclass +class ClassificationResult: + """The outcome of running the classification engine on a document.""" + + category: str + confidence: int # 0 – 100 + matched_rules: list[MatchedRule] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Built-in rules +# --------------------------------------------------------------------------- + +BUILTIN_RULES: list[ClassificationRule] = [ + # ── Invoice ─────────────────────────────────────────────────────────── + ClassificationRule("builtin_invoice_filename", "invoice", RULE_TYPE_FILENAME, r"(?i)invoice|rechnung|facture"), + ClassificationRule( + "builtin_invoice_content", + "invoice", + RULE_TYPE_CONTENT, + "invoice number|invoice total|amount due|rechnung|rechnungsnummer|total amount|bill to", + ), + ClassificationRule("builtin_invoice_metadata", "invoice", RULE_TYPE_METADATA, "document_type=Invoice"), + ClassificationRule( + "builtin_invoice_kommunikationsart", "invoice", RULE_TYPE_METADATA, "kommunikationsart=Rechnung" + ), + # ── Contract ────────────────────────────────────────────────────────── + ClassificationRule("builtin_contract_filename", "contract", RULE_TYPE_FILENAME, r"(?i)contract|vertrag|agreement"), + ClassificationRule( + "builtin_contract_content", + "contract", + RULE_TYPE_CONTENT, + "hereby agrees|terms and conditions|vertrag|agreement between|party agrees|effective date", + ), + ClassificationRule("builtin_contract_metadata", "contract", RULE_TYPE_METADATA, "document_type=Contract"), + ClassificationRule( + "builtin_contract_kommunikationsart", "contract", RULE_TYPE_METADATA, "kommunikationsart=Vertrag" + ), + # ── Receipt ─────────────────────────────────────────────────────────── + ClassificationRule("builtin_receipt_filename", "receipt", RULE_TYPE_FILENAME, r"(?i)receipt|quittung|beleg"), + ClassificationRule( + "builtin_receipt_content", + "receipt", + RULE_TYPE_CONTENT, + "receipt|quittung|payment received|thank you for your purchase|transaction id", + ), + ClassificationRule("builtin_receipt_metadata", "receipt", RULE_TYPE_METADATA, "document_type=Receipt"), + ClassificationRule( + "builtin_receipt_kommunikationsart", "receipt", RULE_TYPE_METADATA, "kommunikationsart=Quittung" + ), + # ── Letter ──────────────────────────────────────────────────────────── + ClassificationRule("builtin_letter_filename", "letter", RULE_TYPE_FILENAME, r"(?i)letter|brief|schreiben"), + ClassificationRule( + "builtin_letter_content", + "letter", + RULE_TYPE_CONTENT, + "dear sir|dear madam|sehr geehrte|to whom it may concern|sincerely|mit freundlichen", + ), + # ── Report ──────────────────────────────────────────────────────────── + ClassificationRule("builtin_report_filename", "report", RULE_TYPE_FILENAME, r"(?i)report|bericht"), + ClassificationRule( + "builtin_report_content", + "report", + RULE_TYPE_CONTENT, + "executive summary|table of contents|annual report|quarterly report|findings", + ), + # ── Bank statement ──────────────────────────────────────────────────── + ClassificationRule( + "builtin_bank_filename", + "bank_statement", + RULE_TYPE_FILENAME, + r"(?i)bank.?statement|kontoauszug", + ), + ClassificationRule( + "builtin_bank_content", + "bank_statement", + RULE_TYPE_CONTENT, + "account statement|kontoauszug|opening balance|closing balance|account number", + ), + ClassificationRule( + "builtin_bank_kommunikationsart", "bank_statement", RULE_TYPE_METADATA, "kommunikationsart=Kontoauszug" + ), + # ── Tax document ────────────────────────────────────────────────────── + ClassificationRule("builtin_tax_filename", "tax_document", RULE_TYPE_FILENAME, r"(?i)tax|steuer|steuerbescheid"), + ClassificationRule( + "builtin_tax_content", + "tax_document", + RULE_TYPE_CONTENT, + "tax return|steuerbescheid|taxable income|finanzamt|tax assessment", + ), + # ── Insurance ───────────────────────────────────────────────────────── + ClassificationRule( + "builtin_insurance_filename", "insurance", RULE_TYPE_FILENAME, r"(?i)insurance|versicherung|police" + ), + ClassificationRule( + "builtin_insurance_content", + "insurance", + RULE_TYPE_CONTENT, + "insurance policy|versicherung|policennummer|coverage|premium|deductible", + ), + # ── Payslip ─────────────────────────────────────────────────────────── + ClassificationRule( + "builtin_payslip_filename", "payslip", RULE_TYPE_FILENAME, r"(?i)payslip|gehaltsabrechnung|lohnabrechnung" + ), + ClassificationRule( + "builtin_payslip_content", + "payslip", + RULE_TYPE_CONTENT, + "gross salary|net salary|gehaltsabrechnung|lohnabrechnung|bruttolohn|nettolohn", + ), +] + + +# --------------------------------------------------------------------------- +# Confidence scoring +# --------------------------------------------------------------------------- + +#: Base confidence for each rule type when it matches. +_CONFIDENCE_MAP: dict[str, int] = { + RULE_TYPE_FILENAME: 60, + RULE_TYPE_CONTENT: 70, + RULE_TYPE_METADATA: 90, +} + +#: Extra confidence per additional matching rule of the same category (capped). +_CONFIDENCE_BONUS_PER_EXTRA_RULE = 10 + + +# --------------------------------------------------------------------------- +# Matching helpers +# --------------------------------------------------------------------------- + + +def _match_filename(rule: ClassificationRule, filename: str) -> bool: + """Return True if *rule.pattern* (regex) matches anywhere in *filename*.""" + if not filename: + return False + flags = 0 if rule.case_sensitive else re.IGNORECASE + return bool(re.search(rule.pattern, filename, flags)) + + +def _match_content(rule: ClassificationRule, text: str) -> bool: + """Return True if any keyword in *rule.pattern* appears in *text*. + + Keywords are separated by ``|`` (pipe). + """ + if not text: + return False + keywords = [kw.strip() for kw in rule.pattern.split("|") if kw.strip()] + text_lower = text if rule.case_sensitive else text.lower() + return any((kw if rule.case_sensitive else kw.lower()) in text_lower for kw in keywords) + + +def _match_metadata(rule: ClassificationRule, metadata: dict[str, Any] | None) -> bool: + """Return True if *rule.pattern* (``field=value``) matches *metadata*. + + Pattern format: ``field_name=expected_value``. + """ + if not metadata: + return False + if "=" not in rule.pattern: + return False + field_name, expected_value = rule.pattern.split("=", 1) + actual = metadata.get(field_name.strip()) + if actual is None: + return False + if rule.case_sensitive: + return str(actual) == expected_value.strip() + return str(actual).lower() == expected_value.strip().lower() + + +_MATCHERS = { + RULE_TYPE_FILENAME: _match_filename, + RULE_TYPE_CONTENT: _match_content, + RULE_TYPE_METADATA: _match_metadata, +} + + +def _evaluate_rule( + rule: ClassificationRule, + filename: str, + text: str, + metadata: dict[str, Any] | None, +) -> MatchedRule | None: + """Evaluate a single rule against the document. Return a :class:`MatchedRule` on match.""" + matcher = _MATCHERS.get(rule.rule_type) + if matcher is None: + return None + + # Dispatch to the appropriate matcher based on rule type + if rule.rule_type == RULE_TYPE_FILENAME: + matched = matcher(rule, filename) + elif rule.rule_type == RULE_TYPE_CONTENT: + matched = matcher(rule, text) + elif rule.rule_type == RULE_TYPE_METADATA: + matched = matcher(rule, metadata) + else: + matched = False + + if matched: + return MatchedRule( + rule_name=rule.name, + rule_type=rule.rule_type, + category=rule.category, + confidence=_CONFIDENCE_MAP.get(rule.rule_type, 50), + ) + return None + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def classify_document( + filename: str = "", + text: str = "", + metadata: dict[str, Any] | None = None, + custom_rules: list[ClassificationRule] | None = None, +) -> ClassificationResult: + """Classify a document by evaluating built-in and custom rules. + + Rules are evaluated in priority order (highest first, then built-in before + custom for the same priority). The category with the most rule matches + wins; ties are broken by cumulative confidence. + + Args: + filename: Original filename of the document. + text: Extracted / OCR text of the document. + metadata: Previously-extracted AI metadata dict (e.g. from ``ai_metadata``). + custom_rules: Optional list of user-defined :class:`ClassificationRule` objects. + + Returns: + A :class:`ClassificationResult` with the best matching category, + overall confidence score, and the list of rules that fired. + """ + all_rules = list(BUILTIN_RULES) + if custom_rules: + all_rules.extend(custom_rules) + + # Sort by priority descending (higher priority first) + all_rules.sort(key=lambda r: r.priority, reverse=True) + + matches: list[MatchedRule] = [] + for rule in all_rules: + result = _evaluate_rule(rule, filename, text, metadata) + if result is not None: + matches.append(result) + + if not matches: + return ClassificationResult(category="unknown", confidence=0, matched_rules=[]) + + # Aggregate by category: pick the one with the most matches, then highest + # cumulative confidence as tiebreaker. + category_scores: dict[str, list[MatchedRule]] = {} + for m in matches: + category_scores.setdefault(m.category, []).append(m) + + best_category = max( + category_scores, + key=lambda cat: (len(category_scores[cat]), sum(m.confidence for m in category_scores[cat])), + ) + + best_matches = category_scores[best_category] + base_confidence = max(m.confidence for m in best_matches) + bonus = min( + (len(best_matches) - 1) * _CONFIDENCE_BONUS_PER_EXTRA_RULE, + 100 - base_confidence, + ) + final_confidence = min(base_confidence + bonus, 100) + + return ClassificationResult( + category=best_category, + confidence=final_confidence, + matched_rules=best_matches, + ) + + +def db_rule_to_engine_rule(db_rule: Any) -> ClassificationRule: + """Convert a database ``ClassificationRuleModel`` row to an engine :class:`ClassificationRule`. + + Args: + db_rule: A SQLAlchemy model instance with ``name``, ``category``, + ``rule_type``, ``pattern``, ``priority``, and ``case_sensitive`` attributes. + + Returns: + A :class:`ClassificationRule` dataclass instance. + """ + return ClassificationRule( + name=db_rule.name, + category=db_rule.category, + rule_type=db_rule.rule_type, + pattern=db_rule.pattern, + priority=db_rule.priority, + case_sensitive=getattr(db_rule, "case_sensitive", False), + ) diff --git a/app/views/files.py b/app/views/files.py index 58efa132..5aeae09e 100644 --- a/app/views/files.py +++ b/app/views/files.py @@ -396,9 +396,7 @@ _STEP_TYPE_TO_STAGES: dict[str, list[str]] = { "embed_metadata": ["embed_metadata_into_pdf"], "compute_embedding": ["compute_embedding"], "send_to_destinations": ["finalize_document_storage", "send_to_all_destinations"], - # "classify" is defined in PIPELINE_STEP_TYPES but has no Celery log stages yet. - # When a classify task is implemented, add its stage key(s) here. - "classify": [], + "classify": ["classify_document"], } # These internal bookkeeping stages are always shown in the flow regardless of diff --git a/migrations/versions/027_add_classification_rules.py b/migrations/versions/027_add_classification_rules.py new file mode 100644 index 00000000..9864c0b2 --- /dev/null +++ b/migrations/versions/027_add_classification_rules.py @@ -0,0 +1,46 @@ +"""Add classification_rules table for custom document classification rules. + +Revision ID: 027_add_classification_rules +Revises: 026_add_scheduled_jobs +Create Date: 2026-03-09 +""" + +from typing import Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "027_add_classification_rules" +down_revision: Union[str, None] = "026_add_scheduled_jobs" +depends_on: Union[str, None] = None + + +def upgrade() -> None: + """Create classification_rules table.""" + op.create_table( + "classification_rules", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("owner_id", sa.String(), nullable=True), + sa.Column("name", sa.String(255), nullable=False), + sa.Column("category", sa.String(100), nullable=False), + sa.Column("rule_type", sa.String(50), nullable=False), + sa.Column("pattern", sa.String(1000), nullable=False), + sa.Column("priority", sa.Integer(), nullable=False, server_default="0"), + sa.Column("case_sensitive", sa.Boolean(), nullable=False, server_default="0"), + sa.Column("enabled", sa.Boolean(), nullable=False, server_default="1"), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("owner_id", "name", name="uq_classification_rules_owner_name"), + ) + op.create_index("ix_classification_rules_id", "classification_rules", ["id"]) + op.create_index("ix_classification_rules_owner_id", "classification_rules", ["owner_id"]) + op.create_index("ix_classification_rules_category", "classification_rules", ["category"]) + + +def downgrade() -> None: + """Drop classification_rules table.""" + op.drop_index("ix_classification_rules_category", "classification_rules") + op.drop_index("ix_classification_rules_owner_id", "classification_rules") + op.drop_index("ix_classification_rules_id", "classification_rules") + op.drop_table("classification_rules") diff --git a/tests/conftest.py b/tests/conftest.py index ae6db91e..9951ee6c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -61,6 +61,7 @@ from app.main import app as fastapi_app # noqa: E402 # Import models to register them with SQLAlchemy Base from app.models import ( # noqa: F401, E402 ApiToken, + ClassificationRuleModel, DocumentMetadata, FileRecord, Pipeline, diff --git a/tests/test_api_classification_rules.py b/tests/test_api_classification_rules.py new file mode 100644 index 00000000..a519c6a4 --- /dev/null +++ b/tests/test_api_classification_rules.py @@ -0,0 +1,290 @@ +"""Tests for the classification rules API endpoints. + +Covers CRUD operations, validation, and access control for +``/api/classification-rules``. +""" + +import pytest + +from app.models import ClassificationRuleModel + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_rule(db_session, owner_id="anonymous", **overrides): + """Insert a ClassificationRuleModel and return it.""" + defaults = { + "owner_id": owner_id, + "name": "test_rule", + "category": "invoice", + "rule_type": "filename_pattern", + "pattern": r"(?i)invoice", + "priority": 0, + "case_sensitive": False, + "enabled": True, + } + defaults.update(overrides) + rule = ClassificationRuleModel(**defaults) + db_session.add(rule) + db_session.commit() + db_session.refresh(rule) + return rule + + +# --------------------------------------------------------------------------- +# Categories & Rule Types endpoints +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestCategoriesEndpoint: + """Tests for GET /api/classification-rules/categories.""" + + def test_list_categories(self, client): + """Should return a dict of built-in categories.""" + r = client.get("/api/classification-rules/categories") + assert r.status_code == 200 + data = r.json() + assert isinstance(data, dict) + assert "invoice" in data + assert "contract" in data + assert "receipt" in data + assert "unknown" in data + + +@pytest.mark.unit +class TestRuleTypesEndpoint: + """Tests for GET /api/classification-rules/rule-types.""" + + def test_list_rule_types(self, client): + """Should return a list of valid rule types.""" + r = client.get("/api/classification-rules/rule-types") + assert r.status_code == 200 + data = r.json() + assert isinstance(data, list) + assert len(data) == 3 + type_values = {item["type"] for item in data} + assert "filename_pattern" in type_values + assert "content_keyword" in type_values + assert "metadata_match" in type_values + + +# --------------------------------------------------------------------------- +# CRUD Operations +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestClassificationRuleCRUD: + """Full CRUD test-suite for classification rules.""" + + def test_list_rules_empty(self, client): + """List returns an empty array when no rules exist.""" + r = client.get("/api/classification-rules/") + assert r.status_code == 200 + assert r.json() == [] + + def test_create_rule(self, client): + """POST should create a new classification rule.""" + r = client.post( + "/api/classification-rules/", + json={ + "name": "My Invoice Rule", + "category": "invoice", + "rule_type": "filename_pattern", + "pattern": r"(?i)rechnung", + "priority": 10, + }, + ) + assert r.status_code == 201 + data = r.json() + assert data["name"] == "My Invoice Rule" + assert data["category"] == "invoice" + assert data["rule_type"] == "filename_pattern" + assert data["priority"] == 10 + assert data["enabled"] is True + assert data["id"] is not None + + def test_create_rule_invalid_type_rejected(self, client): + """Creating a rule with an invalid rule_type should be rejected.""" + r = client.post( + "/api/classification-rules/", + json={ + "name": "Bad Rule", + "category": "test", + "rule_type": "invalid_type", + "pattern": "test", + }, + ) + assert r.status_code == 400 + + def test_create_duplicate_name_rejected(self, client): + """Creating two rules with the same name should be rejected.""" + payload = { + "name": "Dupe Rule", + "category": "invoice", + "rule_type": "filename_pattern", + "pattern": "test", + } + r1 = client.post("/api/classification-rules/", json=payload) + assert r1.status_code == 201 + r2 = client.post("/api/classification-rules/", json=payload) + assert r2.status_code == 409 + + def test_get_rule(self, client): + """GET should return a specific rule by ID.""" + create_resp = client.post( + "/api/classification-rules/", + json={ + "name": "Get Test Rule", + "category": "contract", + "rule_type": "content_keyword", + "pattern": "agreement|terms", + }, + ) + rule_id = create_resp.json()["id"] + + r = client.get(f"/api/classification-rules/{rule_id}") + assert r.status_code == 200 + assert r.json()["name"] == "Get Test Rule" + assert r.json()["category"] == "contract" + + def test_get_nonexistent_rule(self, client): + """GET for a nonexistent rule should return 404.""" + r = client.get("/api/classification-rules/99999") + assert r.status_code == 404 + + def test_update_rule(self, client): + """PUT should update an existing rule.""" + create_resp = client.post( + "/api/classification-rules/", + json={ + "name": "Update Test", + "category": "receipt", + "rule_type": "filename_pattern", + "pattern": "receipt", + }, + ) + rule_id = create_resp.json()["id"] + + r = client.put( + f"/api/classification-rules/{rule_id}", + json={"category": "invoice", "priority": 50}, + ) + assert r.status_code == 200 + assert r.json()["category"] == "invoice" + assert r.json()["priority"] == 50 + # Name should be unchanged + assert r.json()["name"] == "Update Test" + + def test_update_nonexistent_rule(self, client): + """PUT for a nonexistent rule should return 404.""" + r = client.put("/api/classification-rules/99999", json={"category": "test"}) + assert r.status_code == 404 + + def test_update_invalid_rule_type_rejected(self, client): + """PUT with an invalid rule_type should be rejected.""" + create_resp = client.post( + "/api/classification-rules/", + json={ + "name": "Invalid Update", + "category": "test", + "rule_type": "filename_pattern", + "pattern": "test", + }, + ) + rule_id = create_resp.json()["id"] + + r = client.put( + f"/api/classification-rules/{rule_id}", + json={"rule_type": "bad_type"}, + ) + assert r.status_code == 400 + + def test_delete_rule(self, client): + """DELETE should remove the rule.""" + create_resp = client.post( + "/api/classification-rules/", + json={ + "name": "Delete Test", + "category": "test", + "rule_type": "content_keyword", + "pattern": "test", + }, + ) + rule_id = create_resp.json()["id"] + + r = client.delete(f"/api/classification-rules/{rule_id}") + assert r.status_code == 204 + + # Verify it's gone + r2 = client.get(f"/api/classification-rules/{rule_id}") + assert r2.status_code == 404 + + def test_delete_nonexistent_rule(self, client): + """DELETE for a nonexistent rule should return 404.""" + r = client.delete("/api/classification-rules/99999") + assert r.status_code == 404 + + def test_list_rules_after_create(self, client): + """List should return created rules.""" + client.post( + "/api/classification-rules/", + json={ + "name": "List Rule 1", + "category": "invoice", + "rule_type": "filename_pattern", + "pattern": "test1", + }, + ) + client.post( + "/api/classification-rules/", + json={ + "name": "List Rule 2", + "category": "contract", + "rule_type": "content_keyword", + "pattern": "test2", + }, + ) + r = client.get("/api/classification-rules/") + assert r.status_code == 200 + assert len(r.json()) == 2 + + def test_create_rule_with_all_fields(self, client): + """Create a rule providing all optional fields.""" + r = client.post( + "/api/classification-rules/", + json={ + "name": "Full Rule", + "category": "tax_document", + "rule_type": "metadata_match", + "pattern": "department=finance", + "priority": 100, + "case_sensitive": True, + "enabled": False, + }, + ) + assert r.status_code == 201 + data = r.json() + assert data["case_sensitive"] is True + assert data["enabled"] is False + assert data["priority"] == 100 + + def test_create_rule_defaults(self, client): + """Create a rule with minimal fields to test defaults.""" + r = client.post( + "/api/classification-rules/", + json={ + "name": "Minimal Rule", + "category": "invoice", + "rule_type": "filename_pattern", + "pattern": "test", + }, + ) + assert r.status_code == 201 + data = r.json() + assert data["priority"] == 0 + assert data["case_sensitive"] is False + assert data["enabled"] is True diff --git a/tests/test_classification_rules.py b/tests/test_classification_rules.py new file mode 100644 index 00000000..1005022e --- /dev/null +++ b/tests/test_classification_rules.py @@ -0,0 +1,422 @@ +"""Tests for the rule-based document classification engine. + +Covers the classification engine logic in ``app/utils/classification_rules.py``: +built-in rules, custom rules, confidence scoring, and edge cases. +""" + +import pytest + +from app.utils.classification_rules import ( + BUILTIN_CATEGORIES, + BUILTIN_RULES, + RULE_TYPE_CONTENT, + RULE_TYPE_FILENAME, + RULE_TYPE_METADATA, + ClassificationResult, + ClassificationRule, + MatchedRule, + classify_document, + db_rule_to_engine_rule, +) + +# --------------------------------------------------------------------------- +# Built-in categories & rules smoke tests +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestBuiltinCategories: + """Verify the pre-built categories and rules are sane.""" + + def test_builtin_categories_not_empty(self): + """There must be at least one built-in category.""" + assert len(BUILTIN_CATEGORIES) > 0 + + def test_unknown_category_exists(self): + """The 'unknown' fallback category must be present.""" + assert "unknown" in BUILTIN_CATEGORIES + + def test_core_categories_present(self): + """Invoice, contract, and receipt categories must exist.""" + for cat in ("invoice", "contract", "receipt"): + assert cat in BUILTIN_CATEGORIES, f"Missing built-in category: {cat}" + + def test_builtin_rules_not_empty(self): + """There must be at least one built-in rule.""" + assert len(BUILTIN_RULES) > 0 + + def test_all_builtin_rules_reference_valid_types(self): + """Every built-in rule must use a valid rule_type.""" + valid_types = {RULE_TYPE_FILENAME, RULE_TYPE_CONTENT, RULE_TYPE_METADATA} + for rule in BUILTIN_RULES: + assert rule.rule_type in valid_types, f"Rule {rule.name!r} has invalid type {rule.rule_type!r}" + + +# --------------------------------------------------------------------------- +# ClassificationRule dataclass validation +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestClassificationRuleValidation: + """Test ClassificationRule dataclass validation.""" + + def test_valid_rule_types(self): + """Valid rule types should not raise.""" + for rt in (RULE_TYPE_FILENAME, RULE_TYPE_CONTENT, RULE_TYPE_METADATA): + rule = ClassificationRule(name="test", category="test", rule_type=rt, pattern="test") + assert rule.rule_type == rt + + def test_invalid_rule_type_raises(self): + """An invalid rule_type should raise ValueError.""" + with pytest.raises(ValueError, match="Invalid rule_type"): + ClassificationRule(name="test", category="test", rule_type="invalid", pattern="test") + + +# --------------------------------------------------------------------------- +# Filename pattern matching +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestFilenamePatternMatching: + """Test classification via filename patterns.""" + + def test_invoice_filename(self): + """A filename containing 'invoice' should classify as invoice.""" + result = classify_document(filename="2024-03-01_Invoice_Acme.pdf") + assert result.category == "invoice" + assert result.confidence > 0 + + def test_german_invoice_filename(self): + """A filename containing 'Rechnung' should classify as invoice.""" + result = classify_document(filename="Rechnung_2024.pdf") + assert result.category == "invoice" + assert result.confidence > 0 + + def test_contract_filename(self): + """A filename containing 'contract' should classify as contract.""" + result = classify_document(filename="Service_Contract_2024.pdf") + assert result.category == "contract" + + def test_receipt_filename(self): + """A filename containing 'receipt' should classify as receipt.""" + result = classify_document(filename="Payment_Receipt.pdf") + assert result.category == "receipt" + + def test_unrecognised_filename(self): + """A generic filename with no keywords should return 'unknown'.""" + result = classify_document(filename="document_12345.pdf") + assert result.category == "unknown" + assert result.confidence == 0 + + def test_empty_filename(self): + """An empty filename should not match any rule.""" + result = classify_document(filename="") + assert result.category == "unknown" + + +# --------------------------------------------------------------------------- +# Content keyword matching +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestContentKeywordMatching: + """Test classification via content keywords.""" + + def test_invoice_content(self): + """Text containing 'invoice number' should classify as invoice.""" + result = classify_document(text="Please pay the invoice number 12345. Amount due: $500") + assert result.category == "invoice" + assert result.confidence > 0 + + def test_contract_content(self): + """Text containing 'terms and conditions' should classify as contract.""" + result = classify_document(text="The parties hereby agree to the following terms and conditions.") + assert result.category == "contract" + + def test_receipt_content(self): + """Text containing 'payment received' should classify as receipt.""" + result = classify_document(text="Thank you. Payment received for order #789.") + assert result.category == "receipt" + + def test_bank_statement_content(self): + """Text containing 'account statement' should classify as bank_statement.""" + result = classify_document(text="Monthly account statement. Opening balance: $1,000.") + assert result.category == "bank_statement" + + def test_empty_text(self): + """Empty text should not match any content rule.""" + result = classify_document(text="") + assert result.category == "unknown" + + def test_case_insensitive_matching(self): + """Content matching should be case-insensitive by default.""" + result = classify_document(text="INVOICE NUMBER 12345") + assert result.category == "invoice" + + +# --------------------------------------------------------------------------- +# Metadata matching +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestMetadataMatching: + """Test classification via metadata field matching.""" + + def test_document_type_invoice(self): + """metadata document_type=Invoice should classify as invoice.""" + result = classify_document(metadata={"document_type": "Invoice"}) + assert result.category == "invoice" + assert result.confidence >= 90 + + def test_document_type_contract(self): + """metadata document_type=Contract should classify as contract.""" + result = classify_document(metadata={"document_type": "Contract"}) + assert result.category == "contract" + + def test_kommunikationsart_rechnung(self): + """German classification metadata should classify as invoice.""" + result = classify_document(metadata={"kommunikationsart": "Rechnung"}) + assert result.category == "invoice" + + def test_no_metadata(self): + """None metadata should not match.""" + result = classify_document(metadata=None) + assert result.category == "unknown" + + def test_empty_metadata(self): + """Empty metadata dict should not match.""" + result = classify_document(metadata={}) + assert result.category == "unknown" + + def test_metadata_case_insensitive(self): + """Metadata matching should be case-insensitive by default.""" + result = classify_document(metadata={"document_type": "invoice"}) + assert result.category == "invoice" + + +# --------------------------------------------------------------------------- +# Combined matching / confidence boosting +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestCombinedMatching: + """Test that multiple matching rules boost confidence.""" + + def test_filename_and_content_boost(self): + """Filename + content matching should produce higher confidence than either alone.""" + filename_only = classify_document(filename="Invoice_2024.pdf") + combined = classify_document(filename="Invoice_2024.pdf", text="Invoice number: 12345. Amount due: $500.") + assert combined.confidence >= filename_only.confidence + assert len(combined.matched_rules) > len(filename_only.matched_rules) + + def test_all_three_signals(self): + """Filename + content + metadata should produce highest confidence.""" + result = classify_document( + filename="Invoice_Acme.pdf", + text="Invoice number: 12345. Amount due: $500.", + metadata={"document_type": "Invoice"}, + ) + assert result.category == "invoice" + assert result.confidence >= 90 + + def test_conflicting_signals_most_matches_wins(self): + """When filename says 'invoice' but content says 'contract', most matches wins.""" + result = classify_document( + filename="Invoice.pdf", + text="The parties hereby agree to the following terms and conditions. " + "This agreement between Company A and Company B is effective immediately.", + ) + # Content has more keyword matches for contract, but filename matches invoice. + # Either is acceptable as long as the result is deterministic. + assert result.category in ("invoice", "contract") + + +# --------------------------------------------------------------------------- +# Custom rules +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestCustomRules: + """Test user-defined custom classification rules.""" + + def test_custom_rule_matches(self): + """A custom filename rule should match when pattern hits.""" + custom = [ + ClassificationRule( + name="custom_hr_doc", + category="hr_document", + rule_type=RULE_TYPE_FILENAME, + pattern=r"(?i)employee|hiring|hr", + ) + ] + result = classify_document(filename="Employee_Handbook.pdf", custom_rules=custom) + assert result.category == "hr_document" + + def test_custom_content_rule(self): + """A custom content keyword rule should match.""" + custom = [ + ClassificationRule( + name="custom_medical", + category="medical", + rule_type=RULE_TYPE_CONTENT, + pattern="diagnosis|prescription|patient record", + ) + ] + result = classify_document(text="Patient record for Jane Doe. Diagnosis: common cold.", custom_rules=custom) + assert result.category == "medical" + + def test_custom_metadata_rule(self): + """A custom metadata rule should match.""" + custom = [ + ClassificationRule( + name="custom_legal", + category="legal", + rule_type=RULE_TYPE_METADATA, + pattern="department=legal", + ) + ] + result = classify_document(metadata={"department": "legal"}, custom_rules=custom) + assert result.category == "legal" + + def test_custom_rule_overrides_builtin(self): + """Custom rules with more matches should override built-in rules.""" + custom = [ + ClassificationRule( + name="custom_internal_invoice", + category="internal_invoice", + rule_type=RULE_TYPE_FILENAME, + pattern=r"(?i)invoice", + priority=100, + ), + ClassificationRule( + name="custom_internal_invoice_content", + category="internal_invoice", + rule_type=RULE_TYPE_CONTENT, + pattern="invoice number", + priority=100, + ), + ] + result = classify_document( + filename="Invoice_2024.pdf", + text="Invoice number: 12345", + custom_rules=custom, + ) + # Both builtin and custom rules for "invoice" patterns match, but custom + # has "internal_invoice" as category. The category with more total matches wins. + assert result.category in ("invoice", "internal_invoice") + assert result.confidence > 0 + + +# --------------------------------------------------------------------------- +# db_rule_to_engine_rule converter +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestDbRuleConversion: + """Test the database model to engine rule converter.""" + + def test_converts_basic_fields(self): + """All basic fields should be mapped correctly.""" + + class FakeDbRule: + name = "test_rule" + category = "invoice" + rule_type = RULE_TYPE_FILENAME + pattern = r"(?i)invoice" + priority = 10 + case_sensitive = True + + engine_rule = db_rule_to_engine_rule(FakeDbRule()) + assert engine_rule.name == "test_rule" + assert engine_rule.category == "invoice" + assert engine_rule.rule_type == RULE_TYPE_FILENAME + assert engine_rule.pattern == r"(?i)invoice" + assert engine_rule.priority == 10 + assert engine_rule.case_sensitive is True + + def test_defaults_case_sensitive_to_false(self): + """When case_sensitive is missing, default to False.""" + + class FakeDbRule: + name = "test" + category = "test" + rule_type = RULE_TYPE_CONTENT + pattern = "test" + priority = 0 + + engine_rule = db_rule_to_engine_rule(FakeDbRule()) + assert engine_rule.case_sensitive is False + + +# --------------------------------------------------------------------------- +# ClassificationResult +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestClassificationResult: + """Test the ClassificationResult dataclass.""" + + def test_default_matched_rules(self): + """matched_rules should default to an empty list.""" + result = ClassificationResult(category="test", confidence=50) + assert result.matched_rules == [] + + def test_with_matched_rules(self): + """matched_rules should be populated when provided.""" + match = MatchedRule(rule_name="test", rule_type=RULE_TYPE_FILENAME, category="invoice", confidence=60) + result = ClassificationResult(category="invoice", confidence=60, matched_rules=[match]) + assert len(result.matched_rules) == 1 + assert result.matched_rules[0].rule_name == "test" + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestEdgeCases: + """Test edge cases in the classification engine.""" + + def test_no_inputs_at_all(self): + """No filename, text, or metadata should return 'unknown'.""" + result = classify_document() + assert result.category == "unknown" + assert result.confidence == 0 + assert result.matched_rules == [] + + def test_metadata_pattern_without_equals(self): + """A metadata pattern without '=' should not match.""" + custom = [ + ClassificationRule( + name="bad_pattern", + category="test", + rule_type=RULE_TYPE_METADATA, + pattern="no_equals_sign", + ) + ] + result = classify_document(metadata={"no_equals_sign": "value"}, custom_rules=custom) + assert result.category == "unknown" + + def test_confidence_capped_at_100(self): + """Confidence should never exceed 100.""" + # Create many rules that all match to test the cap + custom = [ + ClassificationRule( + name=f"flood_{i}", + category="flood", + rule_type=RULE_TYPE_CONTENT, + pattern="test keyword", + ) + for i in range(20) + ] + result = classify_document(text="test keyword is here", custom_rules=custom) + assert result.confidence <= 100 diff --git a/tests/test_classify_document.py b/tests/test_classify_document.py new file mode 100644 index 00000000..60fa2fe5 --- /dev/null +++ b/tests/test_classify_document.py @@ -0,0 +1,232 @@ +"""Tests for the classify_document Celery task. + +Covers the ``classify_document_task`` in ``app/tasks/classify_document.py``. +""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from app.models import ClassificationRuleModel, FileRecord +from app.tasks.classify_document import _load_custom_rules + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_file_record(db_session, **overrides): + """Insert a minimal FileRecord and return it.""" + defaults = { + "owner_id": "test-user", + "filehash": "abc123", + "original_filename": "Invoice_2024.pdf", + "local_filename": "/tmp/test.pdf", + "file_size": 1024, + "mime_type": "application/pdf", + "ocr_text": "Invoice number: 12345. Amount due: $500.", + "ai_metadata": None, + } + defaults.update(overrides) + fr = FileRecord(**defaults) + db_session.add(fr) + db_session.commit() + db_session.refresh(fr) + return fr + + +def _make_rule(db_session, **overrides): + """Insert a ClassificationRuleModel and return it.""" + defaults = { + "owner_id": None, + "name": "test_rule", + "category": "test_category", + "rule_type": "filename_pattern", + "pattern": r"(?i)test", + "priority": 0, + "case_sensitive": False, + "enabled": True, + } + defaults.update(overrides) + rule = ClassificationRuleModel(**defaults) + db_session.add(rule) + db_session.commit() + db_session.refresh(rule) + return rule + + +# --------------------------------------------------------------------------- +# _load_custom_rules +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestLoadCustomRules: + """Test the custom rule loading helper.""" + + @patch("app.tasks.classify_document.SessionLocal") + def test_loads_enabled_rules(self, mock_session_local): + """Should load enabled rules from the database.""" + mock_rule = MagicMock() + mock_rule.name = "rule1" + mock_rule.category = "invoice" + mock_rule.rule_type = "filename_pattern" + mock_rule.pattern = r"(?i)invoice" + mock_rule.priority = 10 + mock_rule.case_sensitive = False + + mock_db = MagicMock() + mock_query = MagicMock() + mock_db.query.return_value = mock_query + mock_query.filter.return_value = mock_query + mock_query.order_by.return_value = mock_query + mock_query.all.return_value = [mock_rule] + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + rules = _load_custom_rules(owner_id="test-user") + assert len(rules) == 1 + assert rules[0].name == "rule1" + assert rules[0].category == "invoice" + + +# --------------------------------------------------------------------------- +# classify_document_task (integration-style with mocked DB and Celery) +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestClassifyDocumentTask: + """Test the Celery classify_document_task.""" + + @patch("app.tasks.classify_document.log_task_progress") + @patch("app.tasks.classify_document._load_custom_rules", return_value=[]) + @patch("app.tasks.classify_document.SessionLocal") + def test_classify_invoice_file(self, mock_session_local, mock_load_rules, mock_log): + """Should classify a file with invoice filename and text as 'invoice'.""" + mock_file = MagicMock(spec=FileRecord) + mock_file.id = 1 + mock_file.original_filename = "Invoice_2024.pdf" + mock_file.ocr_text = "Invoice number: 12345. Amount due: $500." + mock_file.ai_metadata = None + mock_file.owner_id = "test-user" + + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_file + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + from app.tasks.classify_document import classify_document_task + + # Call the underlying function directly via .run(), bypassing Celery + result = classify_document_task.run(1, owner_id="test-user") + + assert result["status"] == "success" + assert result["category"] == "invoice" + assert result["confidence"] > 0 + + # Verify ai_metadata was updated + assert mock_file.ai_metadata is not None + metadata = json.loads(mock_file.ai_metadata) + assert "classification" in metadata + assert metadata["classification"]["category"] == "invoice" + + @patch("app.tasks.classify_document.log_task_progress") + @patch("app.tasks.classify_document.SessionLocal") + def test_classify_file_not_found(self, mock_session_local, mock_log): + """Should return error when file record is not found.""" + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = None + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + from app.tasks.classify_document import classify_document_task + + result = classify_document_task.run(99999) + assert result["status"] == "error" + + @patch("app.tasks.classify_document.log_task_progress") + @patch("app.tasks.classify_document._load_custom_rules", return_value=[]) + @patch("app.tasks.classify_document.SessionLocal") + def test_classify_preserves_existing_metadata(self, mock_session_local, mock_load_rules, mock_log): + """Should preserve existing ai_metadata fields and add classification.""" + existing_meta = json.dumps({"document_type": "Invoice", "tags": ["finance"]}) + + mock_file = MagicMock(spec=FileRecord) + mock_file.id = 2 + mock_file.original_filename = "doc.pdf" + mock_file.ocr_text = "" + mock_file.ai_metadata = existing_meta + mock_file.owner_id = "test-user" + + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_file + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + from app.tasks.classify_document import classify_document_task + + classify_document_task.run(2) + + # Check that existing fields are preserved + metadata = json.loads(mock_file.ai_metadata) + assert metadata["tags"] == ["finance"] + assert metadata["document_type"] == "Invoice" + assert "classification" in metadata + + @patch("app.tasks.classify_document.log_task_progress") + @patch("app.tasks.classify_document._load_custom_rules", return_value=[]) + @patch("app.tasks.classify_document.SessionLocal") + def test_classify_sets_document_type_when_missing(self, mock_session_local, mock_load_rules, mock_log): + """Should set document_type from classification when not already present.""" + mock_file = MagicMock(spec=FileRecord) + mock_file.id = 3 + mock_file.original_filename = "Invoice_2024.pdf" + mock_file.ocr_text = "Invoice number: 12345" + mock_file.ai_metadata = json.dumps({"tags": ["test"]}) + mock_file.owner_id = "test-user" + + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_file + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + from app.tasks.classify_document import classify_document_task + + classify_document_task.run(3) + + metadata = json.loads(mock_file.ai_metadata) + assert metadata["document_type"] == "Invoice" + + @patch("app.tasks.classify_document.log_task_progress") + @patch("app.tasks.classify_document._load_custom_rules", return_value=[]) + @patch("app.tasks.classify_document.SessionLocal") + def test_classify_unknown_document(self, mock_session_local, mock_load_rules, mock_log): + """Should classify as 'unknown' when no rules match.""" + mock_file = MagicMock(spec=FileRecord) + mock_file.id = 4 + mock_file.original_filename = "random_file.pdf" + mock_file.ocr_text = "Lorem ipsum dolor sit amet." + mock_file.ai_metadata = None + mock_file.owner_id = "test-user" + + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_file + mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db) + mock_session_local.return_value.__exit__ = MagicMock(return_value=False) + + from app.tasks.classify_document import classify_document_task + + result = classify_document_task.run(4) + + assert result["category"] == "unknown" + assert result["confidence"] == 0 + + def test_classify_document_task_is_celery_task(self): + """Task should be registered as a Celery task.""" + from app.tasks.classify_document import classify_document_task + + assert hasattr(classify_document_task, "apply_async") + assert hasattr(classify_document_task, "delay") + assert callable(classify_document_task)