PTSRaiAgent/app/llm/rag/qa.py

"""Grounded answers from the knowledge base: every sentence shown is quoted from a cited chunk.

Pipeline (deterministic except for the optional patient rephrasing):

1. **Dose guard (patient).** Dose questions are never answered — the dose is the doctor's decision.
2. **Retrieve.** Top `answer.max_chunks` chunks for the audience (`team` documents never reach patients
   or doctors, see `config/rag/sources.yaml`).
3. **Sentences.** Chunks are split into sentences; KR boilerplate is stripped (literature refs "[1, 15]").
   Doctors also get table rows as "column: value; ..." lines and evidence grades. For patients, sentences
   about prescribing/dosing or with forbidden content (`QUOTE_FILTER`) are dropped.
4. **Passage.** In each chunk the best sentence (idf-weighted overlap with the query) is extended with
   adjacent relevant sentences into one contiguous passage — the answer reads as written in the source,
   not as a collage. The best passage, plus a second one if it scores close, is kept.
5. **No-answer gate.** If the passages cover less than `answer.min_coverage` of the question's idf mass
   (e.g. "рецепт борща"), there is no answer -> `None`.
6. **Rephrase (patient only).** The excerpt may be simplified via `gateway.phrase("rag_answer")`; a
   rephrasing that adds content (`med_qa.adds_content`) is rejected and the excerpt is shown verbatim.
   Doctors always get the verbatim text.
7. **Framing.** Intro, citations "[КР ПТСР 2023, стр. 28]" and the doctor note from `config/texts/rag.yaml`
   are added outside the LLM, so they are always present.

Safety screening of the incoming message is the caller's job (guardrail #2): `answer()` expects text that
already passed `app.core.safety` and `privacy.scrub()`. Logs carry only codes and counts.
"""

import logging
import re
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path

import yaml
from pydantic import BaseModel

from app.config import get_settings
from app.core.types import ScrubbedText
from app.llm.gateway import FORBIDDEN_OUTPUT, MAX_REPLY_CHARS, get_gateway
from app.llm.med_qa import DOSE_QUESTION, DOSING_SENTENCE, adds_content
from app.llm.rag.corpus import TABLE_DIVIDER, AnswerConfig, Audience, Chunk, load_config, split_sentences
from app.llm.rag.index import BM25Index, get_index, tokens

log = logging.getLogger(__name__)

PHRASE_INTENT = "rag_answer"

