FabledCurator/alembic/versions/0003_fc2b_ml_pipeline.py

"""fc2b: ML pipeline — allowlist, aliases, centroids, ml_settings

Revision ID: 0003
Revises: 0002
Create Date: 2026-05-15

"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector

revision: str = "0003"
down_revision: Union[str, None] = "0002"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # 3.1 rename wd14_* -> tagger_*
    op.alter_column("image_record", "wd14_predictions", new_column_name="tagger_predictions")
    op.alter_column(
        "image_record", "wd14_model_version", new_column_name="tagger_model_version"
    )

    # 3.2 tag_allowlist
    op.create_table(
        "tag_allowlist",
        sa.Column("tag_id", sa.Integer(), nullable=False),
        sa.Column(
            "min_confidence", sa.Float(), nullable=False, server_default="0.95"
        ),
        sa.Column(
            "added_at", sa.DateTime(timezone=True), nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["tag_id"], ["tag.id"], name="fk_tag_allowlist_tag_id_tag",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("tag_id", name="pk_tag_allowlist"),
        sa.CheckConstraint(
            "min_confidence > 0 AND min_confidence <= 1",
            name="ck_tag_allowlist_confidence_range",
        ),
    )

    # 3.3 tag_suggestion_rejection
    op.create_table(
        "tag_suggestion_rejection",
        sa.Column("image_record_id", sa.Integer(), nullable=False),
        sa.Column("tag_id", sa.Integer(), nullable=False),
        sa.Column(
            "rejected_at", sa.DateTime(timezone=True), nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["image_record_id"], ["image_record.id"],
            name="fk_tsr_image_record_id_image_record", ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["tag_id"], ["tag.id"], name="fk_tsr_tag_id_tag", ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint(
            "image_record_id", "tag_id", name="pk_tag_suggestion_rejection"
        ),
    )
    op.create_index(
        "ix_tag_suggestion_rejection_tag", "tag_suggestion_rejection", ["tag_id"]
    )

    # 3.4 tag_alias
    op.create_table(
        "tag_alias",
        sa.Column("alias_string", sa.String(length=255), nullable=False),
        sa.Column("alias_category", sa.String(length=32), nullable=False),
        sa.Column("canonical_tag_id", sa.Integer(), nullable=False),
        sa.Column(
            "created_at", sa.DateTime(timezone=True), nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["canonical_tag_id"], ["tag.id"],
            name="fk_tag_alias_canonical_tag_id_tag", ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint(
            "alias_string", "alias_category", name="pk_tag_alias"
        ),
    )
    op.create_index("ix_tag_alias_canonical", "tag_alias", ["canonical_tag_id"])

    # 3.5 tag_reference_embedding (centroids)
    op.create_table(
        "tag_reference_embedding",
        sa.Column("tag_id", sa.Integer(), nullable=False),
        sa.Column("embedding", Vector(1152), nullable=False),
        sa.Column("reference_count", sa.Integer(), nullable=False),
        sa.Column("model_version", sa.String(length=128), nullable=False),
        sa.Column(
            "updated_at", sa.DateTime(timezone=True), nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["tag_id"], ["tag.id"],
            name="fk_tag_reference_embedding_tag_id_tag", ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("tag_id", name="pk_tag_reference_embedding"),
    )

    # 3.6 ml_settings singleton
    op.create_table(
        "ml_settings",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column(
            "suggestion_threshold_artist", sa.Float(), nullable=False,
            server_default="0.30",
        ),
        sa.Column(
            "suggestion_threshold_character", sa.Float(), nullable=False,
            server_default="0.50",
        ),
        sa.Column(
            "suggestion_threshold_copyright", sa.Float(), nullable=False,
            server_default="0.50",
        ),
        sa.Column(
            "suggestion_threshold_general", sa.Float(), nullable=False,
            server_default="0.95",
        ),
        sa.Column(
            "centroid_similarity_threshold", sa.Float(), nullable=False,
            server_default="0.55",
        ),
        sa.Column(
            "min_reference_images", sa.Integer(), nullable=False,
            server_default="5",
        ),
        sa.Column(
            "tagger_model_version", sa.String(length=128), nullable=False,
            server_default="camie-tagger-v2",
        ),
        sa.Column(
            "embedder_model_version", sa.String(length=128), nullable=False,
            server_default="siglip-so400m-patch14-384",
        ),
        sa.Column(
            "updated_at", sa.DateTime(timezone=True), nullable=False,
            server_default=sa.func.now(),
        ),
        sa.PrimaryKeyConstraint("id", name="pk_ml_settings"),
        sa.CheckConstraint("id = 1", name="ck_ml_settings_singleton"),
    )
    op.execute("INSERT INTO ml_settings (id) VALUES (1)")


def downgrade() -> None:
    op.drop_table("ml_settings")
    op.drop_table("tag_reference_embedding")
    op.drop_index("ix_tag_alias_canonical", table_name="tag_alias")
    op.drop_table("tag_alias")
    op.drop_index(
        "ix_tag_suggestion_rejection_tag", table_name="tag_suggestion_rejection"
    )
    op.drop_table("tag_suggestion_rejection")
    op.drop_table("tag_allowlist")
    op.alter_column(
        "image_record", "tagger_model_version", new_column_name="wd14_model_version"
    )
    op.alter_column(
        "image_record", "tagger_predictions", new_column_name="wd14_predictions"
    )