decision/engine/models.py

"""
Core data models for the Decision Engine.

All models are immutable dataclasses to prevent accidental mutation
during safety-critical evaluation pipelines.
"""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class ConfidenceTier(enum.Enum):
    """Confidence tier for taxonomy trigger matches."""
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    NONE = "none"

    def __ge__(self, other: ConfidenceTier) -> bool:
        order = {self.HIGH: 3, self.MEDIUM: 2, self.LOW: 1, self.NONE: 0}
        return order[self] >= order[other]

    def __gt__(self, other: ConfidenceTier) -> bool:
        order = {self.HIGH: 3, self.MEDIUM: 2, self.LOW: 1, self.NONE: 0}
        return order[self] > order[other]

    def __le__(self, other: ConfidenceTier) -> bool:
        order = {self.HIGH: 3, self.MEDIUM: 2, self.LOW: 1, self.NONE: 0}
        return order[self] <= order[other]

    def __lt__(self, other: ConfidenceTier) -> bool:
        order = {self.HIGH: 3, self.MEDIUM: 2, self.LOW: 1, self.NONE: 0}
        return order[self] < order[other]


class RiskClass(enum.Enum):
    """Clinical risk classification tier."""
    R0 = "R0"  # No clinical concern
    R1 = "R1"  # Low-level concern, routine follow-up
    R2 = "R2"  # Moderate concern, warm transfer to nurse
    R3 = "R3"  # Emergent, instruct to call 911 + transfer

    @property
    def is_emergent(self) -> bool:
        return self == RiskClass.R3

    @property
    def requires_nurse(self) -> bool:
        return self in (RiskClass.R2, RiskClass.R3)

    @property
    def severity_rank(self) -> int:
        return {self.R0: 0, self.R1: 1, self.R2: 2, self.R3: 3}[self]

    def __ge__(self, other: RiskClass) -> bool:
        return self.severity_rank >= other.severity_rank

    def __gt__(self, other: RiskClass) -> bool:
        return self.severity_rank > other.severity_rank

    def __le__(self, other: RiskClass) -> bool:
        return self.severity_rank <= other.severity_rank

    def __lt__(self, other: RiskClass) -> bool:
        return self.severity_rank < other.severity_rank


class MatchType(enum.Enum):
    """Type of lexical match that fired."""
    PHRASE = "phrase"
    REGEX = "regex"
    PARTIAL = "partial"


class NegationAction(enum.Enum):
    """What to do when negation is detected."""
    SUPPRESS = "suppress"              # Remove domain nomination entirely
    DOWNGRADE_CONFIDENCE = "downgrade_confidence"  # Drop confidence one tier


class TurnOutcome(enum.Enum):
    """Possible outcomes after a conversational turn."""
    PROCEED = "proceed"
    CLARIFY = "clarify"
    ESCALATE = "escalate"
    HANDOFF = "handoff"
    END = "end"


class FlowType(enum.Enum):
    """Type of conversation flow."""
    SCREENING = "screening"
    ESCALATION = "escalation"
    CLARIFICATION = "clarification"
    RESOLUTION = "resolution"
    HANDOFF = "handoff"
    PRIMITIVE = "primitive"


class CallPhase(enum.Enum):
    """Phase of the current call."""
    SESSION_START = "session_start"
    ACTIVE_CALL = "active_call"
    ENDED = "ended"


# ---------------------------------------------------------------------------
# Trigger & Match Models
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class TriggerMatch:
    """A single lexical match from a taxonomy trigger."""
    domain: str
    confidence_tier: ConfidenceTier
    match_type: MatchType
    matched_text: str          # The phrase/pattern that matched
    matched_span: str          # The actual text span from patient input
    priority: int              # Domain priority from triggers.yaml


@dataclass(frozen=True)
class NegationResult:
    """Result of negation analysis on a domain match."""
    is_negated: bool
    action: NegationAction
    matched_pattern: str = ""  # Which negation pattern matched
    original_confidence: ConfidenceTier = ConfidenceTier.NONE
    adjusted_confidence: ConfidenceTier = ConfidenceTier.NONE


@dataclass(frozen=True)
class DomainNomination:
    """A nominated clinical domain with combined confidence from ML + lexical signals."""
    domain: str
    display_name: str
    confidence_tier: ConfidenceTier
    priority: int                          # From triggers.yaml priority field
    target_risk_class: Optional[RiskClass] # Default risk from rules.yaml
    trigger_matches: tuple                 # Tuple of TriggerMatch (frozen)
    negation_result: Optional[NegationResult] = None
    ml_label: Optional[str] = None         # DriveHealthBERT predicted label
    ml_confidence: Optional[float] = None  # DriveHealthBERT probability
    recommended_flow: Optional[str] = None # From rules.yaml decision_rules

    @property
    def is_negated(self) -> bool:
        return self.negation_result is not None and self.negation_result.is_negated

    @property
    def is_safety_critical(self) -> bool:
        """R3 domains or hard-escalate domains."""
        return self.target_risk_class == RiskClass.R3