LIT_REFS = re.compile(r"\s*\[\d+(?:\s*[-–,]\s*\d+)*\s*,?\]")  # "[1, 15, 16]", "[48]", "[1-3]"
GRADE = re.compile(r"\s*Уровень убедительности рекомендаций\s+[ABCАВС]\b.*$", re.IGNORECASE)
BULLET = re.compile(r"^\s*(?:[-–—●•▪◦*]|\d+[.)])\s+")
COMMENT_LABEL = re.compile(r"^Комментари[йи]\s*[:.]\s*")
# KR sentences about prescribing are for the doctor: a patient reading "рекомендуется назначение
# пароксетина" or "дозировка титруется" out of context takes it as advice (SPEC §7: no drug/dose advice).
PRESCRIBING = re.compile(
    r"назнач|дозиров|титр|мг\b|мг/|доз[аеуыо]\b|дозы\b|/сут|в сутки|раза? в день|разово"
    r"|\d+\s*(?:[-–]\s*\d+\s*)?таблет",
    re.IGNORECASE,
)
# "что такое флэшбек" -> the term; its definition "Флэшбек (flashback – англ.) – ..." is preferred.
DEFINITION_QUESTION = re.compile(r"^\s*что (?:такое|значит|означает|это за)\s+(.+?)[?.!\s]*$", re.IGNORECASE)
LEADING_PLACEHOLDERS = re.compile(r"^(?:\s*\[[A-Z_]+\][\s,.:;!-]*)+")  # "[NAME], что такое ПТСР?"
DEFINITION_HEAD = re.compile(r"\s[–—-]\s(?![^(]*\))|\s[–—-]?\s?это\s")  # dash outside "(… – англ.)", or "это"
DEFINITION_HEAD_CHARS = 60
DEFINITION_BONUS = 2.0  # multiplier for a sentence that starts with the asked term (picks it within a chunk)
# The gateway's output filter targets the bot's own words. Two of its patterns also hit neutral quotes from
# the KR — the definition "ПТСР – это расстройство, которое..." and "диагноз ставится на основании..." —
# and would cut exactly the patient-information text. Quotes are checked against the rest of the filter
# (dose/drug imperatives, "у вас депрессия", confidentiality promises).
QUOTE_EXEMPT_PREFIXES = ("диагноз", "(это|похоже на|признаки)")
QUOTE_FILTER = tuple(p for p in FORBIDDEN_OUTPUT if not p.pattern.startswith(QUOTE_EXEMPT_PREFIXES))
MIN_SENTENCE_CHARS = 25
LEAD_IN_CHARS = 80
# A best sentence opening with a reference back ("Это помогает...") needs the sentence it refers to.
ANAPHORA = re.compile(
    r"^(Это|Эти|Этот|Эта|Он|Она|Оно|Они|Такие|Такой|Данн\w+|При этом|Кроме того|Однако|Поэтому)\b"
)  # a best sentence this short ("Что такое ПТСР?") pulls in the next one as its answer


class Citation(BaseModel):
    doc: str  # document id from sources.yaml
    short: str  # "КР ПТСР 2023"
    title: str
    section: str
    pages: str  # "28" | "28–29"

    @property
    def label(self) -> str:
        return f"{self.short}, стр. {self.pages}"


class RagAnswer(BaseModel):
    text: str  # framed answer: intro + body + sources (+ doctor note for patients)
    excerpt: str  # the quoted passages, verbatim
    citations: list[Citation]
    audience: Audience
    coverage: float  # share of the question covered by the passages (0..1)
    phrased_by_llm: bool


class AudienceTexts(BaseModel):
    fallback: str
    intro: str
    sources_label: str
    doctor_note: str = ""


@dataclass(frozen=True)
class Passage:
    chunk: Chunk
    text: str
    score: float
    definition: bool = False  # starts with the term asked in "что такое X?"


@lru_cache
def load_rag_texts(path: Path | None = None) -> dict[str, AudienceTexts]:
    path = path or get_settings().config_dir / "texts" / "rag.yaml"
    raw = yaml.safe_load(path.read_text(encoding="utf-8"))
    return {k: AudienceTexts.model_validate(v) for k, v in raw.items()}


def fallback_text(audience: Audience = "patient") -> str:
    """What the caller sends when `answer()` returns `None`."""
    return load_rag_texts()[audience].fallback


def _clean(sentence: str, *, keep_grade: bool) -> str:
    s = BULLET.sub("", sentence)
    s = COMMENT_LABEL.sub("", s)
    if not keep_grade:  # evidence grades matter to doctors, to a patient they are noise
        s = GRADE.sub("", s)
    s = LIT_REFS.sub("", s)
    return " ".join(s.split()).strip(" ;")


def _cells(line: str) -> list[str]:
    return [c.strip() for c in line.strip().strip("|").split("|")]


def _row_text(header: list[str], row: list[str]) -> str:
    pairs = [f"{h}: {v}" if h else v for h, v in zip(header, row, strict=False) if v]
    return LIT_REFS.sub("", "; ".join(pairs))


