resolutionflow/backend/app/models/kb_import.py

import uuid
from datetime import datetime, timezone
from typing import Optional, Any, TYPE_CHECKING
from sqlalchemy import String, Text, DateTime, ForeignKey, Boolean, Integer, Float, CheckConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from sqlalchemy.dialects.postgresql import UUID, JSONB
from app.core.database import Base

if TYPE_CHECKING:
    from app.models.account import Account
    from app.models.user import User
    from app.models.tree import Tree


class KBImport(Base):
    __tablename__ = "kb_imports"
    __table_args__ = (
        CheckConstraint(
            "source_format IN ('txt', 'paste', 'docx', 'pdf', 'html', 'md')",
            name="ck_kb_imports_source_format",
        ),
        CheckConstraint(
            "target_type IN ('troubleshooting', 'procedural')",
            name="ck_kb_imports_target_type",
        ),
        CheckConstraint(
            "status IN ('processing', 'ready', 'committed', 'failed')",
            name="ck_kb_imports_status",
        ),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    account_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("accounts.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    created_by: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("users.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    source_filename: Mapped[Optional[str]] = mapped_column(
        String(500), nullable=True
    )
    source_format: Mapped[str] = mapped_column(String(20), nullable=False)
    source_text: Mapped[str] = mapped_column(Text, nullable=False)
    source_metadata: Mapped[Optional[dict[str, Any]]] = mapped_column(
        JSONB, nullable=True
    )
    target_type: Mapped[str] = mapped_column(String(20), nullable=False)
    status: Mapped[str] = mapped_column(
        String(20), nullable=False, default="processing"
    )
    confidence_avg: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    processing_time_ms: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    ai_tokens_input: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    ai_tokens_output: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    tree_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("trees.id", ondelete="SET NULL"),
        nullable=True,
    )
    batch_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True), nullable=True, index=True
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        onupdate=lambda: datetime.now(timezone.utc),
    )

    # Relationships
    account: Mapped["Account"] = relationship("Account", foreign_keys=[account_id])
    created_by_user: Mapped["User"] = relationship("User", foreign_keys=[created_by])
    tree: Mapped[Optional["Tree"]] = relationship("Tree", foreign_keys=[tree_id])
    nodes: Mapped[list["KBImportNode"]] = relationship(
        "KBImportNode",
        back_populates="kb_import",
        cascade="all, delete-orphan",
        order_by="KBImportNode.node_order",
    )


class KBImportNode(Base):
    __tablename__ = "kb_import_nodes"
    __table_args__ = (
        CheckConstraint(
            "node_type IN ('question', 'resolution', 'step', 'section_header', 'warning', 'action')",
            name="ck_kb_import_nodes_node_type",
        ),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    kb_import_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("kb_imports.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    node_order: Mapped[int] = mapped_column(Integer, nullable=False)
    node_type: Mapped[str] = mapped_column(String(20), nullable=False)
    content: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    parent_node_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("kb_import_nodes.id", ondelete="SET NULL"),
        nullable=True,
    )
    source_excerpt: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    confidence_score: Mapped[float] = mapped_column(Float, nullable=False)
    user_edited: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    user_approved: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        onupdate=lambda: datetime.now(timezone.utc),
    )

    # Relationships
    kb_import: Mapped["KBImport"] = relationship(
        "KBImport", back_populates="nodes"
    )
    parent: Mapped[Optional["KBImportNode"]] = relationship(
        "KBImportNode",
        remote_side="KBImportNode.id",
        foreign_keys=[parent_node_id],
    )