resolutionflow/backend/app/models/ai_session.py

"""AI-powered troubleshooting session model.

Represents a complete FlowPilot interaction from intake to resolution/escalation.
This is the central entity of the FlowPilot-First pivot.
"""
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
import sqlalchemy as sa
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.user import User
    from app.models.team import Team
    from app.models.account import Account
    from app.models.tree import Tree
    from app.models.psa_connection import PsaConnection
    from app.models.session_branch import SessionBranch
    from app.models.fork_point import ForkPoint
    from app.models.session_handoff import SessionHandoff
    from app.models.session_resolution_output import SessionResolutionOutput


class AISession(Base):
    """A FlowPilot-guided troubleshooting session.

    Lifecycle: active → resolved | escalated | abandoned
    Sessions may be paused and resumed (e.g., escalation handoff).
    """
    __tablename__ = "ai_sessions"
    __table_args__ = (
        CheckConstraint(
            "intake_type IN ('free_text', 'psa_ticket', 'screenshot', 'log_paste', 'combined')",
            name="ck_ai_sessions_intake_type",
        ),
        CheckConstraint(
            "status IN ('active', 'paused', 'resolved', 'escalated', 'requesting_escalation', 'abandoned')",
            name="ck_ai_sessions_status",
        ),
        CheckConstraint(
            "confidence_tier IN ('guided', 'exploring', 'discovery')",
            name="ck_ai_sessions_confidence_tier",
        ),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    user_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("users.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    account_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("accounts.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    team_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("teams.id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )

    # ── Session type ──
    session_type: Mapped[str] = mapped_column(
        String(10), nullable=False, default="guided", index=True,
        comment="Session type: guided (FlowPilot) or chat (assistant)",
    )
    title: Mapped[Optional[str]] = mapped_column(
        String(255), nullable=True,
        comment="Display title for chat sessions; guided sessions use problem_summary",
    )

    # ── Intake ──
    intake_type: Mapped[str] = mapped_column(
        String(20), nullable=False, default="free_text"
    )
    intake_content: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False, default=dict,
        comment="Original intake data: {text, image_urls, log_content, ticket_data}",
    )
    problem_summary: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="AI-generated one-line problem summary from intake",
    )
    problem_domain: Mapped[Optional[str]] = mapped_column(
        String(100), nullable=True,
        comment="Classified domain: active_directory, networking, m365, hardware, etc.",
    )

    # ── Session state ──
    status: Mapped[str] = mapped_column(
        String(30), nullable=False, default="active", index=True,
    )
    confidence_tier: Mapped[str] = mapped_column(
        String(20), nullable=False, default="discovery",
        comment="Current AI confidence: guided (>80%), exploring (40-80%), discovery (<40%)",
    )
    confidence_score: Mapped[float] = mapped_column(
        Float, nullable=False, default=0.0,
        comment="Numeric confidence 0.0-1.0 for internal tracking",
    )

    # ── Flow matching ──
    matched_flow_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("trees.id", ondelete="SET NULL"),
        nullable=True,
        comment="If following an existing flow, which one",
    )
    match_score: Mapped[Optional[float]] = mapped_column(
        Float, nullable=True,
        comment="Similarity score of the matched flow (0.0-1.0)",
    )

    # ── PSA link ──
    psa_ticket_id: Mapped[Optional[str]] = mapped_column(
        String(100), nullable=True,
        comment="External PSA ticket ID if session was started from a ticket",
    )
    psa_connection_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("psa_connections.id", ondelete="SET NULL"),
        nullable=True,
    )
    ticket_data: Mapped[Optional[dict[str, Any]]] = mapped_column(
        JSONB, nullable=True,
        comment="Snapshot of PSA ticket data at session start",
    )

    # ── Triage / Cockpit Header ──
    client_name: Mapped[Optional[str]] = mapped_column(
        String(255), nullable=True,
        comment="MSP client name for incident header (AI-inferred or manual)",
    )
    asset_name: Mapped[Optional[str]] = mapped_column(
        String(255), nullable=True,
        comment="Device, asset, or user being worked on",
    )
    issue_category: Mapped[Optional[str]] = mapped_column(
        String(100), nullable=True,
        comment="Human-readable category (e.g. DNS / Networking)",
    )
    triage_hypothesis: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="Current working hypothesis — AI-updated + engineer-editable",
    )
    evidence_items: Mapped[Optional[list[dict[str, Any]]]] = mapped_column(
        JSONB, nullable=True,
        comment='What We Know list: [{"text": str, "status": "confirmed"|"ruled_out"|"pending"}]',
    )

    # ── Resolution / Escalation ──
    resolution_summary: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="What fixed the issue (set on resolution)",
    )
    resolution_action: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="The specific action/step that resolved the issue",
    )
    escalation_reason: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="Why escalated (set on escalation)",
    )
    escalation_package: Mapped[Optional[dict[str, Any]]] = mapped_column(
        JSONB, nullable=True,
        comment="Context package for receiving engineer: steps_tried, hypotheses, suggestions",
    )
    escalated_to_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("users.id", ondelete="SET NULL"),
        nullable=True,
    )

    # ── Feedback ──
    session_rating: Mapped[Optional[int]] = mapped_column(
        Integer, nullable=True,
        comment="1-5 engineer feedback rating",
    )
    session_feedback: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="Optional feedback text from engineer",
    )

    # ── Knowledge Flywheel ──
    analysis_status: Mapped[Optional[str]] = mapped_column(
        String(20), nullable=True,
        comment="Knowledge Flywheel status: null (N/A), pending, completed, failed",
    )

    # ── AI tracking ──
    total_input_tokens: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
    )
    total_output_tokens: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
    )
    step_count: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
    )

    # ── Timestamps ──
    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),
    )
    resolved_at: Mapped[Optional[datetime]] = mapped_column(
        DateTime(timezone=True), nullable=True,
    )

    # ── LLM conversation context ──
    system_prompt_snapshot: Mapped[Optional[str]] = mapped_column(
        Text, nullable=True,
        comment="Snapshot of the system prompt used (for debugging/training)",
    )
    conversation_messages: Mapped[list[dict[str, Any]]] = mapped_column(
        JSONB, nullable=False, default=list,
        comment="Full LLM message history for context continuity",
    )
    pending_task_lane: Mapped[Optional[dict[str, Any]]] = mapped_column(
        JSONB, nullable=True,
        comment="Current task lane state: {questions: [...], actions: [...]}",
    )

    # ── Branching ──
    is_branching: Mapped[bool] = mapped_column(
        default=False,
        comment="Whether conversational branching is active for this session",
    )
    active_branch_id: Mapped[Optional[uuid.UUID]] = mapped_column(
        UUID(as_uuid=True), nullable=True,
        comment="Currently viewed branch. No FK — soft pointer to avoid circular FK with session_branches",
    )
    handoff_count: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
        comment="Number of times this session has been handed off",
    )
    total_active_seconds: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
        comment="Cumulative active time in seconds",
    )
    total_parked_seconds: Mapped[int] = mapped_column(
        Integer, nullable=False, default=0,
        comment="Cumulative parked time in seconds",
    )

    # ── Relationships ──
    user: Mapped["User"] = relationship("User", foreign_keys=[user_id])
    account: Mapped["Account"] = relationship("Account")
    team: Mapped[Optional["Team"]] = relationship("Team")
    matched_flow: Mapped[Optional["Tree"]] = relationship("Tree", foreign_keys=[matched_flow_id])
    escalated_to: Mapped[Optional["User"]] = relationship("User", foreign_keys=[escalated_to_id])
    psa_connection: Mapped[Optional["PsaConnection"]] = relationship("PsaConnection")
    steps: Mapped[list["AISessionStep"]] = relationship(
        "AISessionStep", back_populates="session",
        cascade="all, delete-orphan",
        order_by="AISessionStep.step_order",
    )
    branches: Mapped[list["SessionBranch"]] = relationship(
        "SessionBranch",
        foreign_keys="SessionBranch.session_id",
        back_populates="session",
        cascade="all, delete-orphan",
        order_by="SessionBranch.branch_order",
    )
    handoffs: Mapped[list["SessionHandoff"]] = relationship(
        "SessionHandoff",
        back_populates="session",
        cascade="all, delete-orphan",
        order_by="SessionHandoff.created_at",
    )
    resolution_outputs: Mapped[list["SessionResolutionOutput"]] = relationship(
        "SessionResolutionOutput",
        back_populates="session",
        cascade="all, delete-orphan",
    )