""" Database models and initialization for LLM Deployment System. This module defines the database schema using SQLAlchemy and provides utilities for database operations. """ from datetime import datetime, timezone from typing import List, Optional, Dict, Any import json import uuid from sqlalchemy import ( create_engine, Column, Integer, String, Text, DateTime, Boolean, Float, ForeignKey, JSON, Enum, UniqueConstraint ) from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, relationship, Session from sqlalchemy.dialects.postgresql import UUID as PGUUID from utils.config import config from utils.logger import get_logger logger = get_logger(__name__) Base = declarative_base() class TaskStatus(str, Enum): """Enumeration of task statuses.""" PENDING = "pending" SENT = "sent" RECEIVED = "received" EVALUATING = "evaluating" COMPLETED = "completed" FAILED = "failed" class EvaluationStatus(str, Enum): """Enumeration of evaluation statuses.""" PENDING = "pending" RUNNING = "running" PASSED = "passed" FAILED = "failed" ERROR = "error" class Submission(Base): """Model for storing student submissions and API endpoints.""" __tablename__ = "submissions" id = Column(Integer, primary_key=True, index=True) email = Column(String(255), nullable=False, index=True) endpoint = Column(String(500), nullable=False) secret = Column(String(255), nullable=False) github_username = Column(String(255)) github_repo_url = Column(String(500)) 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)) # Relationships tasks = relationship("Task", back_populates="submission", cascade="all, delete-orphan") __table_args__ = ( UniqueConstraint('email', 'endpoint', name='unique_submission'), ) class Task(Base): """Model for storing tasks sent to students.""" __tablename__ = "tasks" id = Column(Integer, primary_key=True, index=True) submission_id = Column(Integer, ForeignKey("submissions.id"), nullable=False) # Task identification task_id = Column(String(255), nullable=False, unique=True, index=True) # e.g., "captcha-solver-a1b2c" round = Column(Integer, nullable=False) # 1 or 2 nonce = Column(String(255), nullable=False, unique=True) # Task content brief = Column(Text, nullable=False) checks = Column(JSON, nullable=False) # List of check strings attachments = Column(JSON, default=list) # List of attachment objects # Task lifecycle status = Column(String(50), default="pending") sent_at = Column(DateTime) received_at = Column(DateTime) # Response data status_code = Column(Integer) # HTTP status code from student's API error_message = Column(Text) 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)) # Relationships submission = relationship("Submission", back_populates="tasks") repos = relationship("Repository", back_populates="task", cascade="all, delete-orphan") class Repository(Base): """Model for storing repository information submitted by students.""" __tablename__ = "repositories" id = Column(Integer, primary_key=True, index=True) task_id = Column(Integer, ForeignKey("tasks.id"), nullable=False) # Repository details repo_url = Column(String(500), nullable=False) commit_sha = Column(String(40), nullable=False) # Git commit SHA pages_url = Column(String(500)) # GitHub Pages URL # Timing submitted_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) # Relationships task = relationship("Task", back_populates="repos") evaluations = relationship("Evaluation", back_populates="repository", cascade="all, delete-orphan") class Evaluation(Base): """Model for storing evaluation results.""" __tablename__ = "evaluations" id = Column(Integer, primary_key=True, index=True) repository_id = Column(Integer, ForeignKey("repositories.id"), nullable=False) # Evaluation details check_name = Column(String(255), nullable=False) # Name of the specific check status = Column(String(50), nullable=False) score = Column(Float) # Score for this check (0.0 to 1.0) reason = Column(Text) # Explanation of the result logs = Column(JSON) # Detailed logs from the evaluation # Metadata evaluated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) duration_seconds = Column(Float) # How long the evaluation took # Relationships repository = relationship("Repository", back_populates="evaluations") class TaskTemplate(Base): """Model for storing task templates used to generate tasks.""" __tablename__ = "task_templates" id = Column(Integer, primary_key=True, index=True) template_id = Column(String(100), nullable=False, unique=True, index=True) # e.g., "captcha-solver" # Template content name = Column(String(255), nullable=False) description = Column(Text) # Round 1 configuration brief_template = Column(Text, nullable=False) checks_template = Column(JSON, nullable=False) # List of check templates attachments_template = Column(JSON, default=list) # Round 2 configuration (optional) round2_brief_template = Column(Text) round2_checks_template = Column(JSON) round2_attachments_template = Column(JSON) # Template metadata is_active = Column(Boolean, default=True) 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)) # Seed configuration for randomization seed_config = Column(JSON, default=dict) # Configuration for seed generation class SystemConfig(Base): """Model for storing system-wide configuration.""" __tablename__ = "system_config" id = Column(Integer, primary_key=True) key = Column(String(255), nullable=False, unique=True, index=True) value = Column(JSON, nullable=False) description = Column(Text) updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) class DatabaseManager: """Database manager class for handling connections and operations.""" def __init__(self): self.engine = None self.SessionLocal = None self._setup_database() def _setup_database(self): """Setup database connection and create tables.""" try: # Create engine self.engine = create_engine( config.get_database_url(), connect_args={"check_same_thread": False} if "sqlite" in config.DATABASE_URL else {}, pool_pre_ping=True, echo=config.FLASK_DEBUG ) # Create tables Base.metadata.create_all(bind=self.engine) # Create session factory self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) logger.info(f"Database initialized: {config.DATABASE_URL}") except Exception as e: logger.error(f"Failed to initialize database: {e}") raise def get_session(self) -> Session: """Get a database session.""" return self.SessionLocal() def create_submission(self, email: str, endpoint: str, secret: str, **kwargs) -> Submission: """Create a new submission record.""" with self.get_session() as session: submission = Submission( email=email, endpoint=endpoint, secret=secret, **kwargs ) session.add(submission) session.commit() session.refresh(submission) return submission def get_submission_by_email(self, email: str) -> Optional[Submission]: """Get submission by email.""" with self.get_session() as session: return session.query(Submission).filter(Submission.email == email).first() def create_task(self, submission_id: int, task_data: Dict[str, Any]) -> Task: """Create a new task record.""" with self.get_session() as session: task = Task(submission_id=submission_id, **task_data) session.add(task) session.commit() session.refresh(task) return task def get_task_by_id(self, task_id: str) -> Optional[Task]: """Get task by task_id.""" with self.get_session() as session: return session.query(Task).filter(Task.task_id == task_id).first() def get_tasks_by_submission(self, submission_id: int, round: Optional[int] = None) -> List[Task]: """Get tasks for a submission, optionally filtered by round.""" with self.get_session() as session: query = session.query(Task).filter(Task.submission_id == submission_id) if round is not None: query = query.filter(Task.round == round) return query.all() def update_task_status(self, task_id: str, status: str, **kwargs) -> bool: """Update task status and other fields.""" with self.get_session() as session: task = session.query(Task).filter(Task.task_id == task_id).first() if not task: return False task.status = status if 'status_code' in kwargs: task.status_code = kwargs['status_code'] if 'error_message' in kwargs: task.error_message = kwargs['error_message'] if status == "sent" and not task.sent_at: task.sent_at = datetime.now(timezone.utc) if status == "received" and not task.received_at: task.received_at = datetime.now(timezone.utc) session.commit() return True def create_repository(self, task_id: int, repo_data: Dict[str, Any]) -> Repository: """Create a new repository record.""" with self.get_session() as session: repo = Repository(task_id=task_id, **repo_data) session.add(repo) session.commit() session.refresh(repo) return repo def get_repositories_by_task(self, task_id: str) -> List[Repository]: """Get repositories for a task.""" with self.get_session() as session: task = session.query(Task).filter(Task.task_id == task_id).first() if not task: return [] return task.repos def add_evaluation(self, repository_id: int, evaluation_data: Dict[str, Any]) -> Evaluation: """Add an evaluation result.""" with self.get_session() as session: evaluation = Evaluation(repository_id=repository_id, **evaluation_data) session.add(evaluation) session.commit() session.refresh(evaluation) return evaluation def get_evaluations_by_repository(self, repository_id: int) -> List[Evaluation]: """Get evaluations for a repository.""" with self.get_session() as session: return session.query(Evaluation).filter(Evaluation.repository_id == repository_id).all() def get_task_templates(self) -> List[TaskTemplate]: """Get all active task templates.""" with self.get_session() as session: return session.query(TaskTemplate).filter(TaskTemplate.is_active == True).all() def create_task_template(self, template_data: Dict[str, Any]) -> TaskTemplate: """Create a new task template.""" with self.get_session() as session: template = TaskTemplate(**template_data) session.add(template) session.commit() session.refresh(template) return template # Global database manager instance db_manager = DatabaseManager() def get_db() -> DatabaseManager: """Get the global database manager instance.""" return db_manager def init_database(): """Initialize the database with default data.""" logger.info("Initializing database with default data...") # Create default task templates templates_data = [ { 'template_id': 'sum-of-sales', 'name': 'Sales Summary Application', 'description': 'Create an app that processes CSV data and displays sales summaries', 'brief_template': 'Publish a single-page site that fetches data.csv from attachments, sums its sales column, sets the title to "Sales Summary {seed}", displays the total inside #total-sales, and loads Bootstrap 5 from jsdelivr.', 'checks_template': [ 'js: document.title === `Sales Summary {seed}`', 'js: !!document.querySelector("link[href*=\'bootstrap\']")', 'js: Math.abs(parseFloat(document.querySelector("#total-sales").textContent) - {result}) < 0.01' ], 'attachments_template': [ { 'name': 'data.csv', 'url': 'data:text/csv;base64,{seed}' } ], 'round2_brief_template': 'Add a Bootstrap table #product-sales that lists each product with its total sales and keeps #total-sales accurate after render.', 'round2_checks_template': [ 'js: document.querySelectorAll("#product-sales tbody tr").length >= 1', 'js: (() => { const rows = [...document.querySelectorAll("#product-sales tbody tr td:last-child")]; const sum = rows.reduce((acc, cell) => acc + parseFloat(cell.textContent), 0); return Math.abs(sum - {result}) < 0.01; })()' ] }, { 'template_id': 'markdown-to-html', 'name': 'Markdown to HTML Converter', 'description': 'Create an app that converts Markdown to HTML with syntax highlighting', 'brief_template': 'Publish a static page that converts input.md from attachments to HTML with marked, renders it inside #markdown-output, and loads highlight.js for code blocks.', 'checks_template': [ 'js: !!document.querySelector("script[src*=\'marked\']")', 'js: !!document.querySelector("script[src*=\'highlight.js\']")', 'js: document.querySelector("#markdown-output").innerHTML.includes("= 2', 'js: document.querySelector("#markdown-source").textContent.trim().length > 0' ] }, { 'template_id': 'github-user-created', 'name': 'GitHub User Information', 'description': 'Create an app that fetches and displays GitHub user information', 'brief_template': 'Publish a Bootstrap page with form id="github-user-{seed}" that fetches a GitHub username, optionally uses ?token=, and displays the account creation date in YYYY-MM-DD UTC inside #github-created-at.', 'checks_template': [ 'js: document.querySelector("#github-user-{seed}").tagName === "FORM"', 'js: document.querySelector("#github-created-at").textContent.includes("20")', 'js: !!document.querySelector("script").textContent.includes("https://api.github.com/users/")' ], 'round2_brief_template': 'Show an aria-live alert #github-status that reports when a lookup starts, succeeds, or fails.', 'round2_checks_template': [ 'js: document.querySelector("#github-status").getAttribute("aria-live") === "polite"', 'js: !!document.querySelector("script").textContent.includes("github-status")' ] } ] with db_manager.get_session() as session: for template_data in templates_data: # Check if template already exists existing = session.query(TaskTemplate).filter( TaskTemplate.template_id == template_data['template_id'] ).first() if not existing: template = TaskTemplate(**template_data) session.add(template) logger.info(f"Created task template: {template_data['template_id']}") session.commit() logger.info("Database initialization completed")