PTSRaiAgent/tests/test_llm_agent.py

import json

import httpx
import pytest

from app.config import Settings
from app.core.privacy import pseudo_id, scrub
from app.core.types import ScrubbedText
from app.integrations.base import Prescription, integrations
from app.integrations.emias_record import emias_records
from app.llm import agent
from app.llm.agent.drug_safety import check_patient, find_drug
from app.llm.agent.providers import MockAgentProvider, OpenRouterProvider
from app.llm.agent.tools import ToolContext, run_tool

ANDREY = pseudo_id("900-000-013 34")
ANNA = pseudo_id("900-000-023 04")
SERGEY = pseudo_id("900-000-037 02")


class _Rx:
    def __init__(self, drug: str) -> None:
        self.drug = drug

    async def get_prescriptions(self, token: str) -> list[Prescription]:
        return [Prescription(drug=self.drug, dose="50 мг", instructions="утром после еды")]


@pytest.fixture
def sertraline(monkeypatch):
    monkeypatch.setattr(integrations, "prescriptions", _Rx("сертралин"))


# ----------------------------------------------------------------------------------------------- drug safety


@pytest.mark.parametrize("name", ["сертралин", "сертралине", "Золофт", "сертралин 50 мг", "sertraline"])
def test_find_drug_forms(name):
    assert find_drug(name).inn == "сертралин"


def test_find_drug_unknown():
    assert find_drug("аспирин") is None


async def test_andrey_sertraline_findings(sertraline):
    rec = await emias_records.get(ANDREY)
    findings = check_patient(find_drug("сертралин"), rec)
    kinds = {(f.kind, f.severity) for f in findings}
    assert ("condition", "high") in kinds  # ulcer history
    assert any(f.kind == "interaction" and "ибупрофен" in f.text for f in findings)


async def test_ibuprofen_sees_current_ssri(sertraline):
    rec = await emias_records.get(ANDREY)
    findings = check_patient(find_drug("ибупрофен"), rec)
    assert any("СИОЗС" in f.text and "сертралин" in f.text for f in findings)


async def test_pregnancy_and_triptan(monkeypatch):
    monkeypatch.setattr(integrations, "prescriptions", _Rx("флуоксетин"))
    anna = await emias_records.get(ANNA)
    assert any(f.kind == "pregnancy" for f in check_patient(find_drug("флуоксетин"), anna))
    monkeypatch.setattr(integrations, "prescriptions", _Rx("пароксетин"))
    sergey = await emias_records.get(SERGEY)
    assert any("Триптаны" in f.text for f in check_patient(find_drug("пароксетин"), sergey))


# ----------------------------------------------------------------------------------------------------- tools


async def test_record_tool_has_no_identifiers(sertraline):
    out = await run_tool(ToolContext(token=ANDREY), "get_patient_record", {})
    data = json.loads(out)
    assert data["found"] and data["current_prescriptions"][0]["drug"] == "сертралин"
    assert "900-000" not in out and "Андрей" not in out


async def test_search_tool_drops_doses():
    out = json.loads(await run_tool(ToolContext(token=ANDREY), "search_guidelines", {"query": "дозы СИОЗС при ПТСР"}))
    for r in out["results"]:
        assert " мг" not in r["text"]


async def test_unknown_tool_is_json_error():
    assert "error" in json.loads(await run_tool(ToolContext(token=ANDREY), "rm_rf", {}))


# ------------------------------------------------------------------------------------------------ agent loop


def _settings(**kw) -> Settings:
    base = {"openrouter_api_key": "test", "chat_agent": "openrouter", "chat_agent_timeout_s": 5.0}
    return Settings(_env_file=None, **{**base, **kw})


def _fake_openrouter(final_text: str, calls: list[dict]) -> tuple[OpenRouterProvider, list[dict]]:
    sent: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        body = json.loads(request.content)
        sent.append(body)
        if len(sent) == 1:
            msg = {"role": "assistant", "content": None, "tool_calls": calls}
        else:
            msg = {"role": "assistant", "content": final_text}
        return httpx.Response(200, json={"choices": [{"message": msg}], "usage": {"total_tokens": 10}})

    client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    return OpenRouterProvider(_settings(), client=client), sent


