PTSRaiAgent/tests/test_dialog_knowledge.py

"""Knowledge base in the dialog (#64): «Вопрос о лечении» menu, questions in idle free text, routing order
(safety → check-in signals → drug question → guidelines → support)."""

from collections.abc import AsyncIterator
from datetime import datetime
from typing import Any

import pytest

from app.config import Settings
from app.core import clock
from app.db import repo, session
from app.db.bootstrap import init_db
from app.db.clinical import Prescription
from app.dialog.contracts import Event, Reply
from app.dialog.engine import Engine
from app.dialog.flows.knowledge import QUESTION_STATE, load_knowledge_texts, looks_like_question
from app.integrations.base import DrugSection, integrations
from app.llm import med_qa as med_qa_module
from app.llm.gateway import LLMGateway
from app.llm.med_qa import load_med_qa_texts
from app.llm.mock import MockTasks
from app.llm.rag import fallback_text
from app.llm.rag import qa as rag_qa
from app.llm.texts import render_phrase
from app.services.contracts import NoopMonitoring

TOKEN = "k" * 32
MEM = "sqlite+aiosqlite:///:memory:"
NOW = datetime(2026, 10, 7, 8, 0)
SUPPORT = render_phrase("support", {})
KB = load_knowledge_texts()


class FakeDrugInfo:
    async def sections(self, inn: str) -> list[DrugSection]:
        text = "Улучшение обычно отмечается через 2–4 недели регулярного приёма."
        return [DrugSection(inn=inn, section="onset", text=text, source="синтетика")]


@pytest.fixture
async def engine(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[Engine]:
    monkeypatch.setattr(clock, "get_settings", lambda: Settings(_env_file=None, demo_clock=True))
    monkeypatch.setattr(integrations, "drug_info", FakeDrugInfo())
    monkeypatch.setattr(med_qa_module, "get_gateway", lambda: LLMGateway(MockTasks()))
    monkeypatch.setattr(rag_qa, "get_gateway", lambda: LLMGateway(MockTasks()))
    clock.set_now(NOW)
    await session.configure(identity_url=MEM, clinical_url=MEM)
    await init_db()
    async with session.get_clinical_session() as s:
        await repo.upsert_profile(
            s, token=TOKEN, doctor_id="doc", modules=["veteran"], sex="m", birth_year=1990, therapy_start=NOW.date()
        )
        s.add(Prescription(token=TOKEN, drug="сертралин", dose="50 мг", schedule={"times": ["08:00"]}))
        await repo.set_dialog_state(s, token=TOKEN, state="idle", data={"consent": True}, now=clock.now())
    yield Engine(gateway=LLMGateway(MockTasks()), monitoring=NoopMonitoring())
    clock.reset()
    await session.configure()


async def press(e: Engine, data: str) -> list[Reply]:
    return await e.handle(TOKEN, Event(type="button", button=data))


async def say(e: Engine, t: str) -> list[Reply]:
    return await e.handle(TOKEN, Event(type="text", text=t))


async def state() -> tuple[str, dict[str, Any]]:
    async with session.get_clinical_session() as s:
        row = await repo.get_dialog_state(s, TOKEN)
        assert row is not None
        return row.state, dict(row.data)


def buttons(replies: list[Reply]) -> list[str]:
    return [b.data for r in replies for row in r.buttons for b in row]


def joined(replies: list[Reply]) -> str:
    return "\n".join(r.text for r in replies)


# --- question detection --------------------------------------------------------------------------------


@pytest.mark.parametrize(
    ("text", "expected"),
    [
        ("что такое флэшбек", True),
        ("Можно ли вылечить ПТСР", True),
        ("кошмары снятся, это нормально?", True),
        ("сегодня 6 из 10", False),
        ("спал плохо", False),
        ("какой-то день", False),  # "как" only as a whole word at the start
        ("как дела?", False),  # small talk
        ("Привет! Как ты?", False),
        ("[NAME], что такое ПТСР", True),  # scrub placeholder before the question
    ],
)
def test_looks_like_question(text: str, expected: bool) -> None:
    assert looks_like_question(text) is expected


# --- menu «Вопрос о лечении» ------------------------------------------------------------------------------


async def test_menu_has_knowledge_button(engine: Engine) -> None:
    r = await engine.handle(TOKEN, Event(type="start"))
    assert "menu:kb_question" in buttons(r)


async def test_menu_question_answers_from_guidelines(engine: Engine) -> None:
    r = await press(engine, "menu:kb_question")
    assert joined(r) == KB.prompt
    assert (await state())[0] == QUESTION_STATE

    r = await say(engine, "что такое флэшбек?")
    text = joined(r)
    assert "Флэшбек (flashback" in text
    assert "Источник: КР ПТСР 2023, стр. 2" in text
    assert rag_qa.load_rag_texts()["patient"].doctor_note in text
    assert "menu:kb_question" in buttons(r)  # back to the menu
    assert (await state())[0] == "idle"


async def test_menu_question_without_answer_sends_to_doctor(engine: Engine) -> None:
    await press(engine, "menu:kb_question")
    r = await say(engine, "какой рецепт борща")
    assert fallback_text("patient") in joined(r)
    assert (await state())[0] == "idle"


@pytest.mark.parametrize("question", ["можно увеличить дозу?", "можно пить полтаблетки?"])
async def test_dose_question_redirected_to_doctor(engine: Engine, question: str) -> None:
    await press(engine, "menu:kb_question")
    r = await say(engine, question)
    assert joined(r).startswith(KB.dose_redirect)
    assert "клинических рекомендациях" not in joined(r)


async def test_button_instead_of_question_leaves_the_flow(engine: Engine) -> None:
    await press(engine, "menu:kb_question")
    r = await press(engine, "menu:contact_doctor")
    assert "врач" in joined(r)
    assert (await state())[0] == "idle"


async def test_scrubbed_question_still_answered(engine: Engine) -> None:
    await press(engine, "menu:kb_question")
    r = await say(engine, "Иван Петров, 89161234567, что такое ПТСР?")
    assert "ПТСР – это расстройство" in joined(r)


# --- free text in idle ------------------------------------------------------------------------------------


async def test_free_text_question_answered_from_guidelines(engine: Engine) -> None:
    r = await say(engine, "помогает ли психотерапия при ПТСР?")
    text = joined(r)
    assert "В клинических рекомендациях Минздрава сказано:" in text
    assert "Источник: КР ПТСР 2023" in text


@pytest.mark.parametrize("message", ["сегодня обычный день", "как дела?"])
async def test_free_text_without_answer_keeps_support(engine: Engine, message: str) -> None:
    r = await say(engine, message)
    assert joined(r) == SUPPORT
    assert "Источник" not in joined(r)


async def test_drug_question_stays_with_instruction(engine: Engine) -> None:
    # SPEC §7 п.5: the patient's own drug is answered only from its instruction, never from the guidelines.
    r = await say(engine, "через сколько подействуют таблетки?")
    text = joined(r)
    assert text.startswith(load_med_qa_texts().intro)
    assert "клинических рекомендациях" not in text


async def test_side_effect_question_goes_to_checkin(engine: Engine) -> None:
    r = await say(engine, "меня тошнит от таблеток, это нормально?")
    assert "клинических рекомендациях" not in joined(r)


async def test_acute_risk_never_reaches_knowledge_base(engine: Engine) -> None:
    r = await say(engine, "зачем жить? хочу умереть")
    assert any(x.kind == "crisis" for x in r)
    assert "клинических рекомендациях" not in joined(r)