PTSRaiAgent/tests/test_llm.py

"""LLM gateway: mock coverage of demo phrases, GigaChat parsing/retries (fake client), policy guards."""

import json
from types import SimpleNamespace
from typing import Any

import pytest

from app.config import Settings
from app.core.types import RiskLevel, ScrubbedText
from app.llm import gateway as gw_module
from app.llm.gateway import ITEM_CONFIDENCE_THRESHOLD, LLMGateway, build_gateway, is_forbidden
from app.llm.giga import MAX_RETRIES, GigaChatProvider, parse_json_object
from app.llm.guard import UnscrubbedTextError, assert_scrubbed
from app.llm.mock import MockTasks
from app.llm.schemas import CheckinFacts, ItemScore, RiskVerdict
from app.llm.tasks import PromptedTasks
from app.llm.texts import load_texts, render_phrase, render_summary

S = ScrubbedText


@pytest.fixture
def gw() -> LLMGateway:
    return LLMGateway(MockTasks())


# --- mock: demo persona replies ---------------------------------------------------------------------


async def test_nightmares_item_and_checkin(gw: LLMGateway) -> None:
    text = S("сплю плохо, кошмары почти каждую ночь")
    phq = await gw.extract_item("phq9_3", "Проблемы со сном", 3, text)
    assert phq is not None and phq.score == 3 and phq.item == "phq9_3"
    pcl = await gw.extract_item("pcl5_2", "Кошмары", 4, text)
    assert pcl is not None and pcl.score == 4
    facts = await gw.extract_checkin(text)
    assert facts is not None and facts.nightmares is True and facts.confidence >= 0.5
    assert facts.side_effects == []


async def test_nausea(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("тошнит"))
    assert facts is not None and facts.side_effects == ["nausea"]
    assert facts.quit_intent is False


async def test_quit_intent(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("от таблеток только хуже, брошу"))
    assert facts is not None
    assert facts.quit_intent is True
    assert facts.side_effects == ["other"]


async def test_quit_drinking_is_not_quit_intent(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("брошу пить пиво"))
    assert facts is not None and facts.quit_intent is False and facts.alcohol == "1-2"


async def test_alcohol_and_insomnia(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("не могу уснуть без пива"))
    assert facts is not None
    assert facts.alcohol == "1-2"
    assert "insomnia" in facts.side_effects


@pytest.mark.parametrize(
    ("text", "bucket"),
    [
        ("выпил пару бутылок", "1-2"),
        ("вчера выпил бокал вина", "1-2"),
        ("выпил пять бутылок пива", "more"),
        ("напился вчера", "more"),
        ("не пил, всё нормально", "none"),
        ("не могу уснуть без алкоголя", "1-2"),
    ],
)
async def test_alcohol_buckets(gw: LLMGateway, text: str, bucket: str) -> None:
    facts = await gw.extract_checkin(S(text))
    assert facts is not None and facts.alcohol == bucket


@pytest.mark.parametrize(
    ("text", "bucket"),
    [("спал 3 часа", "<4"), ("сплю часов 5", None), ("сплю 5 часов", "4-6"), ("почти не спал", "<4")],
)
async def test_sleep_bucket(gw: LLMGateway, text: str, bucket: str | None) -> None:
    facts = await gw.extract_checkin(S(text))
    assert facts is not None and facts.sleep_hours_bucket == bucket


async def test_mood_med_and_question(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("настроение 4, таблетку принял. Можно ли пить сертралин вечером?"))
    assert facts is not None
    assert facts.mood == 4 and facts.med_taken is True and facts.med_question is True
    missed = await gw.extract_checkin(S("забыл принять таблетку"))
    assert missed is not None and missed.med_taken is False


async def test_registry_question_is_not_med_question(gw: LLMGateway) -> None:
    facts = await gw.extract_checkin(S("А меня поставят на учёт?"))
    assert facts is not None and facts.med_question is False and facts.confidence < 0.5


@pytest.mark.parametrize(
    ("text", "level"),
    [
        ("всё бессмысленно, устал от всего", RiskLevel.POSSIBLE),
        ("не хочу жить", RiskLevel.ACUTE),
        ("тошнит", RiskLevel.NONE),
        ("от таблеток только хуже, брошу", RiskLevel.NONE),
        ("не могу уснуть без пива", RiskLevel.NONE),
    ],
)
async def test_risk(gw: LLMGateway, text: str, level: RiskLevel) -> None:
    assert await gw.risk_classify(S(text)) == level


# --- extract_item policy -----------------------------------------------------------------------------


async def test_item_low_confidence_returns_none(gw: LLMGateway) -> None:
    assert await gw.extract_item("phq9_1", "Интерес к делам", 3, S("ну такое")) is None
    assert await gw.extract_item("phq9_1", "Интерес к делам", 3, S("не знаю")) is None