# ---------------------------------------------------------------------------
# Risk Assessment Models
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class RiskRuleMatch:
    """A risk escalation rule that fired."""
    rule_id: str
    domain: str
    risk_class: RiskClass
    conditions_met: Dict[str, Any] = field(default_factory=dict)
    context_conditions: Optional[Dict[str, Any]] = None


@dataclass(frozen=True)
class SuppressionResult:
    """Result of domain suppression evaluation."""
    suppressed: bool
    suppressor_domain: Optional[str] = None
    rule_reason: Optional[str] = None


@dataclass(frozen=True)
class RiskAssessment:
    """Complete risk assessment for a patient utterance."""
    # Primary result
    risk_class: RiskClass
    primary_domain: Optional[str]
    # All nominated domains (ranked by priority)
    domain_nominations: tuple           # Tuple of DomainNomination
    # Which rules fired
    matched_rules: tuple                # Tuple of RiskRuleMatch
    # Suppression results
    suppressions: tuple = ()            # Tuple of SuppressionResult
    # Metadata
    hard_escalate: bool = False         # suicidal_ideation / homicidal_ideation
    safety_override: bool = False       # Lexical safety signal regardless of ML
    ml_label: Optional[str] = None
    ml_confidence: Optional[float] = None
    recommended_flow: Optional[str] = None
    recommended_action: Optional[TurnOutcome] = None

    @property
    def requires_911(self) -> bool:
        return self.risk_class == RiskClass.R3

    @property
    def requires_nurse_transfer(self) -> bool:
        return self.risk_class >= RiskClass.R2

    def to_dict(self) -> Dict[str, Any]:
        """Serialize for API response."""
        return {
            "risk_class": self.risk_class.value,
            "primary_domain": self.primary_domain,
            "requires_911": self.requires_911,
            "requires_nurse_transfer": self.requires_nurse_transfer,
            "hard_escalate": self.hard_escalate,
            "safety_override": self.safety_override,
            "ml_label": self.ml_label,
            "ml_confidence": self.ml_confidence,
            "recommended_flow": self.recommended_flow,
            "recommended_action": self.recommended_action.value if self.recommended_action else None,
            "domain_nominations": [
                {
                    "domain": n.domain,
                    "confidence_tier": n.confidence_tier.value,
                    "priority": n.priority,
                    "target_risk_class": n.target_risk_class.value if n.target_risk_class else None,
                    "is_negated": n.is_negated,
                    "recommended_flow": n.recommended_flow,
                    "trigger_match_count": len(n.trigger_matches),
                }
                for n in self.domain_nominations
            ],
            "matched_rules": [
                {
                    "rule_id": r.rule_id,
                    "domain": r.domain,
                    "risk_class": r.risk_class.value,
                }
                for r in self.matched_rules
            ],
        }


# ---------------------------------------------------------------------------
# Flow Execution Models (Phase 2 — stubs for forward compatibility)
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ResponseSpec:
    """LLM response constraints for a flow step."""
    spec_id: str
    goal: str
    tone: str
    must_include: tuple = ()
    must_ask: tuple = ()
    must_not: tuple = ()
    max_questions: int = 0


@dataclass(frozen=True)
class FlowStep:
    """A single step in a conversation flow."""
    step_number: int
    step_type: str              # "ask", "respond", "handoff"
    response_spec: ResponseSpec
    collects: tuple = ()        # Slot names to extract
    handoff_target: Optional[str] = None
    handoff_urgency: Optional[str] = None


@dataclass(frozen=True)
class ConversationFlow:
    """A complete conversation flow definition."""
    flow_id: str
    flow_type: FlowType
    domain: Optional[str]
    required_slots: tuple = ()
    steps: tuple = ()           # Tuple of FlowStep
    exit_rules: Dict[str, str] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Session State Models (Phase 2 — stubs for forward compatibility)
# ---------------------------------------------------------------------------

@dataclass
class SlotState:
    """Mutable slot tracker for an active conversation."""
    slots: Dict[str, Any] = field(default_factory=dict)

    def set_slot(self, name: str, value: Any) -> None:
        self.slots[name] = value

    def get_slot(self, name: str) -> Optional[Any]:
        return self.slots.get(name)

    def has_slot(self, name: str) -> bool:
        return name in self.slots

    def has_all(self, slot_names: List[str]) -> bool:
        return all(name in self.slots for name in slot_names)


@dataclass
class SessionState:
    """Mutable state for an active call session."""
    session_id: str
    patient_id: Optional[str] = None
    tenant_id: Optional[str] = None
    call_phase: CallPhase = CallPhase.SESSION_START
    current_flow: Optional[str] = None
    current_step: int = 0
    agenda_position: int = 0
    slots: SlotState = field(default_factory=SlotState)
    domain_history: List[str] = field(default_factory=list)
    risk_history: List[RiskAssessment] = field(default_factory=list)
    turn_count: int = 0