def chunk_sentences(chunk: Chunk, *, clinical: bool) -> list[str | None]:
    """Readable sentences of a chunk in reading order. `clinical` (doctor/team) keeps table rows as
    "column: value" lines and evidence grades. `None` marks a dropped row: passages do not span it."""
    out: list[str | None] = []
    lines = chunk.text.splitlines()
    header: list[str] | None = None
    for i, line in enumerate(lines):
        if line.startswith("|"):
            if TABLE_DIVIDER.match(line):
                continue
            if i + 1 < len(lines) and TABLE_DIVIDER.match(lines[i + 1]):
                header = _cells(line)
                continue
            out.append(_row_text(header or [], _cells(line)) if clinical else None)
            continue
        header = None
        out.extend(s for s in (_clean(p, keep_grade=clinical) for p in split_sentences(line)) if s)
    return out


def _allowed(sentence: str | None, audience: Audience) -> bool:
    if sentence is None or len(sentence) < MIN_SENTENCE_CHARS and not sentence.endswith(("?", ":")):
        return False
    if sentence[0].islower():
        return False  # tail of a paragraph cut by the chunk boundary: unreadable without its start
    if audience == "patient":
        return not (
            DOSING_SENTENCE.search(sentence)
            or PRESCRIBING.search(sentence)
            or any(p.search(sentence) for p in QUOTE_FILTER)
        )
    return True


def _defined_term(question: str) -> set[str] | None:
    m = DEFINITION_QUESTION.match(LEADING_PLACEHOLDERS.sub("", question))
    return set(tokens(m.group(1))) or None if m else None


def _is_definition(sentence: str, defined: set[str]) -> bool:
    """ "Флэшбек (flashback – англ.) – возврат..." / "ПТСР – это расстройство...": the term, then a dash."""
    head = DEFINITION_HEAD.split(sentence, maxsplit=1)
    return len(head) == 2 and len(head[0]) <= DEFINITION_HEAD_CHARS and bool(set(tokens(head[0])) & defined)


def _sentence_score(index: BM25Index, sentence: str, terms: dict[str, float], defined: set[str] | None) -> float:
    score = index.overlap(sentence, terms)
    if defined and score > 0 and _is_definition(sentence, defined):
        score *= DEFINITION_BONUS
    return score


def best_passage(
    index: BM25Index,
    chunk: Chunk,
    terms: dict[str, float],
    audience: Audience,
    cfg: AnswerConfig,
    defined: set[str] | None = None,
    cued: set[str] | None = None,
) -> Passage | None:
    """Best sentence of the chunk extended by adjacent relevant sentences (contiguous, within limits)."""
    sents = [s if _allowed(s, audience) else None for s in chunk_sentences(chunk, clinical=audience != "patient")]
    scores = [_sentence_score(index, s, terms, defined) if s else 0.0 for s in sents]
    if not scores or max(scores) <= 0:
        return None
    best = max(range(len(sents)), key=lambda i: (scores[i], -i))
    head = sents[best]
    assert head is not None  # score > 0 implies an allowed sentence
    lo = hi = best
    size = len(head)

    def fits(j: int) -> bool:
        if not 0 <= j < len(sents):
            return False
        s = sents[j]
        return s is not None and size + len(s) + 1 <= cfg.max_chars

    # A short lead-in (question-like heading, "Комментарии:"-style intro) needs the sentence after it.
    def length(j: int) -> int:
        return len(sents[j] or "") + 1

    if (size < LEAD_IN_CHARS or head.endswith(("?", ":"))) and fits(best + 1):
        hi += 1
        size += length(hi)
    if ANAPHORA.match(head) and fits(best - 1):
        lo -= 1
        size += length(lo)
    floor = cfg.neighbor_ratio * scores[best]
    while hi - lo + 1 < cfg.max_sentences:
        nxt = next((j for j in (hi + 1, lo - 1) if fits(j) and scores[j] > 0 and scores[j] >= floor), None)
        if nxt is None:
            break
        lo, hi = min(lo, nxt), max(hi, nxt)
        size += length(nxt)
    text = " ".join(s for s in sents[lo : hi + 1] if s)
    definition = scores[best] > index.overlap(head, terms)  # got DEFINITION_BONUS
    return Passage(chunk, text, index.overlap(text, terms) * index.boost(chunk, audience, cued), definition)


