vent_ai_platform/alembic/versions/0004_extraction_edit_accuracy.py

"""add extraction edit accuracy tracking

Revision ID: 0004_extraction_edit_accuracy
Revises: 0003_business_template_lengths
Create Date: 2026-10-12
"""

from __future__ import annotations

import sqlalchemy as sa
from alembic import op


revision = "0004_extraction_edit_accuracy"
down_revision = "0003_business_template_lengths"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column(
        "document_extractions",
        sa.Column("original_value", sa.Text(), nullable=True),
    )
    op.add_column(
        "document_extractions",
        sa.Column("edited", sa.Boolean(), nullable=False, server_default=sa.false()),
    )
    op.add_column(
        "document_extractions",
        sa.Column("edit_count", sa.Integer(), nullable=False, server_default="0"),
    )
    op.add_column(
        "document_extractions",
        sa.Column(
            "extraction_correct",
            sa.Boolean(),
            nullable=False,
            server_default=sa.true(),
        ),
    )
    op.add_column(
        "document_extractions",
        sa.Column("edit_reason", sa.Text(), nullable=True),
    )
    op.add_column(
        "document_extractions",
        sa.Column(
            "edit_judgement_source",
            sa.String(length=32),
            nullable=False,
            server_default="",
        ),
    )
    op.add_column(
        "document_extractions",
        sa.Column("edit_judgement_confidence", sa.Float(), nullable=True),
    )
    op.add_column(
        "document_extractions",
        sa.Column("edit_judged_at", sa.DateTime(timezone=True), nullable=True),
    )

    op.execute("UPDATE document_extractions SET original_value = '' WHERE original_value IS NULL")
    op.execute("UPDATE document_extractions SET edit_reason = '' WHERE edit_reason IS NULL")
    op.alter_column(
        "document_extractions",
        "original_value",
        existing_type=sa.Text(),
        nullable=False,
    )
    op.alter_column(
        "document_extractions",
        "edit_reason",
        existing_type=sa.Text(),
        nullable=False,
    )
    op.create_index(
        "ix_document_extractions_edited",
        "document_extractions",
        ["edited"],
    )
    op.create_index(
        "ix_document_extractions_extraction_correct",
        "document_extractions",
        ["extraction_correct"],
    )


def downgrade() -> None:
    op.drop_index(
        "ix_document_extractions_extraction_correct",
        table_name="document_extractions",
    )
    op.drop_index(
        "ix_document_extractions_edited",
        table_name="document_extractions",
    )
    op.drop_column("document_extractions", "edit_judged_at")
    op.drop_column("document_extractions", "edit_judgement_source")
    op.drop_column("document_extractions", "edit_judgement_confidence")
    op.drop_column("document_extractions", "edit_reason")
    op.drop_column("document_extractions", "extraction_correct")
    op.drop_column("document_extractions", "edit_count")
    op.drop_column("document_extractions", "edited")
    op.drop_column("document_extractions", "original_value")