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
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 (
    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
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 = "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":
        # TODO(GigaChat): function calling via the `gigachat` package (functions=..., message.function_call).
        log.warning("chat agent: gigachat provider not implemented yet — disabled")
        return None
    return None


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
    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 provider.chat(_wire(messages) if provider.name != "mock" else messages, TOOL_SPECS)
                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": [
                            {
                                "id": c.id,
                                "type": "function",
                                "function": {"name": c.name, "arguments": json.dumps(c.arguments, ensure_ascii=False)},
                            }
                            for c in turn.tool_calls
                        ],
                    }
                )
                for c in turn.tool_calls:
                    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", provider.name, type(exc).__name__ or "timeout")
        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)


def render(answer: AgentAnswer) -> str:
    """Patient-facing text: answer + short sources footer."""
    if not answer.sources:
        return answer.text
    return f"{answer.text}\n\nИсточники: " + "; ".join(answer.sources)