PTSRaiAgent/app/llm/agent/__init__.py

"""Chat agent: tool-using LLM over the knowledge base (RAG), the ЕМИАС record and the drug-safety reference.

    from app.llm.agent import answer_chat
    reply = await answer_chat(token, scrubbed_question)   # None → caller falls back to extractive RAG / med_qa

Guarantees (on top of what the dialog engine already did — scrub, consent, safety check; ACUTE never gets here):
  * only scrubbed text goes to the provider (`assert_scrubbed`), tool results carry no identifiers;
  * tools are bound to the current token server-side; read-only;
  * the final text passes the gateway's forbidden-output filter (diagnoses, dose/drug advice); a violation or any
    provider error → None, so the patient gets the safe extractive answer instead;
  * logs carry tool names and counts only.
"""

import asyncio
import json
import logging
import re
from typing import Any

from pydantic import BaseModel

from app.config import Settings, get_settings
from app.core.types import ScrubbedText
from app.llm.agent.prompts import SYSTEM_PROMPT
from app.llm.agent.providers import (
    GigaChatToolProvider,
    MockAgentProvider,
    OpenRouterProvider,
    ProviderError,
    ToolChatProvider,
)
from app.llm.agent.tools import TOOL_SPECS, ToolContext, run_tool
from app.llm.gateway import is_forbidden
from app.llm.guard import assert_scrubbed

log = logging.getLogger(__name__)

MAX_ANSWER_CHARS = 1200
MAX_TOOL_CALLS_PER_TURN = 4  # seen live: one «как дела?» turn issued 23 parallel guideline searches
# Greetings / small talk: answered without the LLM (instant, free, never drags medical findings into «привет»).
SMALL_TALK = re.compile(
    r"^\W*(прив\w*|здравств\w*|здаров\w*|добр(ый|ое|ого)\s+\w+|хай|hello|hi|ку)\b"
    r"|\b(как\s+(дела|ты|вы|жизнь|поживаешь)|что\s+делаешь|чем\s+занят\w*|спишь|кто\s+ты|ты\s+(бот|робот|кто))\b",
    re.I,
)
SMALL_TALK_MAX_WORDS = 8
SMALL_TALK_REPLY = (
    "Здравствуйте! Я на связи. Могу рассказать о вашем лечении: как действует препарат и когда ждать эффекта, "
    "какие бывают побочные эффекты, что можно и нельзя сочетать с вашими таблетками, что в составе. "
    "Спросите — отвечу по вашей карте, инструкции и клиническим рекомендациям."
)
NUDGE = "Сформулируй ответ пациенту по результатам инструментов выше (3–6 предложений)."
_ALLOWED_MSG_KEYS = {"role", "content", "tool_calls", "tool_call_id"}


class AgentAnswer(BaseModel):
    text: str
    sources: list[str]
    tools: list[str]
    provider: str


def build_provider(settings: Settings | None = None) -> ToolChatProvider | None:
    s = settings or get_settings()
    mode = s.chat_agent
    if mode == "auto":
        mode = "gigachat" if s.gigachat_credentials else "openrouter" if s.openrouter_api_key else "off"
    if mode == "openrouter":
        if not s.openrouter_api_key:
            log.warning("chat agent: openrouter selected without OPENROUTER_API_KEY — disabled")
            return None
        return OpenRouterProvider(s)
    if mode == "mock":
        return MockAgentProvider()
    if mode == "gigachat":
        if not s.gigachat_credentials:
            log.warning("chat agent: gigachat selected without GIGACHAT_CREDENTIALS — disabled")
            return None
        return GigaChatToolProvider(s)
    return None


def is_small_talk(text: str) -> bool:
    """Short greeting / chit-chat without a medical question in it."""
    words = re.findall(r"\w+", text)
    medical = r"(лекарств|таблет|препарат|побоч|доз|лечени|можно|ибупрофен|нурофен|алкогол|пив|сертралин|сон|сплю)"
    if len(words) > SMALL_TALK_MAX_WORDS or re.search(medical, text, re.I):
        return False
    return bool(SMALL_TALK.search(text))


async def _chat_with_retry(provider: ToolChatProvider, messages: list[dict[str, Any]]):
    """One retry: OpenRouter occasionally fails a single call (seen live as intermittent ProviderError)."""
    wire = _wire(messages) if provider.name == "openrouter" else messages  # gigachat needs the tool `name`
    try:
        return await provider.chat(wire, TOOL_SPECS)
    except ProviderError as exc:
        log.warning("chat agent provider=%s retry after %s", provider.name, exc)
        return await provider.chat(wire, TOOL_SPECS)