async def test_item_cues(gw: LLMGateway) -> None:
    for text, score in [("иногда", 1), ("часто", 2), ("нет, не было", 0), ("2", 2)]:
        res = await gw.extract_item("phq9_1", "Интерес к делам", 3, S(text))
        assert res is not None and res.score == score, text


class _FixedTasks(MockTasks):
    def __init__(self, item: ItemScore | None = None, reply: str | None = None) -> None:
        self._item, self._reply = item, reply
        self.calls = 0

    async def extract_item(self, *args: Any) -> ItemScore | None:
        self.calls += 1
        return self._item

    async def phrase(self, intent: str, payload: ScrubbedText, base_text: str) -> str | None:
        self.calls += 1
        return self._reply

    async def summarize_for_doctor(self, structured: ScrubbedText, base_text: str) -> str | None:
        return self._reply

    async def risk_classify(self, text: ScrubbedText) -> RiskVerdict | None:
        raise RuntimeError("boom")


async def test_confidence_threshold_constant() -> None:
    below = ItemScore(item="x", score=2, confidence=ITEM_CONFIDENCE_THRESHOLD - 0.01, evidence="")
    at = below.model_copy(update={"confidence": ITEM_CONFIDENCE_THRESHOLD})
    assert await LLMGateway(_FixedTasks(item=below)).extract_item("x", "q", 3, S("a")) is None
    assert await LLMGateway(_FixedTasks(item=at)).extract_item("x", "q", 3, S("a")) == at


async def test_score_above_item_max_rejected() -> None:
    item = ItemScore(item="x", score=4, confidence=0.9, evidence="")
    assert await LLMGateway(_FixedTasks(item=item)).extract_item("x", "q", 3, S("a")) is None


def test_item_score_validation_context() -> None:
    with pytest.raises(ValueError):
        ItemScore.model_validate({"item": "x", "score": 4, "confidence": 1}, context={"max_score": 3})
    assert ItemScore.model_validate({"item": "x", "score": 4, "confidence": 1}).score == 4


async def test_risk_failure_is_none() -> None:
    assert await LLMGateway(_FixedTasks()).risk_classify(S("что угодно")) == RiskLevel.NONE


# --- scrub guard -------------------------------------------------------------------------------------


LEAKS = ["позвоните +7 (916) 123-45-67", "пишите a.b@mail.ru", "СНИЛС 123-456-789 01"]


@pytest.mark.parametrize("leak", LEAKS)
async def test_unscrubbed_text_never_reaches_provider(leak: str) -> None:
    with pytest.raises(UnscrubbedTextError):
        assert_scrubbed(leak)
    tasks = _FixedTasks(item=ItemScore(item="x", score=1, confidence=0.9, evidence=""))
    assert await LLMGateway(tasks).extract_item("x", "q", 3, S(leak)) is None
    assert tasks.calls == 0


def test_guard_allows_dates_and_numbers() -> None:
    text = "2026-10-03, PHQ-9 = 14, сертралин 50 мг, [PHONE]"
    assert assert_scrubbed(text) == text


# --- phrase / summary --------------------------------------------------------------------------------


def test_templates_pass_forbidden_filter() -> None:
    for intent, text in load_texts().phrases.items():
        assert not is_forbidden(text), intent


@pytest.mark.parametrize(
    "bad",
    [
        "У вас депрессия, это видно.",
        "Ваш диагноз — ПТСР.",
        "Увеличьте дозу до 100 мг.",
        "Отмените сертралин.",
        "Прекратите приём на пару дней.",
        "Не волнуйтесь, никто не узнает.",
    ],
)
async def test_forbidden_reply_falls_back(bad: str) -> None:
    assert is_forbidden(bad)
    reply = await LLMGateway(_FixedTasks(reply=bad)).phrase("quit_intent")
    assert reply == render_phrase("quit_intent", {})


async def test_phrase_ok_and_failure_fallback() -> None:
    ok = "Спасибо, что написали."
    assert await LLMGateway(_FixedTasks(reply=ok)).phrase("support") == ok
    assert await LLMGateway(_FixedTasks(reply=None)).phrase("support") == render_phrase("support", {})
    assert await LLMGateway(_FixedTasks(reply="x" * 1000)).phrase("support") == render_phrase("support", {})


async def test_phrase_mock_renders_template(gw: LLMGateway) -> None:
    assert "сертралин" in await gw.phrase("med_reminder", {"drug": "сертралин"})
    assert await gw.phrase("unknown_intent") == render_phrase("default", {})
    assert "{" not in await gw.phrase("med_reminder")


