PTSRaiAgent/app/core/rules.py

"""Deterministic alert & agent-action engine (ADR-007, guardrail #3).

The LLM never decides whether to alert. Structured data (daily check-ins, doses, scale totals, C-SSRS) is
packed into a `PatientSnapshot`; `evaluate()` runs the rules and returns `Action`s that the orchestrator
executes (create an alert, start the crisis script, schedule a scale, send medication info...). The LLM only
phrases the reply around an action that has already been chosen.

How a rule is defined — the YAML/code split:
  * `config/rules.yaml` declares each rule: `id`, description, thresholds/windows (`params`), the actions it
    produces (type, level, extra params) and `reason_ru` for the doctor. Clinicians edit only this file.
  * This module implements the predicate of each rule as a pure function registered under the same id with
    `@evaluator("<id>")`. It returns `None` when the rule does not fire, or a dict of computed values
    (e.g. `{"total": 22}`) that is merged into action params and used to format `reason_ru`.
  * Loading validates both directions: a YAML id without an evaluator is an error (no silently dead rules).

Determinism: evaluators read time only from `snapshot.as_of` (the virtual "now" from `app.core.clock`),
never from the wall clock, so the same snapshot always yields the same actions.

Dedup: when a rule would emit an `alert` whose code is already open (not acked by the doctor) the whole rule
stays silent. `crisis` actions are never suppressed — the orchestrator is responsible for not restarting an
already running crisis script.

Colours (ADR-006): ACUTE → red; HIGH/MODERATE → yellow; INFO only → grey; no open alerts → green.
"""

from __future__ import annotations

from collections.abc import Callable, Iterable
from datetime import date, datetime, timedelta
from functools import lru_cache
from pathlib import Path
from typing import Any, Literal

import yaml
from pydantic import BaseModel, Field, field_validator, model_validator

from app.config import get_settings
from app.core.types import AlertLevel

# ---------------------------------------------------------------------------------------------- snapshot

Module = Literal["veteran", "depression", "family"]
SleepBucket = Literal["<4", "4-6", "6-8", ">8"]
AlcoholBucket = Literal["none", "1-2", "more"]
DoseStatus = Literal["taken", "missed", "skipped", "pending"]
CssrsRisk = Literal["none", "low", "moderate", "high"]
ActionType = Literal[
    "alert",
    "crisis",
    "schedule_scale",
    "reach_out",
    "med_info",
    "suggest_visit",
    "show_safety_plan",
    "enhanced_monitoring",
]
StatusColor = Literal["red", "yellow", "green", "grey"]

_CSSRS_RANK: dict[str, int] = {"none": 0, "low": 1, "moderate": 2, "high": 3}
_LEVELS: dict[str, AlertLevel] = {
    "info": AlertLevel.INFO,
    "moderate": AlertLevel.MODERATE,
    "high": AlertLevel.HIGH,
    "acute": AlertLevel.ACUTE,
}


class DailyCheckin(BaseModel):
    """One day of the daily check-in (ADR-002). `None` = not asked/not answered that day."""

    day: date
    mood: int | None = Field(default=None, ge=0, le=10)
    med_taken: bool | None = None
    sleep_hours_bucket: SleepBucket | None = None
    alcohol: AlcoholBucket | None = None
    side_effects: list[str] = Field(default_factory=list)  # normalised codes from `extract`, e.g. "nausea"
    quit_intent: bool = False  # "I'll stop taking the pills" detected in free text
    nightmares: bool | None = None


class ScalePoint(BaseModel):
    """A completed scale: only the total (+ PHQ-9 item 9) is needed by the rules."""

    scale: str  # phq9 | pcl5 | gad7 | auditc | ...
    ts: datetime
    total: int = Field(ge=0)
    item9: int | None = Field(default=None, ge=0, le=3)
    scheduled: bool = True  # False = unscheduled (triggered by a rule)


class DoseRecord(BaseModel):
    planned_at: datetime
    status: DoseStatus


