app/models/db.py

"""
SQLModel ORM table definitions and the SQLAlchemy engine singleton.

Tables
------
User         — accounts (email, hashed password, role, active flag).
Job          — file processing jobs with full status/step/retry/error tracking.
UsageLog     — one row per LLM/Whisper/embed API call (tokens, latency, model).
QueryHistory — one row per RAG query (answer, citations, RAGAS scores).

Engine
------
get_engine() returns a module-level singleton so the connection pool is shared
across all Celery tasks and FastAPI requests in the same process.
"""

import enum
import uuid
from datetime import datetime
from typing import Optional

from sqlmodel import Field, SQLModel, create_engine, Session
from sqlalchemy import text


# ── Enums ──────────────────────────────────────────────────────────────────────

class UserRole(str, enum.Enum):
    admin = "admin"
    user = "user"


class JobStatus(str, enum.Enum):
    pending = "PENDING"
    processing = "PROCESSING"
    completed = "COMPLETED"
    failed = "FAILED"
    failed_permanent = "FAILED_PERMANENT"


class ErrorType(str, enum.Enum):
    rate_limit = "RATE_LIMIT"
    invalid_input = "INVALID_INPUT"
    unknown = "UNKNOWN"


# ── Tables ─────────────────────────────────────────────────────────────────────

class User(SQLModel, table=True):
    __tablename__ = "users"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    email: str = Field(unique=True, index=True)
    hashed_password: str
    role: UserRole = Field(default=UserRole.user)
    is_active: bool = Field(default=True)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    last_active_at: Optional[datetime] = Field(default=None)


class Job(SQLModel, table=True):
    __tablename__ = "jobs"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    user_id: uuid.UUID = Field(foreign_key="users.id", index=True)
    filename: str
    file_type: str
    file_path: str
    file_size_bytes: int
    status: JobStatus = Field(default=JobStatus.pending)
    step: Optional[str] = Field(default=None)
    retry_count: int = Field(default=0)
    error_type: Optional[ErrorType] = Field(default=None)
    error_message: Optional[str] = Field(default=None)
    result: Optional[str] = Field(default=None)
    chunk_count: Optional[int] = Field(default=None)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)


class UsageLog(SQLModel, table=True):
    __tablename__ = "usage_logs"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    user_id: Optional[uuid.UUID] = Field(default=None, foreign_key="users.id", index=True)
    job_id: Optional[uuid.UUID] = Field(default=None, foreign_key="jobs.id", index=True)
    endpoint: str
    model: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    latency_ms: int
    query_text: Optional[str] = Field(default=None)
    llm_response_preview: Optional[str] = Field(default=None)
    created_at: datetime = Field(default_factory=datetime.utcnow)


class QueryHistory(SQLModel, table=True):
    __tablename__ = "query_history"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    user_id: uuid.UUID = Field(foreign_key="users.id", index=True)
    question: str
    answer: str
    citations: str
    job_ids_queried: str
    chunk_count_retrieved: int
    avg_similarity_score: float
    confidence_gate_passed: bool
    prompt_tokens: int
    completion_tokens: int
    latency_ms: int
    ragas_scores: Optional[str] = Field(default=None)
    ragas_computed_at: Optional[datetime] = Field(default=None)
    created_at: datetime = Field(default_factory=datetime.utcnow)


# ── Engine helpers ─────────────────────────────────────────────────────────────

_engine = None


def get_engine():
    global _engine
    if _engine is None:
        from app.config import settings
        _engine = create_engine(
            settings.DATABASE_URL,
            echo=False,
            pool_size=10,
            max_overflow=20,
            pool_pre_ping=True,
        )
    return _engine


def create_db_and_tables():
    SQLModel.metadata.create_all(get_engine())