async def _patient_brief(ctx: ToolContext) -> str:
    """One line of context so the model knows what «мои таблетки» means (drug names only, no identifiers)."""
    rec = await ctx.record()
    if rec is None or not rec.current_prescriptions:
        return "\n\nТекущее назначение психиатра: нет данных."
    drugs = ", ".join(p.drug for p in rec.current_prescriptions)
    return f"\n\nТекущее назначение психиатра: {drugs}. «Мои таблетки» = это назначение."


def _wire(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Strip internal keys (tool `name`) before sending."""
    return [{k: v for k, v in m.items() if k in _ALLOWED_MSG_KEYS} for m in messages]


async def answer_chat(
    token: str, question: ScrubbedText, *, provider: ToolChatProvider | None = None, settings: Settings | None = None
) -> AgentAnswer | None:
    s = settings or get_settings()
    provider = provider or build_provider(s)
    if provider is None:
        return None
    try:
        assert_scrubbed(question)
    except ValueError:
        log.warning("chat agent: unscrubbed input rejected")
        return None
    if is_small_talk(question):
        return AgentAnswer(text=SMALL_TALK_REPLY, sources=[], tools=[], provider="small_talk")
    ctx = ToolContext(token=token)
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": SYSTEM_PROMPT + await _patient_brief(ctx)},
        {"role": "user", "content": str(question)},
    ]
    nudged = False
    try:
        async with asyncio.timeout(s.chat_agent_timeout_s * 2):
            for _ in range(s.chat_agent_max_steps):
                turn = await _chat_with_retry(provider, messages)
                if not turn.tool_calls:
                    if not (turn.content or "").strip() and not nudged:
                        # Some models (Gemini Flash) occasionally end with an empty message after tool results.
                        nudged = True
                        messages.append({"role": "user", "content": NUDGE})
                        continue
                    return _finish(turn.content or "", ctx, provider.name)
                messages.append(
                    {
                        "role": "assistant",
                        "content": turn.content,
                        "tool_calls": [  # the calls we actually run (capped) — history must match the results
                            {
                                "id": c.id,
                                "type": "function",
                                "function": {"name": c.name, "arguments": json.dumps(c.arguments, ensure_ascii=False)},
                            }
                            for c in turn.tool_calls[:MAX_TOOL_CALLS_PER_TURN]
                        ],
                    }
                )
                for c in turn.tool_calls[:MAX_TOOL_CALLS_PER_TURN]:
                    result = await run_tool(ctx, c.name, c.arguments)
                    messages.append({"role": "tool", "tool_call_id": c.id, "name": c.name, "content": result})
            log.warning("chat agent: step limit reached tools=%s", ",".join(ctx.used))
            return None
    except (ProviderError, TimeoutError) as exc:
        log.error("chat agent provider=%s failed: %s %s", provider.name, type(exc).__name__, exc)
        return None
    except Exception:  # noqa: BLE001 — the agent is optional: any bug must degrade to the extractive answer, not a 500
        log.exception("chat agent provider=%s crashed", provider.name)
        return None


def _finish(text: str, ctx: ToolContext, provider: str) -> AgentAnswer | None:
    text = text.strip()
    if not text:
        return None
    if is_forbidden(text):
        log.warning("chat agent: forbidden output filtered provider=%s", provider)
        return None
    if len(text) > MAX_ANSWER_CHARS:
        text = text[:MAX_ANSWER_CHARS].rsplit(" ", 1)[0] + "…"
    log.info("chat agent ok provider=%s tools=%s sources=%d", provider, ",".join(ctx.used), len(ctx.sources))
    return AgentAnswer(text=text, sources=ctx.sources[:4], tools=ctx.used, provider=provider)


_MD_BULLET = re.compile(r"^\s*[*-]\s+", re.M)
_MD_MARKS = re.compile(r"\*\*|__|^#+\s*", re.M)
SOURCES_MIN_CHARS = 160  # short replies (greetings, small talk) go without the sources footer


def render(answer: AgentAnswer) -> str:
    """Patient-facing text: answer (markdown stripped to plain text with «•» bullets) + short sources footer."""
    text = _MD_MARKS.sub("", _MD_BULLET.sub("• ", answer.text)).strip()
    if not answer.sources or len(text) < SOURCES_MIN_CHARS:
        return text
    return f"{text}\n\nИсточники: " + "; ".join(answer.sources)