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


async def test_region_block_string_error_degrades():
    # Seen live from the RF server: 403 {"success": false, "error": "Access denied by security policy."}
    client = httpx.AsyncClient(
        transport=httpx.MockTransport(lambda r: httpx.Response(403, json={"success": False, "error": "Access denied"}))
    )
    provider = OpenRouterProvider(_settings(), client=client)
    assert await agent.answer_chat(ANDREY, scrub("вопрос"), provider=provider, settings=_settings()) is None


async def test_unexpected_provider_bug_degrades():
    class Broken:
        name = "broken"

        async def chat(self, messages, tools):
            raise AttributeError("boom")

    assert await agent.answer_chat(ANDREY, scrub("вопрос"), provider=Broken(), settings=_settings()) is None


# ------------------------------------------------------------------------------------------------- gigachat


class _FakeGiga:
    """Mimics gigachat.GigaChat.achat: first a function_call, then the final text."""

    def __init__(self):
        self.chats = []

    async def achat(self, chat):
        from types import SimpleNamespace as NS

        self.chats.append(chat)
        if len(self.chats) == 1:
            fc = NS(name="check_drug_for_patient", arguments={"drug": "ибупрофен"})
            msg = NS(content="", function_call=fc)
        else:
            msg = NS(
                content="Ибупрофен вместе с сертралином повышает риск кровотечений. Обсудите с врачом.",
                function_call=None,
            )
        return NS(choices=[NS(message=msg)], usage=NS(total_tokens=42))


async def test_gigachat_tool_loop(sertraline):
    from app.llm.agent.providers import GigaChatToolProvider

    fake = _FakeGiga()
    s = Settings(_env_file=None, chat_agent="gigachat", gigachat_credentials="x", chat_agent_timeout_s=5.0)
    ans = await agent.answer_chat(
        ANDREY, scrub("Можно ибупрофен?"), provider=GigaChatToolProvider(s, client=fake), settings=s
    )
    assert ans is not None and "кровотеч" in ans.text and ans.provider == "gigachat"
    second = fake.chats[1]
    roles = [str(m.role.value if hasattr(m.role, "value") else m.role) for m in second.messages]
    assert roles == ["system", "user", "assistant", "function"]
    assert second.messages[2].function_call.name == "check_drug_for_patient"
    assert second.messages[3].name == "check_drug_for_patient"
    assert {f.name for f in second.functions} == {
        "search_guidelines",
        "get_patient_record",
        "get_drug_info",
        "check_drug_for_patient",
    }


def test_build_provider_prefers_gigachat():
    from app.llm.agent.providers import GigaChatToolProvider

    s = Settings(_env_file=None, chat_agent="auto", gigachat_credentials="x", openrouter_api_key="y")
    assert isinstance(agent.build_provider(s), GigaChatToolProvider)


async def test_web_demo_token_shares_andreys_record():
    rec = await emias_records.get("tok-andrey")
    assert rec is not None and any("язвенная" in c.name.lower() for c in rec.conditions)


@pytest.mark.parametrize("text", ["прив", "Привет!", "как дела у тебя ? что делаешь ? спишь ?", "добрый вечер"])
async def test_small_talk_without_llm(text):
    provider, sent = _fake_openrouter("x", _CALLS)
    ans = await agent.answer_chat(ANDREY, scrub(text), provider=provider, settings=_settings())
    assert ans is not None and ans.provider == "small_talk" and sent == []


async def test_medical_question_is_not_small_talk():
    assert not agent.is_small_talk("привет, можно мне ибупрофен?")


async def test_tool_calls_are_capped(sertraline):
    many = [
        {"id": f"s{i}", "type": "function", "function": {"name": "search_guidelines", "arguments": '{"query": "x"}'}}
        for i in range(20)
    ]
    provider, sent = _fake_openrouter("Ответ по рекомендациям, обсудите с врачом.", many)
    ans = await agent.answer_chat(ANDREY, scrub("что такое ПТСР?"), provider=provider, settings=_settings())
    assert ans is not None and len(ans.tools) == agent.MAX_TOOL_CALLS_PER_TURN
    tool_msgs = [m for m in sent[1]["messages"] if m["role"] == "tool"]
    assistant = [m for m in sent[1]["messages"] if m["role"] == "assistant"][0]
    assert len(tool_msgs) == len(assistant["tool_calls"]) == agent.MAX_TOOL_CALLS_PER_TURN