src/document/schemas/classification.py

"""
Document Classification Schemas

Pydantic models for document type classification and categorization.
"""

from enum import Enum
from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field

from .core import EvidenceRef


class DocumentType(str, Enum):
    """
    Common document types for classification.
    Extensible for domain-specific types.
    """
    # Legal & Business
    CONTRACT = "contract"
    INVOICE = "invoice"
    RECEIPT = "receipt"
    PURCHASE_ORDER = "purchase_order"
    AGREEMENT = "agreement"
    NDA = "nda"
    TERMS_OF_SERVICE = "terms_of_service"

    # Technical & Scientific
    PATENT = "patent"
    RESEARCH_PAPER = "research_paper"
    TECHNICAL_REPORT = "technical_report"
    SPECIFICATION = "specification"
    DATASHEET = "datasheet"
    USER_MANUAL = "user_manual"

    # Financial
    FINANCIAL_REPORT = "financial_report"
    BANK_STATEMENT = "bank_statement"
    TAX_FORM = "tax_form"
    BALANCE_SHEET = "balance_sheet"
    INCOME_STATEMENT = "income_statement"

    # Identity & Administrative
    ID_DOCUMENT = "id_document"
    PASSPORT = "passport"
    DRIVERS_LICENSE = "drivers_license"
    CERTIFICATE = "certificate"
    FORM = "form"
    APPLICATION = "application"

    # Medical
    MEDICAL_RECORD = "medical_record"
    PRESCRIPTION = "prescription"
    LAB_REPORT = "lab_report"
    INSURANCE_CLAIM = "insurance_claim"

    # General
    LETTER = "letter"
    EMAIL = "email"
    MEMO = "memo"
    PRESENTATION = "presentation"
    SPREADSHEET = "spreadsheet"
    REPORT = "report"
    ARTICLE = "article"
    BOOK = "book"

    # Catch-all
    OTHER = "other"
    UNKNOWN = "unknown"


class ClassificationScore(BaseModel):
    """Score for a single document type classification."""
    document_type: DocumentType = Field(..., description="Document type")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Classification confidence")
    reasoning: Optional[str] = Field(default=None, description="Reasoning for classification")


class DocumentClassification(BaseModel):
    """
    Document classification result with confidence scores.
    """
    document_id: str = Field(..., description="Document identifier")

    # Primary classification
    primary_type: DocumentType = Field(..., description="Most likely document type")
    primary_confidence: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence in primary classification"
    )

    # All classification scores
    scores: List[ClassificationScore] = Field(
        default_factory=list,
        description="Scores for all considered types"
    )

    # Evidence
    evidence: List[EvidenceRef] = Field(
        default_factory=list,
        description="Evidence supporting classification"
    )

    # Classification metadata
    method: str = Field(
        default="llm",
        description="Classification method used (llm/rule-based/hybrid)"
    )
    model_used: Optional[str] = Field(default=None, description="Model used for classification")

    # Warnings and flags
    is_confident: bool = Field(
        default=True,
        description="Whether classification meets confidence threshold"
    )
    warnings: List[str] = Field(default_factory=list, description="Classification warnings")
    needs_human_review: bool = Field(
        default=False,
        description="Whether human review is recommended"
    )

    # Additional attributes detected
    attributes: Dict[str, Any] = Field(
        default_factory=dict,
        description="Additional detected attributes (language, domain, etc.)"
    )

    def get_top_k(self, k: int = 3) -> List[ClassificationScore]:
        """Get top k classifications by confidence."""
        sorted_scores = sorted(self.scores, key=lambda x: x.confidence, reverse=True)
        return sorted_scores[:k]

    def is_type(self, doc_type: DocumentType, min_confidence: float = 0.5) -> bool:
        """Check if document is classified as a specific type with minimum confidence."""
        for score in self.scores:
            if score.document_type == doc_type and score.confidence >= min_confidence:
                return True
        return False


class DocumentCategoryRule(BaseModel):
    """
    Rule for rule-based document classification.
    """
    name: str = Field(..., description="Rule name")
    document_type: DocumentType = Field(..., description="Target document type")

    # Matching criteria
    title_keywords: List[str] = Field(
        default_factory=list,
        description="Keywords to match in title"
    )
    content_keywords: List[str] = Field(
        default_factory=list,
        description="Keywords to match in content"
    )
    required_sections: List[str] = Field(
        default_factory=list,
        description="Required section headings"
    )
    file_patterns: List[str] = Field(
        default_factory=list,
        description="Filename patterns (regex)"
    )

    # Confidence adjustment
    base_confidence: float = Field(
        default=0.8,
        ge=0.0,
        le=1.0,
        description="Base confidence when rule matches"
    )
    keyword_boost: float = Field(
        default=0.05,
        ge=0.0,
        le=0.2,
        description="Confidence boost per matched keyword"
    )

    # Priority
    priority: int = Field(
        default=0,
        description="Rule priority (higher = checked first)"
    )


class ClassificationConfig(BaseModel):
    """
    Configuration for document classification.
    """
    # Confidence thresholds
    min_confidence: float = Field(
        default=0.6,
        ge=0.0,
        le=1.0,
        description="Minimum confidence for classification"
    )
    human_review_threshold: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Below this, flag for human review"
    )

    # Classification method
    use_llm: bool = Field(default=True, description="Use LLM for classification")
    use_rules: bool = Field(default=True, description="Use rule-based classification")
    hybrid_mode: str = Field(
        default="llm_primary",
        description="Hybrid mode: llm_primary, rules_primary, or ensemble"
    )

    # Custom rules
    custom_rules: List[DocumentCategoryRule] = Field(
        default_factory=list,
        description="Custom classification rules"
    )

    # Document types to consider
    enabled_types: List[DocumentType] = Field(
        default_factory=lambda: list(DocumentType),
        description="Document types to consider"
    )

    # Evidence requirements
    require_evidence: bool = Field(
        default=True,
        description="Require evidence for classification"
    )
    max_evidence_snippets: int = Field(
        default=3,
        description="Maximum evidence snippets to include"
    )