baodan/api/insurance/models/insurance_document.py

"""Document IR 持久化模型。"""
import json

from sqlalchemy import BigInteger, Column, Float, ForeignKey, Index, Integer, String, Text, TIMESTAMP, UniqueConstraint, func

from insurance.db.compat import db


DOCUMENT_ID_TYPE = BigInteger().with_variant(Integer, "sqlite")


def _json_value(value, fallback):
    if not value:
        return fallback
    try:
        return json.loads(value)
    except (TypeError, ValueError):
        return fallback


class InsuranceDocument(db.Model):
    __tablename__ = "insurance_documents"
    __table_args__ = (
        UniqueConstraint("user_id", "sha256", name="uq_insurance_documents_user_sha256"),
        Index("idx_insurance_documents_user_status", "user_id", "status"),
        Index("idx_insurance_documents_expires", "expires_at"),
    )

    id = Column(DOCUMENT_ID_TYPE, primary_key=True, autoincrement=True)
    user_id = Column(String(64), nullable=False)
    sha256 = Column(String(64), nullable=False)
    storage_key = Column(String(500), nullable=False)
    original_name = Column(String(255), nullable=False)
    mime_type = Column(String(100), nullable=False, default="application/pdf")
    file_size = Column(BigInteger, nullable=False)
    page_count = Column(Integer, nullable=False)
    document_type = Column(String(50), nullable=False, default="plan_illustration")
    pdf_kind = Column(String(20), nullable=False)
    status = Column(String(20), nullable=False)
    parser_version = Column(String(100), nullable=False)
    ocr_version = Column(String(100), nullable=True)
    profile_code = Column(String(100), nullable=True)
    profile_version = Column(String(100), nullable=True)
    quality_json = Column(Text, nullable=True)
    processing_log_json = Column(Text, nullable=True)
    expires_at = Column(TIMESTAMP, nullable=True)
    created_at = Column(TIMESTAMP, server_default=func.now())
    updated_at = Column(TIMESTAMP, server_default=func.now(), onupdate=func.now())

    def to_public_dict(self) -> dict:
        quality = _json_value(self.quality_json, {})
        return {
            "id": self.id,
            "sha256": self.sha256,
            "originalName": self.original_name,
            "mimeType": self.mime_type,
            "fileSize": self.file_size,
            "pageCount": self.page_count,
            "documentType": self.document_type,
            "pdfKind": self.pdf_kind,
            "status": self.status,
            "parserVersion": self.parser_version,
            "ocrVersion": self.ocr_version,
            "profileCode": self.profile_code,
            "profileVersion": self.profile_version,
            "quality": quality,
            "expiresAt": self.expires_at.isoformat() if self.expires_at else None,
            "createdAt": self.created_at.isoformat() if self.created_at else None,
            "updatedAt": self.updated_at.isoformat() if self.updated_at else None,
        }


class InsuranceDocumentPage(db.Model):
    __tablename__ = "insurance_document_pages"
    __table_args__ = (
        UniqueConstraint("document_id", "page_number", name="uq_insurance_document_pages_number"),
        Index("idx_insurance_document_pages_document", "document_id"),
    )

    id = Column(DOCUMENT_ID_TYPE, primary_key=True, autoincrement=True)
    document_id = Column(DOCUMENT_ID_TYPE, ForeignKey("insurance_documents.id", ondelete="CASCADE"), nullable=False)
    page_number = Column(Integer, nullable=False)
    width = Column(Float, nullable=False)
    height = Column(Float, nullable=False)
    native_text_quality = Column(Float, nullable=True)
    ocr_quality = Column(Float, nullable=True)
    page_class = Column(String(30), nullable=False)
    text_blocks_json = Column(Text, nullable=True)
    ocr_blocks_json = Column(Text, nullable=True)
    images_json = Column(Text, nullable=True)
    created_at = Column(TIMESTAMP, server_default=func.now())

    def to_public_dict(self) -> dict:
        return {
            "documentId": self.document_id,
            "pageNumber": self.page_number,
            "width": self.width,
            "height": self.height,
            "pageClass": self.page_class,
            "nativeTextQuality": self.native_text_quality,
            "ocrQuality": self.ocr_quality,
            "textBlocks": _json_value(self.text_blocks_json, []),
            "ocrBlocks": _json_value(self.ocr_blocks_json, []),
            "images": _json_value(self.images_json, []),
        }


class InsuranceDocumentTable(db.Model):
    __tablename__ = "insurance_document_tables"
    __table_args__ = (
        UniqueConstraint("document_id", "table_id", name="uq_insurance_document_tables_table"),
        Index("idx_insurance_document_tables_document", "document_id"),
    )

    id = Column(DOCUMENT_ID_TYPE, primary_key=True, autoincrement=True)
    document_id = Column(DOCUMENT_ID_TYPE, ForeignKey("insurance_documents.id", ondelete="CASCADE"), nullable=False)
    table_id = Column(String(100), nullable=False)
    scenario_type = Column(String(50), nullable=False, default="unknown")
    headers_json = Column(Text, nullable=True)
    rows_json = Column(Text, nullable=True)
    source_pages_json = Column(Text, nullable=True)
    continuation_of = Column(String(100), nullable=True)
    bbox_json = Column(Text, nullable=True)
    conflicts_json = Column(Text, nullable=True)
    created_at = Column(TIMESTAMP, server_default=func.now())

    def to_public_dict(self) -> dict:
        return {
            "documentId": self.document_id,
            "tableId": self.table_id,
            "scenarioType": self.scenario_type,
            "headers": _json_value(self.headers_json, []),
            "rows": _json_value(self.rows_json, []),
            "sourcePages": _json_value(self.source_pages_json, []),
            "continuationOf": self.continuation_of,
            "bbox": _json_value(self.bbox_json, None),
            "conflicts": _json_value(self.conflicts_json, []),
        }