class PatientSnapshot(BaseModel):
    """Everything the rules may look at, assembled by the caller from the clinical DB (pseudonymised)."""

    token: str
    as_of: datetime  # virtual "now" (app.core.clock.now()); all windows are counted back from here
    age: int = Field(ge=0)
    sex: Literal["m", "f"]
    modules: set[Module] = Field(default_factory=set)
    therapy_day: int = 0  # day 0 = visit where therapy was prescribed
    checkins: list[DailyCheckin] = Field(default_factory=list)
    scale_results: list[ScalePoint] = Field(default_factory=list)
    doses: list[DoseRecord] = Field(default_factory=list)
    cssrs_risk: CssrsRisk | None = None
    weapon_access: bool | None = None
    last_patient_message_at: datetime | None = None
    open_alert_codes: set[str] = Field(default_factory=set)  # alerts not yet acked by the doctor

    @field_validator("as_of")
    @classmethod
    def _aware(cls, v: datetime) -> datetime:
        if v.tzinfo is None:
            raise ValueError("as_of must be timezone-aware (take it from app.core.clock.now())")
        return v


class Action(BaseModel):
    """One output of the rule engine, executed by the orchestrator."""

    type: ActionType
    level: AlertLevel | None = None
    code: str  # stable reason code = rule id (dedup key, stored with the alert)
    params: dict[str, Any] = Field(default_factory=dict)
    reason_ru: str = ""  # short text for the doctor dashboard

    @property
    def creates_alert(self) -> bool:
        """`alert` and `crisis` actions are persisted as an alert record (crisis = acute alert)."""
        return self.type in ("alert", "crisis")


# ---------------------------------------------------------------------------------------------- config


class ActionSpec(BaseModel):
    type: ActionType
    level: Literal["info", "moderate", "high", "acute"] | None = None
    params: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def _alert_needs_level(self) -> ActionSpec:
        if self.type in ("alert", "crisis") and self.level is None:
            raise ValueError(f"action '{self.type}' requires a level")
        return self


class RuleSpec(BaseModel):
    id: str
    description: str = ""
    enabled: bool = True
    params: dict[str, Any] = Field(default_factory=dict)
    actions: list[ActionSpec] = Field(min_length=1)
    reason_ru: str = ""


class RulesConfig(BaseModel):
    rules: list[RuleSpec]

    @model_validator(mode="after")
    def _check_ids(self) -> RulesConfig:
        ids = [r.id for r in self.rules]
        dupes = {i for i in ids if ids.count(i) > 1}
        if dupes:
            raise ValueError(f"duplicate rule ids: {sorted(dupes)}")
        unknown = set(ids) - set(_REGISTRY)
        if unknown:
            raise ValueError(f"rules without an evaluator in app/core/rules.py: {sorted(unknown)}")
        return self


def load_rules(path: Path | None = None) -> RulesConfig:
    """Parse and validate a rules YAML (default: `<config_dir>/rules.yaml`, cached)."""
    if path is None:
        return _default_rules()
    return RulesConfig.model_validate(yaml.safe_load(path.read_text(encoding="utf-8")))


@lru_cache
def _default_rules() -> RulesConfig:
    return load_rules(get_settings().config_dir / "rules.yaml")


# ---------------------------------------------------------------------------------------------- registry

Context = dict[str, Any]
Evaluator = Callable[[PatientSnapshot, dict[str, Any]], Context | None]
_REGISTRY: dict[str, Evaluator] = {}


def evaluator(rule_id: str) -> Callable[[Evaluator], Evaluator]:
    """Register the predicate implementing YAML rule `rule_id`."""

    def deco(fn: Evaluator) -> Evaluator:
        if rule_id in _REGISTRY:
            raise ValueError(f"evaluator for '{rule_id}' already registered")
        _REGISTRY[rule_id] = fn
        return fn

    return deco


def registered_rule_ids() -> set[str]:
    return set(_REGISTRY)


# ---------------------------------------------------------------------------------------------- helpers


def _today(s: PatientSnapshot) -> date:
    return s.as_of.date()


def _checkins_by_day(s: PatientSnapshot) -> dict[date, DailyCheckin]:
    return {c.day: c for c in s.checkins if c.day <= _today(s)}


def _anchor_day(s: PatientSnapshot) -> date:
    """Last day a "N days in a row" streak may end on.

    Today if today's check-in exists, otherwise yesterday — so a morning evaluation (before today's check-in)
    still sees yesterday's streak, but an old streak does not keep firing forever.
    """
    today = _today(s)
    return today if today in _checkins_by_day(s) else today - timedelta(days=1)


def _streak(s: PatientSnapshot, days: int, pred: Callable[[DailyCheckin], bool]) -> list[DailyCheckin] | None:
    """Check-ins of `days` consecutive calendar days ending at the anchor day, if all satisfy `pred`.

    A missing day breaks the streak: we never assume a symptom on a day the patient did not report.
    """
    by_day = _checkins_by_day(s)
    end = _anchor_day(s)
    run: list[DailyCheckin] = []
    for i in range(days):
        c = by_day.get(end - timedelta(days=i))
        if c is None or not pred(c):
            return None
        run.append(c)
    return run


