feat(classify): add rule-based document classification engine, task, and API
Implements the classify pipeline step with: - Classification rules engine (app/utils/classification_rules.py) with pre-built categories (invoice, contract, receipt, letter, report, bank_statement, tax_document, insurance, payslip) and support for filename patterns, content keywords, and metadata matching rules - Celery task (app/tasks/classify_document.py) that runs as a pipeline step - CRUD API (app/api/classification_rules.py) for managing custom rules - ClassificationRuleModel in app/models.py with migration 027 - Updated pipeline step config_schema and stage mapping - Comprehensive tests for engine, API, and task Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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.).",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"),)
|
||||
|
||||
@@ -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
|
||||
@@ -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),
|
||||
)
|
||||
+1
-3
@@ -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
|
||||
|
||||
@@ -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")
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user