backend/db_models/evaluation.py

"""
Evaluation model for AegisLM SaaS Backend.

Production-ready SQLAlchemy model for storing
AI evaluation results and job tracking.
"""

from datetime import datetime
from enum import Enum
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Text, JSON, ForeignKey, Float
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func

from core.database import Base


class EvaluationStatus(str, Enum):
    """Evaluation job status."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class Evaluation(Base):
    """Evaluation model for tracking AI assessment jobs."""
    
    __tablename__ = "evaluations"
    
    # Primary key
    id = Column(Integer, primary_key=True, index=True)
    
    # Foreign key
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
    
    # Job tracking
    job_id = Column(String(255), unique=True, index=True, nullable=False)
    status = Column(String(50), default=EvaluationStatus.PENDING, nullable=False, index=True)
    
    # Model configuration
    model_config = Column(JSON, nullable=False)
    pipeline_config = Column(JSON, nullable=False)
    
    # Results storage
    result_json = Column(JSON, nullable=True)
    error_message = Column(Text, nullable=True)
    
    # Experiment tracking
    experiment_run_id = Column(String(255), nullable=True, index=True)  # Link to experiment table
    
    # Audit tracking
    config_hash = Column(String(64), nullable=True, index=True)  # SHA-256 hash of configuration
    result_checksum = Column(String(64), nullable=True, index=True)  # SHA-256 checksum of result
    audit_trail_id = Column(String(255), nullable=True, index=True)  # Link to audit trail
    
    # Metrics
    total_attacks = Column(Integer, default=0, nullable=False)
    successful_attacks = Column(Integer, default=0, nullable=False)
    success_rate = Column(Float, nullable=True)  # Store as float (0.0 to 1.0) for efficiency
    execution_time_ms = Column(Integer, nullable=True)
    queue_delay_ms = Column(Integer, nullable=True)  # Time spent in queue
    total_latency_ms = Column(Integer, nullable=True)  # Total time including queue
    
    # Audit fields
    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False, index=True)
    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False)
    started_at = Column(DateTime(timezone=True), nullable=True)
    completed_at = Column(DateTime(timezone=True), nullable=True)
    
    # Relationships
    user = relationship("User", back_populates="evaluations")
    
    def __repr__(self):
        return f"<Evaluation(id={self.id}, job_id={self.job_id}, status={self.status})>"
    
    def to_dict(self):
        """Convert evaluation to dictionary."""
        return {
            "id": self.id,
            "user_id": self.user_id,
            "job_id": self.job_id,
            "status": self.status,
            "model_config": self.model_config,
            "pipeline_config": self.pipeline_config,
            "result_json": self.result_json,
            "error_message": self.error_message,
            "experiment_run_id": self.experiment_run_id,
            "config_hash": self.config_hash,
            "result_checksum": self.result_checksum,
            "audit_trail_id": self.audit_trail_id,
            "total_attacks": self.total_attacks,
            "successful_attacks": self.successful_attacks,
            "success_rate": self.success_rate,
            "execution_time_ms": self.execution_time_ms,
            "queue_delay_ms": self.queue_delay_ms,
            "total_latency_ms": self.total_latency_ms,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
        }
    
    def update_status(self, status: EvaluationStatus, error_message: str = None):
        """Update evaluation status and timestamps."""
        self.status = status.value
        self.updated_at = datetime.utcnow()
        
        if status == EvaluationStatus.RUNNING and not self.started_at:
            self.started_at = datetime.utcnow()
        elif status in [EvaluationStatus.COMPLETED, EvaluationStatus.FAILED, EvaluationStatus.CANCELLED]:
            self.completed_at = datetime.utcnow()
            if error_message:
                self.error_message = error_message