def _history(s: PatientSnapshot, scale: str) -> list[ScalePoint]:
    """Results of `scale` up to `as_of`, oldest first (results "from the future" are ignored)."""
    return sorted((p for p in s.scale_results if p.scale == scale and p.ts <= s.as_of), key=lambda p: p.ts)


def _recent_item9(s: PatientSnapshot, threshold: int, recent_hours: float) -> int | None:
    """Max PHQ-9 item 9 among PHQ-9s completed within `recent_hours`, if it reaches `threshold`.

    Limited to recent results so one old positive answer does not restart the crisis protocol on every
    evaluation; after that the open acute alert keeps the patient red until the doctor acts.
    """
    since = s.as_of - timedelta(hours=recent_hours)
    values = [p.item9 for p in _history(s, "phq9") if p.ts > since and p.item9 is not None]
    top = max(values, default=None)
    return top if top is not None and top >= threshold else None


# ---------------------------------------------------------------------------------------------- evaluators
# Each function implements the predicate of the YAML rule with the same id; thresholds come from `p`.


@evaluator("phq9_item9")
def _phq9_item9(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    # Item 9 is collected only by buttons (never LLM-interpreted), so the value is the patient's own answer.
    item9 = _recent_item9(s, p["threshold"], p["recent_hours"])
    return None if item9 is None else {"item9": item9}


@evaluator("cssrs_high")
def _cssrs_high(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    return {} if s.cssrs_risk == "high" else None


@evaluator("weapon_access_risk")
def _weapon_access_risk(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    # Access to means alone is not a crisis; combined with any suicide-risk signal it is.
    if not s.weapon_access:
        return None
    item9_signal = _recent_item9(s, p["item9_threshold"], p["recent_hours"]) is not None
    cssrs_signal = _CSSRS_RANK.get(s.cssrs_risk or "none", 0) >= _CSSRS_RANK[p["cssrs_min"]]
    return {} if item9_signal or cssrs_signal else None


@evaluator("cssrs_moderate")
def _cssrs_moderate(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    return {} if s.cssrs_risk == "moderate" else None


def _latest_at_least(scale: str) -> Evaluator:
    """Latest total of `scale` >= params.threshold."""

    def fn(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
        hist = _history(s, scale)
        if hist and hist[-1].total >= p["threshold"]:
            return {"total": hist[-1].total}
        return None

    return fn


def _rise_between_last_two(scale: str) -> Evaluator:
    """Latest total minus the previous total of `scale` >= params.min_rise (worsening)."""

    def fn(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
        hist = _history(s, scale)
        if len(hist) < 2:
            return None
        prev, last = hist[-2].total, hist[-1].total
        if last - prev >= p["min_rise"]:
            return {"previous": prev, "total": last, "rise": last - prev}
        return None

    return fn


evaluator("phq9_severe")(_latest_at_least("phq9"))
evaluator("pcl5_severe")(_latest_at_least("pcl5"))
evaluator("phq9_rise")(_rise_between_last_two("phq9"))
evaluator("pcl5_rise")(_rise_between_last_two("pcl5"))


@evaluator("auditc_positive")
def _auditc_positive(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    # Sex-specific cut-off: >= 4 for men, >= 3 for women (lower tolerance / higher risk at equal intake).
    hist = _history(s, "auditc")
    threshold = p["threshold_m"] if s.sex == "m" else p["threshold_f"]
    if hist and hist[-1].total >= threshold:
        return {"total": hist[-1].total, "threshold": threshold}
    return None


@evaluator("quit_intent")
def _quit_intent(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    # Fires on the day the intent is voiced, and on a later skipped dose within the window
    # (the intent turned into action).
    anchor = _anchor_day(s)
    today = _checkins_by_day(s).get(anchor)
    if today is None:
        return None
    if today.quit_intent:
        return {}
    since = anchor - timedelta(days=p["window_days"])
    intent_before = any(c.quit_intent for c in s.checkins if since <= c.day < anchor)
    return {} if today.med_taken is False and intent_before else None


@evaluator("missed_doses_7d")
def _missed_doses(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    # Counts distinct calendar days (not doses) within the window with a missed/skipped dose,
    # from dose records and from the check-in answer "did not take" — whichever source has the data.
    window = p["window_days"]
    first = _today(s) - timedelta(days=window - 1)
    tz = s.as_of.tzinfo
    days = {
        d.planned_at.astimezone(tz).date()
        for d in s.doses
        if d.status in ("missed", "skipped") and d.planned_at <= s.as_of
    }
    days |= {c.day for c in s.checkins if c.med_taken is False}
    count = sum(1 for d in days if first <= d <= _today(s))
    return {"count": count} if count >= p["threshold"] else None


@evaluator("side_effect_3d")
def _side_effect_3d(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    run = _streak(s, p["consecutive_days"], lambda c: bool(c.side_effects))
    if run is None:
        return None
    effects = sorted({e for c in run for e in c.side_effects})
    return {"effects": ", ".join(effects)}


@evaluator("sleep_lt4_3n")
def _sleep_lt4(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    return None if _streak(s, p["consecutive_days"], lambda c: c.sleep_hours_bucket == "<4") is None else {}


@evaluator("alcohol_3d")
def _alcohol_3d(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    drank = _streak(s, p["consecutive_days"], lambda c: c.alcohol is not None and c.alcohol != "none")
    return None if drank is None else {}


@evaluator("mood_low_3d")
def _mood_low_3d(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    low = _streak(s, p["consecutive_days"], lambda c: c.mood is not None and c.mood <= p["threshold"])
    if low is None:
        return None
    # Do not ask for an unscheduled PHQ-9 again if one was already taken recently.
    since = s.as_of - timedelta(days=p["cooldown_days"])
    if any(not r.scheduled and r.ts > since for r in _history(s, "phq9")):
        return None
    return {}


@evaluator("no_contact_3d")
def _no_contact(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    if s.last_patient_message_at is None:
        return None
    silence = s.as_of - s.last_patient_message_at
    return {"silent_days": silence.days} if silence > timedelta(days=p["days"]) else None


@evaluator("young_adult_start")
def _young_adult_start(s: PatientSnapshot, p: dict[str, Any]) -> Context | None:
    if s.age < p["age_below"] and 0 <= s.therapy_day <= p["max_therapy_day"]:
        return {}
    return None


# ---------------------------------------------------------------------------------------------- engine


class _KeepMissing(dict[str, Any]):
    def __missing__(self, key: str) -> str:
        return "{" + key + "}"


def _build_actions(rule: RuleSpec, ctx: Context) -> list[Action]:
    reason = rule.reason_ru.format_map(_KeepMissing({**rule.params, **ctx}))
    return [
        Action(
            type=spec.type,
            level=_LEVELS[spec.level] if spec.level else None,
            code=rule.id,
            params={**ctx, **spec.params},
            reason_ru=reason,
        )
        for spec in rule.actions
    ]


def evaluate(snapshot: PatientSnapshot, config: RulesConfig | None = None) -> list[Action]:
    """Run all enabled rules (YAML order) on the snapshot. Pure and deterministic."""
    cfg = config or load_rules()
    out: list[Action] = []
    for rule in cfg.rules:
        if not rule.enabled:
            continue
        ctx = _REGISTRY[rule.id](snapshot, rule.params)
        if ctx is None:
            continue
        actions = _build_actions(rule, ctx)
        is_crisis = any(a.type == "crisis" for a in actions)
        already_open = any(a.type == "alert" and a.code in snapshot.open_alert_codes for a in actions)
        if already_open and not is_crisis:
            continue
        out.extend(actions)
    return out


# ---------------------------------------------------------------------------------------------- colours


class PatientStatus(BaseModel):
    color: StatusColor
    top_level: AlertLevel | None


def color(levels: Iterable[AlertLevel | None]) -> StatusColor:
    """Dashboard colour for a set of open alert levels (ADR-006)."""
    present = {lvl for lvl in levels if lvl is not None}
    if AlertLevel.ACUTE in present:
        return "red"
    if present & {AlertLevel.HIGH, AlertLevel.MODERATE}:
        return "yellow"
    if AlertLevel.INFO in present:
        return "grey"
    return "green"


def patient_status(open_alert_levels: Iterable[AlertLevel | None]) -> PatientStatus:
    """Colour + highest open level for a patient row on the dashboard."""
    levels = [lvl for lvl in open_alert_levels if lvl is not None]
    return PatientStatus(color=color(levels), top_level=max(levels, default=None))