_CALLS = [
    {"id": "a", "type": "function", "function": {"name": "get_patient_record", "arguments": "{}"}},
    {
        "id": "b",
        "type": "function",
        "function": {"name": "check_drug_for_patient", "arguments": '{"drug": "ибупрофен"}'},
    },
]


async def test_agent_openrouter_tool_loop(sertraline):
    text = (
        "В вашей карте отмечена язвенная болезнь, а ибупрофен с сертралином повышает риск кровотечений. "
        "Обсудите с врачом."
    )
    provider, sent = _fake_openrouter(text, _CALLS)
    ans = await agent.answer_chat(
        ANDREY, scrub("Можно ибупрофен? Мой телефон 89161234567"), provider=provider, settings=_settings()
    )
    assert ans is not None and ans.text == text
    assert ans.tools == ["get_patient_record", "check_drug_for_patient"]
    assert any("ЕМИАС" in s for s in ans.sources)
    first, second = sent
    assert "89161234567" not in json.dumps(first, ensure_ascii=False)  # scrubbed before the provider
    tool_msgs = [m for m in second["messages"] if m["role"] == "tool"]
    assert len(tool_msgs) == 2 and all("name" not in m for m in tool_msgs)
    assert first["tools"] and first["model"]


async def test_forbidden_output_rejected(sertraline):
    provider, _ = _fake_openrouter("Увеличьте дозу сертралина вдвое.", _CALLS)
    assert await agent.answer_chat(ANDREY, scrub("что делать"), provider=provider, settings=_settings()) is None


async def test_provider_error_returns_none():
    client = httpx.AsyncClient(transport=httpx.MockTransport(lambda r: httpx.Response(500)))
    provider = OpenRouterProvider(_settings(), client=client)
    assert await agent.answer_chat(ANDREY, scrub("вопрос"), provider=provider, settings=_settings()) is None


async def test_unscrubbed_input_never_sent():
    provider, sent = _fake_openrouter("ok", [])
    assert await agent.answer_chat(ANDREY, ScrubbedText("звоните 8 916 123-45-67"), provider=provider) is None
    assert sent == []


async def test_mock_agent_answers_offline(sertraline):
    ans = await agent.answer_chat(
        ANDREY, scrub("Можно ли мне ибупрофен вместе с моими таблетками?"), provider=MockAgentProvider()
    )
    assert ans is not None and "кровотеч" in ans.text and ans.provider == "mock"


def test_build_provider_modes():
    assert agent.build_provider(Settings(_env_file=None, chat_agent="auto")) is None
    assert isinstance(agent.build_provider(_settings(chat_agent="auto")), OpenRouterProvider)
    assert isinstance(agent.build_provider(Settings(_env_file=None, chat_agent="mock")), MockAgentProvider)
    assert agent.build_provider(Settings(_env_file=None, chat_agent="openrouter")) is None  # no key
    assert agent.build_provider(Settings(_env_file=None, chat_agent="off")) is None


async def test_check_defaults_to_current_prescription(sertraline):
    out = json.loads(await run_tool(ToolContext(token=ANDREY), "check_drug_for_patient", {"drug": "мои таблетки"}))
    assert out["found"] and out["inn"] == "сертралин"
    assert "лактоз" in out["intolerance_check"] and "нет" in out["intolerance_check"]
    assert any("Алкоголь" in t for t in out["general_interactions"])


async def test_empty_final_answer_is_nudged(sertraline):
    replies = iter(
        [
            {"role": "assistant", "content": None, "tool_calls": _CALLS},
            {"role": "assistant", "content": ""},
            {"role": "assistant", "content": "Алкоголь во время лечения не рекомендуется. Обсудите с врачом."},
        ]
    )
    sent = []

    def handler(request):
        sent.append(json.loads(request.content))
        return httpx.Response(200, json={"choices": [{"message": next(replies)}]})

    provider = OpenRouterProvider(_settings(), client=httpx.AsyncClient(transport=httpx.MockTransport(handler)))
    ans = await agent.answer_chat(ANDREY, scrub("пиво можно?"), provider=provider, settings=_settings())
    assert ans is not None and "Алкоголь" in ans.text
    assert sent[0]["tool_choice"] == "required" and "tool_choice" not in sent[1]
    assert "сертралин" in sent[0]["messages"][0]["content"]  # patient brief in the system prompt