Pavel Eliseev
c42ef621e5
CI / ruff · pytest (pull_request) Successful in 55s
Add gated TSQ, Russian C-SSRS Recent, durable responses and notifications, diaries, schedules and audited clinical review. Preserve manual assessment when the Russian follow-up source is unavailable. Refs #52
360 lines
17 KiB
Python
360 lines
17 KiB
Python
"""Dialog engine (issue #10, ADR-004): `Engine.handle(token, event) -> list[Reply]`, channel-agnostic.
|
||
|
||
Pipeline for one incoming event (SPEC §3, guardrails #1, #2, #3):
|
||
|
||
text event
|
||
1. scrub() raw text → ScrubbedText. Raw text is never stored, logged or sent anywhere.
|
||
2. consent gate no valid consent → no LLM, nothing stored; only the local crisis lexicon runs
|
||
(protecting life does not need consent) and the core flow shows consent/refusal.
|
||
3. safety.check() lexicon FIRST. Lexicon ACUTE → the LLM is not called at all. Otherwise the LLM risk
|
||
classifier is an auxiliary signal combined as max(lexicon, llm): it can only raise risk.
|
||
ACUTE → fixed crisis script verbatim (once per open crisis), state `crisis`, structured
|
||
`safety_events` fact for lane B, MonitoringService.after_data_change() (lane B alerts).
|
||
POSSIBLE → C-SSRS screen (scale session) and back to the previous state (or crisis if high).
|
||
4. FSM the current state's handler (`app/dialog/flows/*`, `app/dialog/scale_session.py`).
|
||
button / system / start events skip 1–3 (no free text) and go through the same FSM.
|
||
|
||
After the FSM step: if clinical data were written → commit → `after_data_change()` → the returned *dialog*
|
||
actions are executed (`app/dialog/actions.py`). The engine never creates alerts (ADR-007).
|
||
Then a queued scale starts if the FSM rests (idle; only the C-SSRS while in crisis).
|
||
Persist: inbound message (scrubbed) + outbound replies + `dialog_state` — only with active consent.
|
||
|
||
Logs carry only the token, states and event codes (guardrail #8).
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime
|
||
from typing import Any
|
||
from weakref import WeakValueDictionary
|
||
|
||
from sqlalchemy.exc import InterfaceError, OperationalError
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core import clock, safety
|
||
from app.core.intake_runtime import intake_enabled
|
||
from app.core.privacy import scrub
|
||
from app.core.safety import SafetyResult, crisis_reply
|
||
from app.core.types import RiskLevel, ScrubbedText
|
||
from app.db import repo
|
||
from app.db.session import get_clinical_session
|
||
from app.dialog import scale_session
|
||
from app.dialog.actions import execute
|
||
from app.dialog.contracts import Event, Reply, register_engine
|
||
from app.dialog.crisis import crisis_open, no_consent_crisis_text, open_crisis
|
||
from app.dialog.flows import load_flows
|
||
from app.dialog.flows.core import close_crisis, revoke
|
||
from app.dialog.fsm import NO_CONSENT_STATES, RESTING_STATES, Ctx, Step, apply_data, enter_hook_for, handler_for
|
||
from app.dialog.texts import DialogTexts, load_dialog_texts
|
||
from app.llm.gateway import LLMGateway, get_gateway
|
||
from app.services import questionnaire as intake_service
|
||
from app.services.contracts import MonitoringService, get_monitoring
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
# Enter hooks may chain (scale done → next onboarding state → next scale); a cycle must not hang a request.
|
||
MAX_ENTER_CHAIN = 8
|
||
_patient_locks = WeakValueDictionary()
|
||
# DB unreachable: answer safely instead of failing the channel (connectivity errors only — bugs must surface).
|
||
_DB_DOWN = (OSError, OperationalError, InterfaceError)
|
||
|
||
|
||
@dataclass
|
||
class _Run:
|
||
replies: list[Reply] = field(default_factory=list)
|
||
notify: bool = False
|
||
|
||
|
||
def _consented(ctx: Ctx) -> bool:
|
||
return ctx.data.get("consent") is True and ctx.state not in NO_CONSENT_STATES
|
||
|
||
|
||
class Engine:
|
||
"""Real `DialogEngine`. `gateway` / `monitoring` are injectable for tests; by default resolved per call so
|
||
a monitoring implementation registered later (lane B, #38) is picked up."""
|
||
|
||
def __init__(self, *, gateway: LLMGateway | None = None, monitoring: MonitoringService | None = None) -> None:
|
||
self._gateway = gateway
|
||
self._monitoring = monitoring
|
||
|
||
async def handle(self, token: str, event: Event) -> list[Reply]:
|
||
texts = load_dialog_texts()
|
||
try:
|
||
lock = _patient_locks.setdefault(token, asyncio.Lock())
|
||
async with lock:
|
||
async with get_clinical_session() as session:
|
||
await intake_service.lock_patient(session, token)
|
||
replies = await self._handle(token, event, session, texts)
|
||
await intake_service.dispatch_notices(token)
|
||
return replies
|
||
except _DB_DOWN as exc:
|
||
log.error("dialog db unavailable token=%s event=%s error=%s", token, event.type, type(exc).__name__)
|
||
return _degraded(event, texts)
|
||
|
||
# ------------------------------------------------------------------------------------------ pipeline
|
||
|
||
async def _handle(self, token: str, event: Event, session: AsyncSession, texts: DialogTexts) -> list[Reply]:
|
||
row = await repo.get_dialog_state(session, token)
|
||
ctx = Ctx(
|
||
token=token,
|
||
state=row.state if row is not None else "new",
|
||
data=dict(row.data or {}) if row is not None else {},
|
||
now=clock.now(),
|
||
session=session,
|
||
gateway=self._gateway or get_gateway(),
|
||
monitoring=self._monitoring or get_monitoring(),
|
||
texts=texts,
|
||
)
|
||
state_before = ctx.state
|
||
run = _Run()
|
||
|
||
await self._apply(ctx, await self._route(ctx, event), run)
|
||
|
||
if run.notify:
|
||
# Monitoring (lane B) reads the DB in its own session: commit our writes first.
|
||
await self._save_state(ctx)
|
||
await session.commit()
|
||
actions = await ctx.monitoring.after_data_change(token)
|
||
await self._apply(ctx, await execute(ctx, actions), run)
|
||
|
||
await self._apply(ctx, await self._start_pending(ctx), run)
|
||
|
||
if _consented(ctx):
|
||
for reply in run.replies:
|
||
# Outgoing texts may quote the safety plan (names of relatives) — scrub before storing too.
|
||
await repo.add_message(
|
||
session, token=token, direction="out", text=scrub(reply.text), ts=ctx.now, fsm_state=ctx.state
|
||
)
|
||
await self._save_state(ctx)
|
||
task_id = event.params.get("questionnaire_task")
|
||
if event.type == "system" and task_id and run.replies:
|
||
from app.db.intake import QuestionnaireTask
|
||
|
||
task = await session.get(QuestionnaireTask, task_id)
|
||
if task and task.token == token:
|
||
task.delivery = [r.model_dump() for r in run.replies]
|
||
log.info(
|
||
"dialog token=%s event=%s code=%s state=%s->%s replies=%d",
|
||
token,
|
||
event.type,
|
||
_event_code(event),
|
||
state_before,
|
||
ctx.state,
|
||
len(run.replies),
|
||
)
|
||
return run.replies
|
||
|
||
async def _route(self, ctx: Ctx, event: Event) -> Step:
|
||
consented = _consented(ctx)
|
||
if event.type == "text":
|
||
# (1) scrub before anything else: from here on only ScrubbedText exists for this message.
|
||
ctx.text = scrub(event.text or "")
|
||
if consented and _is_revoke_command(event.text or "", ctx.texts):
|
||
return await revoke(ctx)
|
||
if not consented:
|
||
return await self._no_consent_text(ctx, event)
|
||
await self._store_in(ctx, ctx.text)
|
||
return await self._safety(ctx, event)
|
||
|
||
if event.type == "button":
|
||
if consented and event.button == "consent:revoke":
|
||
return await revoke(ctx)
|
||
if consented and event.button == "intake:help":
|
||
from app.db.intake import QuestionnaireRecord
|
||
from app.dialog.flows.intake import urgent
|
||
|
||
row = await ctx.session.get(QuestionnaireRecord, ctx.data.get("questionnaire_record", ""))
|
||
return await urgent(ctx, row)
|
||
if consented:
|
||
await self._store_in(ctx, scrub(f"[button] {event.button}")) # button data comes from the client
|
||
return await self._dispatch(ctx, event)
|
||
|
||
if event.type == "system":
|
||
if event.system == "doctor_ack":
|
||
return close_crisis(ctx)
|
||
if not consented:
|
||
return Step() # nothing is pushed to people without consent
|
||
if event.system == "scale_due" and event.params.get("scale"):
|
||
if intake_enabled() and ctx.data.get("adult") is True:
|
||
scale = str(event.params["scale"])
|
||
kind = "cssrs_since_contact" if scale in ("cssrs_screen", "cssrs") else scale
|
||
if kind in ("phq9", "pcl5", "rus_audit_s", "cssrs_since_contact"):
|
||
await intake_service.queue_task(
|
||
ctx.session, ctx.token, kind, f"external:{kind}:{ctx.now.date()}", ctx.now
|
||
)
|
||
return Step()
|
||
scheduled = bool(event.params.get("scheduled", True))
|
||
pending = scale_session.add_pending(ctx.data, str(event.params["scale"]), scheduled)
|
||
return Step(data={"pending_scales": pending}) # starts now if resting, else after the current step
|
||
return await self._dispatch(ctx, event)
|
||
|
||
return await self._dispatch(ctx, event) # start
|
||
|
||
async def _no_consent_text(self, ctx: Ctx, event: Event) -> Step:
|
||
assert ctx.text is not None
|
||
# Lexicon only: without consent the text must not reach the LLM (third-party processing).
|
||
result = safety.check(ctx.text)
|
||
if result.level == RiskLevel.ACUTE:
|
||
log.warning("safety acute without consent token=%s rules=%s", ctx.token, result.matched_rules)
|
||
return Step(replies=[Reply(text=no_consent_crisis_text(crisis_reply(result.categories)), kind="crisis")])
|
||
return await self._dispatch(ctx, event)
|
||
|
||
async def _safety(self, ctx: Ctx, event: Event) -> Step:
|
||
"""Guardrail #2: the safety decision is made before any LLM sees the message for anything else."""
|
||
assert ctx.text is not None
|
||
lexicon = safety.check(ctx.text)
|
||
if lexicon.level == RiskLevel.ACUTE:
|
||
result = lexicon # an acute message never reaches the LLM — not even the risk classifier
|
||
else:
|
||
result = safety.check(ctx.text, llm_signal=await ctx.gateway.risk_classify(ctx.text))
|
||
|
||
if result.level == RiskLevel.ACUTE:
|
||
handled = "crisis_repeat" if crisis_open(ctx.data) else "crisis_script"
|
||
log.warning("safety acute token=%s rules=%s handled=%s", ctx.token, result.matched_rules, handled)
|
||
if ctx.data.get("questionnaire_record") or (intake_enabled() and ctx.data.get("adult") is True):
|
||
from app.db.intake import QuestionnaireRecord
|
||
from app.dialog.flows.intake import urgent
|
||
|
||
row = await ctx.session.get(QuestionnaireRecord, ctx.data.get("questionnaire_record", ""))
|
||
step = await urgent(ctx, row, code="safety_acute", categories=result.categories)
|
||
else:
|
||
step = open_crisis(ctx, result.categories, source=f"safety:{result.source}")
|
||
step.data["safety_events"] = self._safety_events(ctx, result, handled)
|
||
step.notify_monitoring = True # lane B escalates to the doctor on call
|
||
return step
|
||
|
||
if result.level == RiskLevel.POSSIBLE and (
|
||
ctx.data.get("questionnaire_record") or (intake_enabled() and ctx.data.get("adult") is True)
|
||
):
|
||
from app.db.intake import QuestionnaireRecord
|
||
from app.dialog.flows.intake import manual_review
|
||
|
||
row = await ctx.session.get(QuestionnaireRecord, ctx.data.get("questionnaire_record", ""))
|
||
if row is None:
|
||
row = await intake_service.create_record(ctx, "cssrs_since_contact")
|
||
return await manual_review(ctx, row, "cssrs_since_contact_required")
|
||
|
||
if result.level == RiskLevel.POSSIBLE:
|
||
screen = ctx.texts.settings.risk_screen_scale
|
||
if self._can_screen(ctx, screen):
|
||
step = await scale_session.start_scale(
|
||
ctx, screen, return_state=ctx.state, scheduled=False, intro=ctx.texts.scale.risk_screen_intro
|
||
)
|
||
step.data["risk_screen_at"] = ctx.now.isoformat()
|
||
handled = "risk_screen"
|
||
else:
|
||
step = await self._dispatch(ctx, event)
|
||
handled = "dispatched"
|
||
log.info("safety possible token=%s rules=%s handled=%s", ctx.token, result.matched_rules, handled)
|
||
step.data["safety_events"] = self._safety_events(ctx, result, handled)
|
||
step.notify_monitoring = True
|
||
return step
|
||
|
||
return await self._dispatch(ctx, event)
|
||
|
||
def _can_screen(self, ctx: Ctx, screen: str) -> bool:
|
||
if crisis_open(ctx.data) or scale_session.active_scale(ctx.state) == screen:
|
||
return False
|
||
last = ctx.data.get("risk_screen_at")
|
||
if last is None:
|
||
return True
|
||
hours = (ctx.now - datetime.fromisoformat(last)).total_seconds() / 3600
|
||
return hours >= ctx.texts.settings.cssrs_cooldown_hours
|
||
|
||
def _safety_events(self, ctx: Ctx, result: SafetyResult, handled: str) -> list[dict[str, Any]]:
|
||
"""Structured safety fact for lane B (#38): no text, only codes. Newest last, capped."""
|
||
event = {
|
||
"ts": ctx.now.isoformat(),
|
||
"level": result.level.name.lower(),
|
||
"categories": result.categories,
|
||
"rules": result.matched_rules,
|
||
"source": result.source,
|
||
"handled": handled,
|
||
"state": ctx.state,
|
||
}
|
||
events = [*(ctx.data.get("safety_events") or []), event]
|
||
return events[-ctx.texts.settings.max_safety_events :]
|
||
|
||
async def _dispatch(self, ctx: Ctx, event: Event) -> Step:
|
||
handler = handler_for(ctx.state)
|
||
if handler is None:
|
||
# A state from a removed/renamed flow: recover instead of getting stuck.
|
||
log.error("dialog unknown state token=%s state=%s", ctx.token, ctx.state)
|
||
return Step(next_state="idle" if _consented(ctx) else "new")
|
||
return await handler(ctx, event)
|
||
|
||
async def _apply(self, ctx: Ctx, step: Step, run: _Run, depth: int = 0) -> None:
|
||
run.replies.extend(step.replies)
|
||
apply_data(ctx.data, step.data)
|
||
run.notify = run.notify or step.notify_monitoring
|
||
if not step.next_state or step.next_state == ctx.state:
|
||
return
|
||
ctx.state = step.next_state
|
||
hook = enter_hook_for(ctx.state)
|
||
if hook is None:
|
||
return
|
||
if depth >= MAX_ENTER_CHAIN:
|
||
log.error("dialog enter chain too long token=%s state=%s", ctx.token, ctx.state)
|
||
return
|
||
await self._apply(ctx, await hook(ctx), run, depth + 1)
|
||
|
||
async def _start_pending(self, ctx: Ctx) -> Step:
|
||
"""Start a queued scale when nothing else is being asked. In crisis only the risk screen may run."""
|
||
if intake_enabled():
|
||
return Step()
|
||
pending = list(ctx.data.get("pending_scales") or [])
|
||
if not pending or ctx.state not in RESTING_STATES:
|
||
return Step()
|
||
screen = ctx.texts.settings.risk_screen_scale
|
||
if ctx.state == "crisis":
|
||
pending_screen = [p for p in pending if p["scale"] == screen]
|
||
if not pending_screen:
|
||
return Step()
|
||
item = pending_screen[0]
|
||
else:
|
||
item = pending[0]
|
||
rest = [p for p in pending if p is not item]
|
||
step = await scale_session.start_scale(
|
||
ctx,
|
||
item["scale"],
|
||
return_state=ctx.state,
|
||
scheduled=bool(item.get("scheduled", True)),
|
||
intro=ctx.texts.scale.risk_screen_intro if item["scale"] == screen else None,
|
||
)
|
||
step.data["pending_scales"] = rest or None
|
||
return step
|
||
|
||
async def _store_in(self, ctx: Ctx, text: ScrubbedText) -> None:
|
||
await repo.add_message(ctx.session, token=ctx.token, direction="in", text=text, ts=ctx.now, fsm_state=ctx.state)
|
||
|
||
async def _save_state(self, ctx: Ctx) -> None:
|
||
await repo.set_dialog_state(ctx.session, token=ctx.token, state=ctx.state, data=ctx.data, now=ctx.now)
|
||
|
||
|
||
def _is_revoke_command(raw: str, texts: DialogTexts) -> bool:
|
||
return raw.strip().lower().rstrip(".!") in texts.commands.revoke
|
||
|
||
|
||
def _event_code(event: Event) -> str:
|
||
if event.type == "button":
|
||
return event.button or ""
|
||
if event.type == "system":
|
||
return event.system or ""
|
||
return "-"
|
||
|
||
|
||
def _degraded(event: Event, texts: DialogTexts) -> list[Reply]:
|
||
"""DB down: still answer a crisis message (lexicon only, no "doctor notified" — nobody could be notified)."""
|
||
if event.type == "text" and event.text:
|
||
result = safety.check(scrub(event.text))
|
||
if result.level == RiskLevel.ACUTE:
|
||
return [Reply(text=no_consent_crisis_text(crisis_reply(result.categories)), kind="crisis")]
|
||
if event.type == "system":
|
||
return []
|
||
return [Reply(text=texts.unavailable, kind="info")]
|
||
|
||
|
||
load_flows()
|
||
register_engine(Engine)
|