""" Document Models Document storage, classification, and embeddings """ import uuid from datetime import datetime, timezone from enum import Enum from sqlalchemy import JSON, Column, DateTime from sqlalchemy import Enum as SQLEnum from sqlalchemy import Float, ForeignKey, Integer, String, Text from sqlalchemy.dialects.postgresql import ARRAY from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB from sqlalchemy.dialects.postgresql import UUID JSONB = JSON().with_variant(PG_JSONB, "postgresql") from pgvector.sqlalchemy import Vector from sqlalchemy.orm import relationship from app.models.base import Base, TimestampMixin, UUIDMixin class DocumentCategory(str, Enum): """Mining document categories for classification""" SAFETY_PROTOCOL = "safety_protocol" EQUIPMENT_MANUAL = "equipment_manual" REGULATORY = "regulatory" INCIDENT_REPORT = "incident_report" GEOLOGICAL = "geological" ENVIRONMENTAL = "environmental" TRAINING = "training" PERMIT = "permit" MAINTENANCE = "maintenance" OTHER = "other" class DocumentStatus(str, Enum): """Document processing status""" PENDING = "pending" PROCESSING = "processing" ANALYZING = "analyzing" COMPLETED = "completed" FAILED = "failed" def db_status(member: DocumentStatus) -> str: """ The literal PostgreSQL actually stores for a DocumentStatus. SQLAlchemy's Enum type persists member *names* ('COMPLETED'), not values ('completed'), unless values_callable is set. ORM queries are coerced automatically, but raw SQL is not — so hand-written SQL must compare against this, never against `.value`. Anything that changes how the column is persisted must change this function too; routing every raw-SQL comparison through here means retrieval fails loudly at one place instead of silently returning zero rows everywhere. """ return member.name class ComplianceStatus(str, Enum): """Safety compliance status""" COMPLIANT = "compliant" WARNING = "warning" VIOLATION = "violation" PENDING = "pending" NOT_APPLICABLE = "not_applicable" class Document(Base, UUIDMixin, TimestampMixin): """ Document model with AI-enhanced metadata. Stores file info, classification, safety analysis, and extracted entities. """ __tablename__ = "documents" # Owner user_id = Column( String(255), ForeignKey("users.clerk_user_id"), nullable=False, index=True ) # File information title = Column(String(500), nullable=False) file_name = Column(String(500), nullable=False) file_size = Column(Integer, nullable=False) # bytes file_type = Column(String(100), nullable=False) # MIME type file_url = Column(Text, nullable=False) # Processing status status = Column( SQLEnum(DocumentStatus), default=DocumentStatus.PENDING, nullable=False, index=True, ) processing_error = Column(Text, nullable=True) processed_at = Column(DateTime, nullable=True) # Extracted content content = Column(Text, nullable=True) # Full text content page_count = Column(Integer, nullable=True) # deprecated alias — use total_pages total_pages = Column( Integer, nullable=True ) # authoritative page count from extractor word_count = Column(Integer, nullable=True) # AI Classification category = Column( SQLEnum(DocumentCategory), default=DocumentCategory.OTHER, nullable=True, index=True, ) subcategory = Column(String(100), nullable=True) classification_confidence = Column(Float, nullable=True) # 0.0 - 1.0 # AI Summary summary = Column(Text, nullable=True) # AI-generated summary key_points = Column(JSONB, nullable=True) # List of key points # Safety Analysis safety_score = Column(Float, nullable=True) # 0-100 compliance_status = Column( SQLEnum(ComplianceStatus), default=ComplianceStatus.PENDING, nullable=True ) hazards_detected = Column(JSONB, nullable=True) # List of hazards safety_recommendations = Column(JSONB, nullable=True) # Named Entity Recognition entities = Column(JSONB, nullable=True) # Structure: # { # "equipment": ["Caterpillar D11", "Komatsu PC8000"], # "chemicals": ["methane", "coal dust"], # "locations": ["Mine Site A", "Section 4B"], # "personnel": ["John Smith", "Safety Team"], # "dates": ["2024-01-15", "Q1 2024"], # "regulations": ["MSHA 30 CFR 75.400", "OSHA 1910.134"] # } # Extra Metadata extra_metadata = Column("metadata", JSONB, default=dict) tags = Column(JSONB, default=list) # Relationships user = relationship("User", back_populates="documents") embeddings = relationship( "DocumentEmbedding", back_populates="document", cascade="all, delete-orphan" ) def __repr__(self): return f"" def to_dict(self): """Convert to dictionary for API responses""" return { "id": str(self.id), "title": self.title, "file_name": self.file_name, "file_size": self.file_size, "file_type": self.file_type, "file_url": self.file_url, "status": self.status.value if self.status else None, "category": self.category.value if self.category else None, "subcategory": self.subcategory, "classification_confidence": self.classification_confidence, "summary": self.summary, "key_points": self.key_points, "safety_score": self.safety_score, "compliance_status": ( self.compliance_status.value if self.compliance_status else None ), "hazards_detected": self.hazards_detected, "entities": { k: v if isinstance(v, list) else [] for k, v in (self.entities or {}).items() } or None, "page_count": self.page_count, "word_count": self.word_count, "created_at": ( self.created_at.replace(tzinfo=timezone.utc).isoformat() if self.created_at else None ), "processed_at": ( self.processed_at.replace(tzinfo=timezone.utc).isoformat() if self.processed_at else None ), "total_pages": self.total_pages or self.page_count, } class DocumentEmbedding(Base, UUIDMixin): """ Vector embeddings for document chunks. Used for semantic search and RAG. The embedding column uses pgvector's native Vector(768) type with an HNSW index (see migration 001) for sub-5ms approximate nearest-neighbor search instead of brute-force Python cosine similarity. """ __tablename__ = "document_embeddings" # Parent document document_id = Column( UUID(as_uuid=True), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, index=True, ) # Chunk information chunk_index = Column(Integer, nullable=False) chunk_text = Column(Text, nullable=False) # Vector embedding — native pgvector type with HNSW index (see migration 001) # Replaces the old JSONB column for 10-100x faster similarity search. embedding = Column(Vector(768), nullable=False) embedding_model = Column(String(100), default="text-embedding-004") # Context metadata — powers context-aware answers with page citations section_title = Column(String(500), nullable=True) # e.g. "Safety Procedures" page_numbers = Column( JSONB, nullable=True ) # e.g. [12, 13] — pages this chunk spans # Legacy page columns (kept for backward compat, use page_numbers instead) start_page = Column(Integer, nullable=True) end_page = Column(Integer, nullable=True) extra_metadata = Column("metadata", JSONB, default=dict) # Relationships document = relationship("Document", back_populates="embeddings") def __repr__(self): return f""