app/models/database.py

"""
Database Models & Setup

Defines SQLAlchemy ORM models for storing image records and generation history.
Uses async SQLite via aiosqlite for non-blocking database operations.
"""

import uuid
from datetime import datetime, timezone
from sqlalchemy import Column, String, Text, DateTime, Integer, Float
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase


class Base(DeclarativeBase):
    """Base class for all ORM models."""
    pass


class ImageRecord(Base):
    """
    Stores metadata for every image in the system (uploaded or generated).
    Each record is linked to a ChromaDB embedding via its ID.
    """
    __tablename__ = "image_records"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    filename = Column(String, nullable=False)           # Stored filename on disk
    original_name = Column(String, nullable=True)       # Original upload filename
    description = Column(Text, nullable=True)           # Vision model description
    prompt = Column(Text, nullable=True)                # User-provided prompt
    generated_prompt = Column(Text, nullable=True)      # Final engineered prompt
    image_url = Column(String, nullable=True)           # URL path to access image
    source = Column(String, default="uploaded")          # 'uploaded' or 'generated'
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))


class GenerationHistory(Base):
    """
    Stores the full pipeline history for each generation request.
    Captures every stage: user input → vision → RAG → prompt → output.
    """
    __tablename__ = "generation_history"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    user_input = Column(Text, nullable=False)            # Original user text/prompt
    input_type = Column(String, default="text")           # 'text' or 'image'
    vision_output = Column(Text, nullable=True)           # Vision model description
    retrieved_context = Column(Text, nullable=True)       # RAG-retrieved similar descriptions
    final_prompt = Column(Text, nullable=True)            # Engineered prompt sent to generator
    generated_image_url = Column(String, nullable=True)   # URL of generated image
    similarity_score = Column(Float, nullable=True)       # Best match score from RAG
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))


# --- Database Engine & Session Factory ---

_engine = None
_session_factory = None


async def init_db(database_url: str) -> None:
    """
    Initialize the async database engine and create all tables.
    Called once during application startup.
    """
    global _engine, _session_factory

    _engine = create_async_engine(database_url, echo=False)
    _session_factory = async_sessionmaker(_engine, class_=AsyncSession, expire_on_commit=False)

    # Create tables if they don't exist
    async with _engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


def get_session_factory() -> async_sessionmaker[AsyncSession]:
    """
    Returns the async session factory.
    Must be called after init_db().
    """
    if _session_factory is None:
        raise RuntimeError("Database not initialized. Call init_db() first.")
    return _session_factory