PTSRaiAgent/app/llm/agent/tools.py

"""Read-only tools of the chat agent. Every tool is bound to the current patient's token on the server side:
the model never passes a token, so it cannot read another patient's data (medical secrecy, 323-FZ art. 13).

Tool results contain clinical data only (no name, SNILS, phone) — they are sent to the LLM provider.
"""

import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any

from app.integrations.emias_record import PatientRecord, emias_records
from app.llm.agent.drug_safety import check_patient, find_drug, load_db
from app.llm.rag import retrieve
from app.llm.rag.index import sentences

log = logging.getLogger(__name__)

RAG_TOP_K = 4
PASSAGE_CHARS = 700
# Patients never get dosages (SPEC §7): sentences with doses are dropped before the text reaches the LLM.
_DOSE = re.compile(r"\d\s*(?:-|–)?\s*\d*\s*(?:мг|mg|мкг|г/сут|таблетк)|\bдоз(?:а|у|ы|е|ой|ировк)", re.I)

TOOL_SPECS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "search_guidelines",
            "description": (
                "Поиск по базе знаний: клинические рекомендации Минздрава РФ (депрессия 2024, ПТСР 2023). "
                "Возвращает фрагменты с источником и страницей. Используй для любых общих вопросов о состоянии, "
                "лечении, сроках действия терапии, побочных эффектах, отмене, образе жизни."
            ),
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string", "description": "Поисковый запрос на русском"}},
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_patient_record",
            "description": (
                "Медицинская карта текущего пациента из ЕМИАС: диагнозы (МКБ-10), сопутствующие заболевания, "
                "аллергии и непереносимости, назначения других врачей и текущее назначение психиатра. "
                "Вызывай, когда ответ зависит от особенностей пациента или его лекарств."
            ),
            "parameters": {"type": "object", "properties": {}},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_drug_info",
            "description": (
                "Справка о препарате по инструкции ГРЛС: состав (действующее и вспомогательные вещества), "
                "противопоказания, взаимодействия, беременность. Аргумент — название препарата (МНН или торговое)."
            ),
            "parameters": {
                "type": "object",
                "properties": {"drug": {"type": "string"}},
                "required": ["drug"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "check_drug_for_patient",
            "description": (
                "Детерминированная проверка препарата для текущего пациента: сопутствующие заболевания, "
                "непереносимости по составу, взаимодействия с другими его лекарствами. Используй перед любым "
                "ответом о том, подходит ли пациенту препарат или сочетание препаратов."
            ),
            "parameters": {
                "type": "object",
                "properties": {"drug": {"type": "string"}},
                "required": ["drug"],
            },
        },
    },
]


@dataclass
class ToolContext:
    token: str
    sources: list[str] = field(default_factory=list)  # citations collected for the answer footer
    used: list[str] = field(default_factory=list)
    _record: PatientRecord | None = None
    _record_loaded: bool = False

    async def record(self) -> PatientRecord | None:
        if not self._record_loaded:
            self._record = await emias_records.get(self.token)
            self._record_loaded = True
        return self._record

    def cite(self, label: str) -> None:
        if label and label not in self.sources:
            self.sources.append(label)


async def run_tool(ctx: ToolContext, name: str, args: dict[str, Any]) -> str:
    """Execute a tool; always returns a JSON string (errors included) so the model can recover."""
    ctx.used.append(name)
    log.info("agent tool=%s", name)  # no arguments: they may echo patient text
    try:
        if name == "search_guidelines":
            return _json(_search(ctx, str(args.get("query", ""))))
        if name == "get_patient_record":
            rec = await ctx.record()
            if rec is None:
                return _json({"found": False})
            ctx.cite("ЕМИАС: медицинская карта")
            return _json({"found": True, **rec.model_dump()})
        if name == "get_drug_info":
            return _json(_drug_info(ctx, str(args.get("drug", ""))))
        if name == "check_drug_for_patient":
            return _json(await _check(ctx, str(args.get("drug", ""))))
        return _json({"error": f"unknown tool {name}"})
    except Exception as exc:  # noqa: BLE001 — a tool failure must not crash the chat
        log.error("agent tool=%s failed: %s", name, type(exc).__name__)
        return _json({"error": "tool failed"})


def _search(ctx: ToolContext, query: str) -> dict[str, Any]:
    hits = retrieve(query, audience="patient", top_k=RAG_TOP_K) if query.strip() else []
    out = []
    for h in hits:
        label = h.chunk.citation.strip("[]")
        ctx.cite(label)
        text = " ".join(x for x in sentences(h.passage) if not _DOSE.search(x))[:PASSAGE_CHARS]
        if text:
            out.append({"source": label, "section": h.chunk.section, "text": text})
    return {"results": out, "found": bool(out)}


def _drug_info(ctx: ToolContext, name: str) -> dict[str, Any]:
    entry = find_drug(name)
    if entry is None:
        return {"found": False, "drug": name}
    ctx.cite(entry.source)
    return {
        "found": True,
        "inn": entry.inn,
        "class": entry.drug_class,
        "composition": entry.composition.model_dump() if entry.composition else None,
        "contraindications": entry.contraindications,
        "interactions": [i.text for i in entry.interactions],
        "pregnancy": entry.pregnancy,
        "source": entry.source,
    }


async def _check(ctx: ToolContext, name: str) -> dict[str, Any]:
    entry = find_drug(name)
    if entry is None:
        return {"found": False, "drug": name, "note": "Препарата нет в справочнике — отвечать «уточните у врача»."}
    rec = await ctx.record()
    if rec is None:
        return {"found": True, "inn": entry.inn, "record": False, "findings": []}
    ctx.cite(entry.source)
    ctx.cite("ЕМИАС: медицинская карта")
    findings = check_patient(entry, rec, load_db())
    return {
        "found": True,
        "inn": entry.inn,
        "record": True,
        "findings": [f.model_dump() for f in findings],
        "note": "Пустой список не означает, что препарат безопасен: решение о приёме принимает врач.",
    }


def _json(data: Any) -> str:
    return json.dumps(data, ensure_ascii=False)