vent_ai_platform/app/db/models/diagnosis.py

from __future__ import annotations

from datetime import datetime

from sqlalchemy import Boolean, DateTime, Float, ForeignKey, Integer, JSON, String, Text
from sqlalchemy.orm import Mapped, mapped_column

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


class DiagnosisRule(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "diagnosis_rules"

    mine_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("mines.id", ondelete="CASCADE"), nullable=True, index=True
    )
    code: Mapped[str] = mapped_column(String(64), unique=True, index=True)
    name: Mapped[str] = mapped_column(String(128))
    scene: Mapped[str] = mapped_column(String(32), index=True)
    defect_type: Mapped[str] = mapped_column(String(64), index=True)
    severity: Mapped[str] = mapped_column(String(16), default="medium")
    condition: Mapped[dict] = mapped_column(JSON, default=dict)
    threshold: Mapped[dict] = mapped_column(JSON, default=dict)
    output_template: Mapped[str] = mapped_column(Text, default="")
    recommendation: Mapped[str] = mapped_column(Text, default="")
    knowledge_refs: Mapped[list] = mapped_column(JSON, default=list)
    enabled: Mapped[bool] = mapped_column(Boolean, default=True)
    version: Mapped[str] = mapped_column(String(32), default="V1.0")


class DiagnosisTask(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "diagnosis_tasks"

    mine_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("mines.id", ondelete="CASCADE"), index=True
    )
    task_no: Mapped[str] = mapped_column(String(64), unique=True, index=True)
    scene: Mapped[str] = mapped_column(String(32), index=True)
    trigger_type: Mapped[str] = mapped_column(String(32), default="manual")
    status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    requested_by: Mapped[str] = mapped_column(String(128), default="")
    model_name: Mapped[str] = mapped_column(String(128), default="")
    input_snapshot: Mapped[dict] = mapped_column(JSON, default=dict)
    error: Mapped[str] = mapped_column(Text, default="")


class DiagnosisResult(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "diagnosis_results"

    task_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("diagnosis_tasks.id", ondelete="CASCADE"), index=True
    )
    rule_id: Mapped[str | None] = mapped_column(
        String(32), ForeignKey("diagnosis_rules.id", ondelete="SET NULL"), nullable=True
    )
    defect_code: Mapped[str] = mapped_column(String(64), index=True)
    defect_name: Mapped[str] = mapped_column(String(255))
    scene: Mapped[str] = mapped_column(String(32), index=True)
    severity: Mapped[str] = mapped_column(String(16), index=True)
    confidence: Mapped[float] = mapped_column(Float, default=0)
    status: Mapped[str] = mapped_column(String(32), default="open", index=True)
    root_cause: Mapped[str] = mapped_column(Text, default="")
    conclusion: Mapped[str] = mapped_column(Text, default="")
    recommendation: Mapped[str] = mapped_column(Text, default="")
    related_objects: Mapped[list] = mapped_column(JSON, default=list)
    raw_output: Mapped[dict] = mapped_column(JSON, default=dict)


class DiagnosisEvidence(UUIDPrimaryKeyMixin, TimestampMixin, Base):
    __tablename__ = "diagnosis_evidence"

    result_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("diagnosis_results.id", ondelete="CASCADE"), index=True
    )
    evidence_type: Mapped[str] = mapped_column(String(32), index=True)
    source_id: Mapped[str] = mapped_column(String(64), default="")
    source_name: Mapped[str] = mapped_column(String(255), default="")
    content: Mapped[str] = mapped_column(Text, default="")
    confidence: Mapped[float | None] = mapped_column(Float, nullable=True)
    occurred_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    metadata_json: Mapped[dict] = mapped_column("metadata", JSON, default=dict)