vent_ai_platform/app/db/models/document.py

from __future__ import annotations

from datetime import datetime

from sqlalchemy import (
    Boolean,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    JSON,
    String,
    Text,
)
from sqlalchemy.dialects.mysql import LONGTEXT
from sqlalchemy.orm import Mapped, mapped_column

from app.db.base import Base, TimestampMixin, UUIDPrimaryKeyMixin


LONG_TEXT = Text().with_variant(LONGTEXT, "mysql")


class BusinessDocument(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "business_documents"

    mine_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("mines.id", ondelete="CASCADE"), index=True
    )
    code: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    title: Mapped[str] = mapped_column(String(255), index=True)
    category: Mapped[str] = mapped_column(String(32), index=True)
    subcategory: Mapped[str] = mapped_column(String(64), default="")
    document_type: Mapped[str] = mapped_column(String(128), default="", index=True)
    file_name: Mapped[str] = mapped_column(String(255))
    file_path: Mapped[str] = mapped_column(String(512))
    file_ext: Mapped[str] = mapped_column(String(16), default="")
    file_size: Mapped[int] = mapped_column(Integer, default=0)
    page_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
    row_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
    owner_name: Mapped[str] = mapped_column(String(128), default="")
    department: Mapped[str] = mapped_column(String(128), default="")
    version: Mapped[str] = mapped_column(String(32), default="V1.0")
    status: Mapped[str] = mapped_column(String(32), default="active", index=True)
    parse_status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
    parse_progress: Mapped[int] = mapped_column(Integer, default=0)
    parse_error: Mapped[str] = mapped_column(Text, default="")
    effective_date: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    expire_date: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    is_confidential: Mapped[bool] = mapped_column(Boolean, default=False)
    tags: Mapped[list] = mapped_column(JSON, default=list)
    summary: Mapped[str] = mapped_column(Text, default="")
    metadata_json: Mapped[dict] = mapped_column("metadata", JSON, default=dict)


class DocumentExtraction(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "document_extractions"

    document_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("business_documents.id", ondelete="CASCADE"), index=True
    )
    field_key: Mapped[str] = mapped_column(String(128), index=True)
    field_name: Mapped[str] = mapped_column(String(128))
    field_value: Mapped[str] = mapped_column(Text, default="")
    normalized_value: Mapped[str] = mapped_column(Text, default="")
    original_value: Mapped[str] = mapped_column(Text, default="")
    value_type: Mapped[str] = mapped_column(String(32), default="text")
    unit: Mapped[str] = mapped_column(String(32), default="")
    confidence: Mapped[float | None] = mapped_column(Float, nullable=True)
    page_no: Mapped[int | None] = mapped_column(Integer, nullable=True)
    evidence_text: Mapped[str] = mapped_column(Text, default="")
    edited: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
    edit_count: Mapped[int] = mapped_column(Integer, default=0)
    extraction_correct: Mapped[bool] = mapped_column(Boolean, default=True, index=True)
    edit_reason: Mapped[str] = mapped_column(Text, default="")
    edit_judgement_source: Mapped[str] = mapped_column(String(32), default="")
    edit_judgement_confidence: Mapped[float | None] = mapped_column(Float, nullable=True)
    edit_judgement_status: Mapped[str] = mapped_column(
        String(32), default="not_required", index=True
    )
    edit_judged_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    verified: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
    verified_by: Mapped[str] = mapped_column(String(128), default="")
    metadata_json: Mapped[dict] = mapped_column("metadata", JSON, default=dict)


class DocumentExtractionTemplate(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "document_extraction_templates"

    code: Mapped[str] = mapped_column(String(64), unique=True, index=True)
    name: Mapped[str] = mapped_column(String(128))
    document_type: Mapped[str] = mapped_column(String(128), index=True)
    description: Mapped[str] = mapped_column(Text, default="")
    version: Mapped[str] = mapped_column(String(32), default="V1.0")
    prompt: Mapped[str] = mapped_column(Text)
    extraction_schema: Mapped[dict] = mapped_column(JSON, default=dict)
    enabled: Mapped[bool] = mapped_column(Boolean, default=True, index=True)


class DocumentProcessingRun(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "document_processing_runs"
    __table_args__ = (
        Index(
            "ix_document_processing_runs_document_created",
            "document_id",
            "created_at",
        ),
    )

    document_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("business_documents.id", ondelete="CASCADE"), index=True
    )
    template_id: Mapped[str | None] = mapped_column(
        String(32),
        ForeignKey("document_extraction_templates.id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )
    status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
    stage: Mapped[str] = mapped_column(String(32), default="uploaded")
    progress: Mapped[int] = mapped_column(Integer, default=0)
    document_type: Mapped[str] = mapped_column(String(128), default="", index=True)
    parser_engine: Mapped[str] = mapped_column(String(32), default="mineru_api_lightweight")
    parsed_markdown_path: Mapped[str] = mapped_column(String(512), default="")
    parsed_markdown: Mapped[str] = mapped_column(LONG_TEXT, default="")
    model_output: Mapped[dict] = mapped_column(JSON, default=dict)
    error: Mapped[str] = mapped_column(Text, default="")
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    metadata_json: Mapped[dict] = mapped_column("metadata", JSON, default=dict)