def select(index: BM25Index, question: str, audience: Audience, docs: set[str] | None = None) -> list[Passage]:
    """Best passage, plus a second one from another chunk if it scores within `second_passage_ratio`."""
    cfg = load_config().answer
    hits = index.search(question, audience=audience, docs=docs, top_k=cfg.max_chunks)
    terms = index.query_terms(question)
    defined = _defined_term(question)
    cued = index.cued_docs(question)
    passages = [p for h in hits if (p := best_passage(index, h.chunk, terms, audience, cfg, defined, cued))]
    if not passages:
        return []
    # A definition answers "что такое X?" better than any passage that merely mentions X often.
    passages.sort(key=lambda p: (not p.definition, -p.score))  # stable: retrieval rank breaks ties
    top = passages[0]
    chosen = [top]
    for p in passages[1:]:
        if len(chosen) >= cfg.max_passages:
            break
        if p.score >= cfg.second_passage_ratio * top.score and p.text != top.text:
            chosen.append(p)
    return chosen


def _citation(chunk: Chunk) -> Citation:
    return Citation(
        doc=chunk.doc.id, short=chunk.doc.short, title=chunk.doc.title, section=chunk.section, pages=chunk.pages
    )


async def _rephrase(excerpt: str, question: str) -> tuple[str, bool]:
    if len(excerpt) > MAX_REPLY_CHARS or any(p.search(excerpt) for p in FORBIDDEN_OUTPUT):
        return excerpt, False  # the gateway would reject any reply this long / with these words anyway
    text = await get_gateway().phrase(PHRASE_INTENT, {"excerpt": excerpt, "question": question})
    if " ".join(text.split()) == " ".join(excerpt.split()):
        return excerpt, False
    if adds_content(text, excerpt, question):
        log.warning("rag rephrasing rejected: adds content")
        return excerpt, False
    return text, True


async def answer(
    question: ScrubbedText,
    *,
    audience: Audience = "patient",
    docs: set[str] | None = None,
    index: BM25Index | None = None,
) -> RagAnswer | None:
    """Answer strictly from the knowledge base, with citations; `None` -> caller sends `fallback_text()`."""
    if audience == "patient" and DOSE_QUESTION.search(question.lower()):
        log.info("rag: dose question redirected to doctor")
        return None
    index = index or get_index()
    passages = select(index, question, audience, docs)
    if not passages:
        log.info("rag: no relevant fragment audience=%s", audience)
        return None
    coverage = index.coverage(" ".join(p.text for p in passages), question)
    if coverage < load_config().answer.min_coverage:
        log.info("rag: low coverage audience=%s coverage=%.2f", audience, coverage)
        return None
    citations = [_citation(p.chunk) for p in passages]
    texts = load_rag_texts()[audience]

    if audience == "patient":
        excerpt = " ".join(p.text for p in passages)
        body, by_llm = await _rephrase(excerpt, question)
        sources = "; ".join(dict.fromkeys(c.label for c in citations))
        text = f"{texts.intro} {body}\n\n{texts.sources_label}: {sources}."
    else:  # doctor/team: verbatim quotes, each with its own citation
        excerpt = "\n".join(f"{p.text} [{c.label}]" for p, c in zip(passages, citations, strict=True))
        by_llm = False
        text = f"{texts.intro}\n{excerpt}"
    if texts.doctor_note:
        text += f"\n\n{texts.doctor_note}"
    log.info("rag: answered audience=%s passages=%d coverage=%.2f llm=%s", audience, len(passages), coverage, by_llm)
    return RagAnswer(
        text=text, excerpt=excerpt, citations=citations, audience=audience, coverage=coverage, phrased_by_llm=by_llm
    )