PTSRaiAgent/app/llm/agent/providers.py

"""Tool-calling chat providers for the agent.

* `OpenRouterProvider` — OpenAI-compatible `/chat/completions` with `tools`. TEMPORARY for the prototype on
  synthetic data: OpenRouter is a foreign service, forbidden for real PII (152-FZ art. 18 p. 5). Replace with
  GigaChat (function calling) before any real patient — same `ToolChatProvider` contract.
* `MockAgentProvider` — offline, deterministic: plans tool calls from keywords and composes an answer from the
  tool results. Keeps the demo working without network (ADR-004).
"""

import json
import logging
import re
from typing import Any, Protocol

import httpx
from pydantic import BaseModel, Field

from app.config import Settings

log = logging.getLogger(__name__)


class ToolCall(BaseModel):
    id: str
    name: str
    arguments: dict[str, Any] = Field(default_factory=dict)


class AssistantTurn(BaseModel):
    content: str | None = None
    tool_calls: list[ToolCall] = Field(default_factory=list)


class ToolChatProvider(Protocol):
    name: str

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> AssistantTurn: ...


class ProviderError(RuntimeError):
    pass


class OpenRouterProvider:
    name = "openrouter"

    def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
        self._key = settings.openrouter_api_key
        self._model = settings.openrouter_model
        self._url = settings.openrouter_base_url.rstrip("/") + "/chat/completions"
        self._timeout = settings.chat_agent_timeout_s
        self._client = client

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> AssistantTurn:
        body = {"model": self._model, "messages": messages, "tools": tools, "temperature": 0.1, "max_tokens": 900}
        headers = {"Authorization": f"Bearer {self._key}", "X-Title": "PTSRaiAgent prototype (synthetic data)"}
        try:
            if self._client is not None:
                resp = await self._client.post(self._url, json=body, headers=headers, timeout=self._timeout)
            else:
                async with httpx.AsyncClient(timeout=self._timeout) as client:
                    resp = await client.post(self._url, json=body, headers=headers)
        except httpx.HTTPError as exc:
            raise ProviderError(type(exc).__name__) from exc
        if resp.status_code != 200:
            raise ProviderError(f"http {resp.status_code}")
        try:
            msg = resp.json()["choices"][0]["message"]
        except (ValueError, KeyError, IndexError) as exc:
            raise ProviderError("bad response") from exc
        calls = []
        for c in msg.get("tool_calls") or []:
            fn = c.get("function") or {}
            try:
                args = json.loads(fn.get("arguments") or "{}")
            except json.JSONDecodeError:
                args = {}
            calls.append(ToolCall(id=c.get("id") or fn.get("name", "call"), name=fn.get("name", ""), arguments=args))
        usage = resp.json().get("usage") or {}
        log.info("openrouter ok calls=%d tokens=%s", len(calls), usage.get("total_tokens"))
        return AssistantTurn(content=msg.get("content"), tool_calls=calls)


# ------------------------------------------------------------------------------------------------ offline mock

_DRUG_WORDS = re.compile(
    r"\b(сертралин\w*|пароксетин\w*|эсциталопрам\w*|флуоксетин\w*|ибупрофен\w*|нурофен\w*|золофт\w*)", re.I
)
_MED_HINT = re.compile(r"(таблет|лекарств|препарат|состав|можно ли|совмест|вместе|противопоказ|аллерг|лактоз)", re.I)


class MockAgentProvider:
    """Plans: record → (drug check for each mentioned drug, or the current prescription for medicine questions)
    → guidelines search; then writes a short answer strictly from the tool results."""

    name = "mock"

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> AssistantTurn:
        question = next(m["content"] for m in reversed(messages) if m["role"] == "user")
        done = [m for m in messages if m["role"] == "tool"]
        if not done:
            drugs = [m.group(1).lower() for m in _DRUG_WORDS.finditer(question)]
            calls = [ToolCall(id="c0", name="get_patient_record")]
            calls += [
                ToolCall(id=f"d{i}", name="check_drug_for_patient", arguments={"drug": d}) for i, d in enumerate(drugs)
            ]
            if drugs and re.search(r"состав|лактоз|желатин|аллерг|непереносим", question, re.I):
                calls += [
                    ToolCall(id=f"i{i}", name="get_drug_info", arguments={"drug": d}) for i, d in enumerate(drugs)
                ]
            calls.append(ToolCall(id="s0", name="search_guidelines", arguments={"query": question}))
            return AssistantTurn(tool_calls=calls)
        mine = re.search(r"мо(и|ими|их|ей|я|ё) (таблет|лекарств|препарат)", question, re.I)
        checked = [m for m in done if m.get("name") == "check_drug_for_patient"]
        if (not checked or mine) and len(checked) < 3 and _MED_HINT.search(question):
            record = _tool_json(done, "get_patient_record")
            rx = (record or {}).get("current_prescriptions") or []
            checked_args = {str(json.loads(m["content"]).get("inn", "")) for m in checked}
            rx = [r for r in rx if not any(r["drug"].lower().startswith(c[:-1]) for c in checked_args if c)]
            if rx:
                return AssistantTurn(
                    tool_calls=[ToolCall(id="d9", name="check_drug_for_patient", arguments={"drug": rx[0]["drug"]})]
                )
        return AssistantTurn(content=_compose(done))


def _tool_json(done: list[dict[str, Any]], name: str) -> dict[str, Any] | None:
    for m in done:
        if m.get("name") == name:
            try:
                return json.loads(m["content"])
            except json.JSONDecodeError:
                return None
    return None


def _compose(done: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    seen: set[str] = set()
    for m in done:
        if m.get("name") != "check_drug_for_patient":
            continue
        data = json.loads(m["content"])
        if not data.get("found") or data["inn"] in seen:
            continue
        seen.add(data["inn"])
        findings = data.get("findings") or []
        if findings:
            parts.append(f"По вашей карте и инструкции к препарату «{data['inn']}» есть важное:")
            parts += [f"• {f['text']}" for f in findings[:3]]
        else:
            parts.append(f"В справочнике не нашлось особых ограничений по «{data['inn']}» с учётом вашей карты.")
    for m in done:
        if m.get("name") != "get_drug_info":
            continue
        data = json.loads(m["content"])
        comp = data.get("composition") if data.get("found") else None
        if comp:
            record = _tool_json(done, "get_patient_record") or {}
            allergies = record.get("allergies") or []
            if allergies:
                names = ", ".join(a["name"] for a in allergies)
                hit = {t for a in allergies for t in a.get("tags", [])} & set(comp.get("excipient_tags", []))
                verdict = "это вещество есть" if hit else "этого вещества нет, но состав зависит от производителя"
                parts.append(f"В вашей карте: {names}. В примере состава «{data['inn']}» {verdict}.")
            excipients = ", ".join(comp["excipients"])
            parts.append(f"Состав ({comp['form']}): {comp['active']}; вспомогательные: {excipients}. {comp['note']}")
    search = _tool_json(done, "search_guidelines") or {}
    results = search.get("results") or []
    if results and not parts:  # medicine findings already answer the question; a KR quote would dilute it
        text = results[0]["text"].strip()
        parts.append(f"В клинических рекомендациях: «{text[:400]}{'…' if len(text) > 400 else ''}»")
    if not parts:
        return ""
    parts.append("Любые изменения в лечении обсуждайте с вашим врачом.")
    return "\n".join(parts)