"""FSM framework of the dialog engine: context, step result and the state / trigger registries. States are plain strings, namespaced by module ("veteran.onboarding.pcl5"); core states have short names ("new", "consent", "age_gate", "idle", "crisis", "consent_revoked", ...). A flow module (`app/dialog/flows/.py`) registers its states at import time — the engine never has to be edited to add a scenario: from app.dialog.fsm import Ctx, Step, on_enter, state from app.dialog.scale_session import start_scale @on_enter("veteran.onboarding.start") # runs when the FSM *enters* the state async def _start(ctx: Ctx) -> Step: return await start_scale(ctx, "pcptsd5", return_state="veteran.onboarding.pcl5") @state("veteran.onboarding.trauma") # handles events while the FSM *is in* the state async def _trauma(ctx: Ctx, event: Event) -> Step: ... return Step(replies=[...], next_state="idle", data={"trauma": True}) Name resolution: an exact registration wins, otherwise the longest registered dotted prefix handles the state ("scale.phq9" → the generic "scale" handler). Triggers (`@trigger("button:menu:checkin")`, `@trigger("system:daily_checkin")`, `@trigger("text:idle")`) let flows take over what the core `idle` state does with menu buttons, system events and free text; the core registers defaults, a flow overrides with `replace=True`. Prefix triggers carry data in the button (#56): `@trigger("button:med:taken:*")` handles every button whose data starts with "med:taken:"; the rest of the data ("med:taken:42" → "42") reaches the handler as `trigger_payload(event)`. An exact key always wins over a prefix, a longer prefix over a shorter one. Abandoned steps (#56): a flow marks its question states with `abandonable("daily")`; if the patient stays silent in such a state longer than `settings.abandon_after_hours`, the engine returns the FSM to `idle` before handling the next event (answers already stored stay stored). """ from __future__ import annotations from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from datetime import datetime from typing import Any from sqlalchemy.ext.asyncio import AsyncSession from app.core.types import ScrubbedText from app.db import repo from app.db.clinical import PatientProfile from app.dialog.contracts import Event, Reply from app.dialog.texts import DialogTexts from app.llm.gateway import LLMGateway from app.services.contracts import MonitoringService # States where nothing is being asked: a pending scale may start here (see engine `_start_pending`). RESTING_STATES = frozenset({"idle", "crisis"}) # States without (valid) consent: the engine stores nothing and never calls the LLM here (152-FZ art. 9-10). NO_CONSENT_STATES = frozenset({"new", "consent", "refused", "minor", "consent_revoked"}) @dataclass class Ctx: """Everything a state handler may use. `data` is the persisted `dialog_state.data` (treat as read-only: return changes in `Step.data`, the engine applies them).""" token: str state: str data: dict[str, Any] now: datetime # from app.core.clock, taken once per incoming event (ADR-005) session: AsyncSession # clinical DB unit of work of this event gateway: LLMGateway monitoring: MonitoringService texts: DialogTexts text: ScrubbedText | None = None # scrubbed patient text for text events — the only text handlers see _profile: PatientProfile | None = field(default=None, repr=False) _profile_loaded: bool = field(default=False, repr=False) async def profile(self) -> PatientProfile | None: if not self._profile_loaded: self._profile = await repo.get_profile(self.session, self.token) self._profile_loaded = True return self._profile @dataclass class Step: """Result of a handler. replies — messages to send, in order; next_state — new FSM state (None = stay); entering a different state runs its `@on_enter` hook; data — updates merged into `dialog_state.data` (a value of None deletes the key); notify_monitoring — the handler wrote clinical data: the engine commits and calls `MonitoringService.after_data_change()`, then executes the returned dialog actions. """ replies: list[Reply] = field(default_factory=list) next_state: str | None = None data: dict[str, Any] = field(default_factory=dict) notify_monitoring: bool = False def then(self, other: Step) -> Step: """Sequential composition: replies concatenated, later state/data win.""" return Step( replies=[*self.replies, *other.replies], next_state=other.next_state or self.next_state, data={**self.data, **other.data}, notify_monitoring=self.notify_monitoring or other.notify_monitoring, ) Handler = Callable[[Ctx, Event], Awaitable[Step]] EnterHook = Callable[[Ctx], Awaitable[Step]] _handlers: dict[str, Handler] = {} _enter_hooks: dict[str, EnterHook] = {} _triggers: dict[str, Handler] = {} _prefix_triggers: dict[str, Handler] = {} # "button:med:taken:" → handler (registered as "button:med:taken:*") _abandonable: dict[str, Callable[[Ctx], bool] | None] = {} PREFIX_WILDCARD = "*" PAYLOAD_PARAM = "trigger_payload" def state(name: str) -> Callable[[Handler], Handler]: """Register the event handler of state `name` (or of a dotted prefix, e.g. "scale").""" def deco(fn: Handler) -> Handler: if name in _handlers: raise ValueError(f"state '{name}' already registered") _handlers[name] = fn return fn return deco def on_enter(name: str) -> Callable[[EnterHook], EnterHook]: """Register a hook that runs when the FSM enters state `name` (asks the first question, chains flows).""" def deco(fn: EnterHook) -> EnterHook: if name in _enter_hooks: raise ValueError(f"enter hook of '{name}' already registered") _enter_hooks[name] = fn return fn return deco def trigger(key: str, *, replace: bool = False) -> Callable[[Handler], Handler]: """Register a handler for an idle-menu button / system event / idle free text. Keys: "button:", "system:", "text:idle". The core registers defaults; flows pass `replace=True` to take one over. A key ending in ":*" is a prefix trigger (see module doc). """ table, name = _triggers, key if key.endswith(":" + PREFIX_WILDCARD): table, name = _prefix_triggers, key[: -len(PREFIX_WILDCARD)] def deco(fn: Handler) -> Handler: if name in table and not replace: raise ValueError(f"trigger '{key}' already registered (use replace=True to override)") table[name] = fn return fn return deco def trigger_payload(event: Event) -> str | None: """Data after the matched prefix of a prefix trigger ("med:taken:42" → "42"); None for exact triggers.""" value = event.params.get(PAYLOAD_PARAM) return str(value) if value is not None else None def abandonable(prefix: str, when: Callable[[Ctx], bool] | None = None) -> None: """Mark states `prefix` / `prefix.*` as question steps the engine may abandon after a long silence. `when(ctx)` may veto it for a particular situation (e.g. a scale that is part of an onboarding). """ if prefix in _abandonable: raise ValueError(f"abandonable '{prefix}' already registered") _abandonable[prefix] = when def is_abandonable(ctx: Ctx) -> bool: key = _resolve(ctx.state, _abandonable) if key is None: return False when = _abandonable[key] return when is None or when(ctx) def _resolve(name: str, table: dict[str, Any]) -> str | None: candidate = name while candidate: if candidate in table: return candidate if "." not in candidate: return None candidate = candidate.rsplit(".", 1)[0] return None def handler_for(name: str) -> Handler | None: key = _resolve(name, _handlers) return _handlers[key] if key else None def enter_hook_for(name: str) -> EnterHook | None: key = _resolve(name, _enter_hooks) return _enter_hooks[key] if key else None def get_trigger(key: str) -> Handler | None: """Exact trigger, else the longest matching prefix trigger (wrapped so it sees the payload).""" exact = _triggers.get(key) if exact is not None: return exact matches = [p for p in _prefix_triggers if key.startswith(p) and len(key) > len(p)] if not matches: return None prefix = max(matches, key=len) fn, payload = _prefix_triggers[prefix], key[len(prefix) :] async def with_payload(ctx: Ctx, event: Event) -> Step: return await fn(ctx, event.model_copy(update={"params": {**event.params, PAYLOAD_PARAM: payload}})) return with_payload def is_registered(name: str) -> bool: """Exact registration (handler or enter hook) — used to find a module's onboarding entry point.""" return name in _handlers or name in _enter_hooks def apply_data(data: dict[str, Any], updates: dict[str, Any]) -> None: for key, value in updates.items(): if value is None: data.pop(key, None) else: data[key] = value