async def test_summary_mock_template(gw: LLMGateway) -> None:
    data = {"phq9": [14, 13, 18], "side_effects": ["nausea"], "quit_intent": True, "alcohol": None}
    text = await gw.summarize_for_doctor(data)
    assert text == render_summary(data)
    assert "PHQ-9: 14 → 13 → 18" in text and "Алкоголь" not in text
    assert await gw.summarize_for_doctor({}) == load_texts().summary.empty


# --- GigaChat provider with a fake client (no network) -----------------------------------------------


class FakeClient:
    def __init__(self, replies: list[Any]) -> None:
        self.replies = replies
        self.payloads: list[Any] = []

    async def achat(self, payload: Any) -> Any:
        self.payloads.append(payload)
        reply = self.replies.pop(0)
        if isinstance(reply, Exception):
            raise reply
        return SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content=reply))],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )


def _provider(replies: list[Any]) -> tuple[GigaChatProvider, FakeClient]:
    client = FakeClient(replies)
    return GigaChatProvider(Settings(_env_file=None), client=client), client


def test_parse_json_object() -> None:
    assert parse_json_object('```json\n{"a": 1}\n```') == {"a": 1}
    assert parse_json_object('Вот ответ: {"a": 1} готово') == {"a": 1}
    assert parse_json_object("[1, 2]") is None
    assert parse_json_object("не json") is None


async def test_giga_json_valid_first_try() -> None:
    reply = '{"item": "phq9_3", "score": 2, "confidence": 0.8, "evidence": "плохо"}'
    provider, client = _provider([f"```json\n{reply}\n```"])
    res = await provider.complete_json("sys", S("сплю плохо"), ItemScore, {"max_score": 3})
    assert res == ItemScore(item="phq9_3", score=2, confidence=0.8, evidence="плохо")
    assert len(client.payloads) == 1
    payload = client.payloads[0]
    assert payload.temperature == 0.1
    assert [m.role for m in payload.messages] == ["system", "user"]


async def test_giga_retries_then_succeeds() -> None:
    good = json.dumps({"level": "possible", "confidence": 0.7})
    provider, client = _provider([TimeoutError(), "мусор", good])
    res = await provider.complete_json("sys", S("t"), RiskVerdict)
    assert res == RiskVerdict(level="possible", confidence=0.7)
    assert len(client.payloads) == MAX_RETRIES + 1


async def test_giga_gives_up_after_retries() -> None:
    bad_score = json.dumps({"item": "x", "score": 4, "confidence": 0.9, "evidence": ""})
    provider, client = _provider([bad_score] * 5)
    assert await provider.complete_json("sys", S("t"), ItemScore, {"max_score": 3}) is None
    assert len(client.payloads) == MAX_RETRIES + 1


async def test_giga_text() -> None:
    provider, _ = _provider(["", "  Спасибо, что написали.  "])
    assert await provider.complete_text("sys", S("t")) == "Спасибо, что написали."
    provider, _ = _provider([RuntimeError()] * 3)
    assert await provider.complete_text("sys", S("t")) is None


async def test_prompted_tasks_through_gateway() -> None:
    reply = json.dumps({"item": "вопрос про сон", "score": 3, "confidence": 0.9, "evidence": "кошмары"})
    provider, client = _provider([reply, json.dumps({"level": "acute"})])
    gw = LLMGateway(PromptedTasks(provider))
    res = await gw.extract_item("phq9_3", "Проблемы со сном", 3, S("кошмары каждую ночь"))
    assert res is not None and res.item == "phq9_3" and res.score == 3
    assert "Не ставь диагнозы" in client.payloads[0].messages[0].content
    assert await gw.risk_classify(S("t")) == RiskLevel.ACUTE


# --- factory -----------------------------------------------------------------------------------------


def test_factory_selects_provider() -> None:
    assert build_gateway(Settings(_env_file=None)).provider_name == "mock"
    no_key = Settings(_env_file=None, llm_provider="gigachat", gigachat_credentials="")
    assert build_gateway(no_key).provider_name == "mock"
    with_key = Settings(_env_file=None, llm_provider="gigachat", gigachat_credentials="fake-key")
    assert build_gateway(with_key).provider_name == "gigachat"


async def test_module_shortcuts(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(gw_module, "get_gateway", lambda: LLMGateway(MockTasks()))
    assert await gw_module.risk_classify(S("всё бессмысленно")) == RiskLevel.POSSIBLE
    facts = await gw_module.extract_checkin(S("тошнит"))
    assert isinstance(facts, CheckinFacts)


# --- live (only with a real key) ---------------------------------------------------------------------

_LIVE = Settings()


@pytest.mark.skipif(not _LIVE.gigachat_credentials, reason="GIGACHAT_CREDENTIALS is empty")
async def test_live_gigachat_checkin() -> None:
    facts = await build_gateway(_LIVE).extract_checkin(S("сплю плохо, кошмары почти каждую ночь"))
    assert facts is not None and facts.nightmares is True