breakpilot-compliance/backend-compliance/compliance/db/ai_system_models.py

"""
AI System & Audit Export models — extracted from compliance/db/models.py.

Covers AI Act system registration/classification and the audit export package
tracker. Re-exported from ``compliance.db.models`` for backwards compatibility.

DO NOT change __tablename__, column names, or relationship strings.
"""

import uuid
import enum
from datetime import datetime, timezone

from sqlalchemy import (
    Column, String, Text, Integer, DateTime, Date,
    Enum, JSON, Index, Float,
)

from classroom_engine.database import Base


# ============================================================================
# ENUMS
# ============================================================================

class AIClassificationEnum(str, enum.Enum):
    """AI Act risk classification."""
    PROHIBITED = "prohibited"
    HIGH_RISK = "high-risk"
    LIMITED_RISK = "limited-risk"
    MINIMAL_RISK = "minimal-risk"
    UNCLASSIFIED = "unclassified"


class AISystemStatusEnum(str, enum.Enum):
    """Status of an AI system in compliance tracking."""
    DRAFT = "draft"
    CLASSIFIED = "classified"
    COMPLIANT = "compliant"
    NON_COMPLIANT = "non-compliant"


class ExportStatusEnum(str, enum.Enum):
    """Status of audit export."""
    PENDING = "pending"
    GENERATING = "generating"
    COMPLETED = "completed"
    FAILED = "failed"


# ============================================================================
# MODELS
# ============================================================================

class AISystemDB(Base):
    """
    AI System registry for AI Act compliance.
    Tracks AI systems, their risk classification, and compliance status.
    """
    __tablename__ = 'compliance_ai_systems'

    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = Column(String(300), nullable=False)
    description = Column(Text)
    purpose = Column(String(500))
    sector = Column(String(100))

    # AI Act classification
    classification = Column(Enum(AIClassificationEnum), default=AIClassificationEnum.UNCLASSIFIED)
    status = Column(Enum(AISystemStatusEnum), default=AISystemStatusEnum.DRAFT)

    # Assessment
    assessment_date = Column(DateTime)
    assessment_result = Column(JSON)                                      # Full assessment result
    obligations = Column(JSON)                                            # List of AI Act obligations
    risk_factors = Column(JSON)                                           # Risk factors from assessment
    recommendations = Column(JSON)                                        # Recommendations from assessment

    # Timestamps
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    __table_args__ = (
        Index('ix_ai_system_classification', 'classification'),
        Index('ix_ai_system_status', 'status'),
    )

    def __repr__(self):
        return f"<AISystem {self.name} ({self.classification.value})>"


class AuditExportDB(Base):
    """
    Tracks audit export packages generated for external auditors.
    """
    __tablename__ = 'compliance_audit_exports'

    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))

    export_type = Column(String(50), nullable=False)                     # "full", "controls_only", "evidence_only"
    export_name = Column(String(200))                                    # User-friendly name

    # Scope
    included_regulations = Column(JSON)                                  # List of regulation codes
    included_domains = Column(JSON)                                      # List of control domains
    date_range_start = Column(Date)
    date_range_end = Column(Date)

    # Generation
    requested_by = Column(String(100), nullable=False)
    requested_at = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc))
    completed_at = Column(DateTime)

    # Output
    file_path = Column(String(500))
    file_hash = Column(String(64))                                       # SHA-256 of ZIP
    file_size_bytes = Column(Integer)

    status = Column(Enum(ExportStatusEnum), default=ExportStatusEnum.PENDING)
    error_message = Column(Text)

    # Statistics
    total_controls = Column(Integer)
    total_evidence = Column(Integer)
    compliance_score = Column(Float)

    # Timestamps
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))

    def __repr__(self):
        return f"<AuditExport {self.export_type} by {self.requested_by}>"


__all__ = [
    "AIClassificationEnum",
    "AISystemStatusEnum",
    "ExportStatusEnum",
    "AISystemDB",
    "AuditExportDB",
]