PTSRaiAgent/app/llm/rag/index.py

"""Lexical retrieval over the corpus: Okapi BM25 on Russian Snowball stems + synonym expansion.

Offline and deterministic (ADR-004): no embeddings, no network. Stop words are dropped before stemming;
the section title is indexed `title_weight` times so a match with "3.2 Психофармакотерапия" counts more
than a passing mention in the body.

Re-ranking: the best `rerank_candidates` chunks by BM25 get a bonus for their densest passage — a window
of `window_sentences` consecutive sentences scored by the idf of query terms it contains. A chunk that
answers in one place beats a long one that mentions the words in passing; the passage itself is returned
with the hit and is what `answer()` quotes from.
"""

import math
import re
from collections import Counter
from dataclasses import dataclass
from functools import lru_cache

from app.llm.rag.corpus import Audience, Chunk, RagConfig, load_chunks, load_config, split_sentences
from app.llm.rag.stemmer import stem

_WORD = re.compile(r"[а-яёa-z0-9]+(?:-[а-яёa-z0-9]+)*", re.IGNORECASE)

# Function words and generic KR boilerplate that carry no topical signal.
STOP_WORDS = frozenset(
    """
    а без более бы был была были было быть в вам вас ваш весь во вот все всего всех вы где да даже для до его
    ее если есть еще же за и из или им их к как ко когда кто ли либо мне может можно мой мы на над не него нее
    нет ни них но ну о об однако он она они оно от по под при про с со так также такой там те то тоже только
    том тот у уже хотя чего чем что чтобы эта эти это этот я ли ней нем
    какой какая какое какие каков такое такая такие сколько почему зачем меня мой моя мое мои себя сегодня будет
    """.split()  # noqa: SIM905 — a word list reads better as text
)


def tokens(text: str) -> list[str]:
    """Stems of content words (hyphenated words are kept whole and also split: "когнитивно-поведенческая")."""
    out: list[str] = []
    for word in _WORD.findall(text.lower().replace("ё", "е")):
        parts = [p for p in word.split("-") if p not in STOP_WORDS and (len(p) > 1 or p.isdigit())]
        if "-" in word and len(parts) > 1:
            out.append("-".join(stem(p) for p in parts))  # the compound itself, in any case form
        out.extend(stem(p) for p in parts)
    return out


EXPANSION_WEIGHT = 0.7
SYNONYM_CREDIT = 0.8
PREFIX_LEN = 6


@dataclass(frozen=True)
class Hit:
    chunk: Chunk
    score: float  # BM25 + passage bonus (ranking key)
    bm25: float
    passage: str  # densest window of the chunk for the query


def sentences(text: str) -> list[str]:
    """Sentences of a chunk; every line (table row, list item) is split on its own."""
    return [s for line in text.splitlines() for s in split_sentences(line)]


class SynonymExpander:
    """Adds the stems of every phrase in a synonym group when any phrase of the group is in the query.

    Added stems weigh `EXPANSION_WEIGHT`: they widen recall but should not outrank the patient's own words.
    """

    def __init__(self, groups: list[list[str]]) -> None:
        self._groups = [[tokens(phrase) for phrase in group] for group in groups]
        # stem -> stems of the other phrases of its groups ("отменя" -> {"отмен", "брос", ...})
        self.related: dict[str, set[str]] = {}
        for group in self._groups:
            stems = {t for phrase in group for t in phrase}
            for t in stems:
                self.related.setdefault(t, set()).update(stems - {t})

    def expand(self, query_tokens: list[str]) -> dict[str, float]:
        weights: dict[str, float] = {}
        for t in query_tokens:
            weights[t] = weights.get(t, 0.0) + 1.0
        present = set(query_tokens)
        for group in self._groups:
            if any(phrase and set(phrase) <= present for phrase in group):
                for phrase in group:
                    for t in phrase:
                        if t not in present:  # a 3-word expansion of "птср" is one concept, not three
                            weights[t] = max(weights.get(t, 0.0), EXPANSION_WEIGHT / len(phrase))
        return weights


class BM25Index:
    def __init__(self, chunks: tuple[Chunk, ...], cfg: RagConfig) -> None:
        r = cfg.retrieval
        self._k1, self._b, self._min_score, self._top_k = r.bm25_k1, r.bm25_b, r.min_score, r.top_k
        self._candidates, self._window, self._window_weight = r.rerank_candidates, r.window_sentences, r.window_weight
        self._section_boost = r.section_boost
        self.chunks = chunks
        self._sentences = [[(s, set(tokens(s))) for s in sentences(c.text)] for c in chunks]
        self._expander = SynonymExpander(cfg.synonyms)
        self._tf: list[Counter[str]] = []
        df: Counter[str] = Counter()
        for c in chunks:
            tf = Counter(tokens(c.text) + tokens(c.section) * r.title_weight)
            self._tf.append(tf)
            df.update(tf.keys())
        self._len = [sum(tf.values()) for tf in self._tf]
        self._avg_len = sum(self._len) / max(len(self._len), 1)
        n = len(chunks)
        self._idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()}

    def query_terms(self, query: str) -> dict[str, float]:
        """Query stems -> weight (repeated words count more, synonym expansions less)."""
        return self._expander.expand(tokens(query))

    def _score(self, i: int, terms: dict[str, float]) -> float:
        tf, norm = self._tf[i], self._k1 * (1 - self._b + self._b * self._len[i] / self._avg_len)
        score = 0.0
        for term, weight in terms.items():
            f = tf.get(term)
            if f:
                score += self._idf[term] * f * (self._k1 + 1) / (f + norm) * weight
        return score

    def boost(self, chunk: Chunk, audience: Audience | None) -> float:
        """Section multiplier for the audience (`retrieval.section_boost`), 1.0 by default."""
        boosts = self._section_boost.get(audience or "", {})
        return next((b for prefix, b in boosts.items() if chunk.section.startswith(prefix)), 1.0)

    def overlap(self, text: str, terms: dict[str, float]) -> float:
        """Sum of idf × weight of the query terms present in `text` (sentence relevance for `answer()`)."""
        present = set(tokens(text))
        return sum(self._idf.get(t, 0.0) * w for t, w in terms.items() if t in present)

    def coverage(self, text: str, query: str) -> float:
        """Share of the query's own (not synonym-expanded) idf mass found in `text`, 0..1.

        A query word whose synonym — or a word with the same `PREFIX_LEN`-letter prefix, for stems Snowball
        leaves apart ("скрининговый"/"скрининг") — is in the text counts as `SYNONYM_CREDIT` of itself.
        A word absent from the corpus counts with twice the maximal idf — the strongest sign of an off-topic
        question: "как установить windows" is then mostly uncovered and `answer()` says there is
        no answer instead of quoting the nearest paragraph about "установить диагноз".
        """
        own = set(tokens(query))
        if not own:
            return 0.0
        unseen = 2 * max(self._idf.values(), default=1.0)
        total = sum(self._idf.get(t, unseen) for t in own)
        present = set(tokens(text))
        prefixes = {p[:PREFIX_LEN] for p in present if len(p) >= PREFIX_LEN}
        covered = 0.0
        for t in own:
            if t in present:
                covered += self._idf[t]
            elif self._expander.related.get(t, set()) & present or (
                len(t) >= PREFIX_LEN and t[:PREFIX_LEN] in prefixes
            ):
                covered += SYNONYM_CREDIT * self._idf.get(t, unseen)  # said in the corpus' own words
        return min(1.0, covered / total)

    def _passage(self, i: int, terms: dict[str, float]) -> tuple[float, str]:
        sents = self._sentences[i]
        best = (0.0, sents[0][0] if sents else "")
        for start in range(len(sents)):
            window = sents[start : start + self._window]
            present = set().union(*(toks for _, toks in window))
            score = sum(self._idf.get(t, 0.0) * w for t, w in terms.items() if t in present)
            if score > best[0]:
                best = (score, " ".join(text for text, _ in window))
        return best

    def search(
        self,
        query: str,
        *,
        audience: Audience | None = None,
        docs: set[str] | None = None,
        top_k: int | None = None,
        min_score: float | None = None,
    ) -> list[Hit]:
        """Best chunks for `query`, filtered by audience/doc ids; empty if nothing passes `min_score` (BM25)."""
        terms = self.query_terms(query)
        if not terms:
            return []
        threshold = self._min_score if min_score is None else min_score
        scored = [
            (s, i)
            for i, c in enumerate(self.chunks)
            if (audience is None or audience in c.doc.audiences)
            and (docs is None or c.doc.id in docs)
            and (s := self._score(i, terms)) >= threshold
        ]
        scored.sort(key=lambda x: (-x[0], x[1]))
        hits = []
        for bm25, i in scored[: max(self._candidates, top_k or 0)]:
            bonus, passage = self._passage(i, terms)
            chunk = self.chunks[i]
            hits.append(Hit(chunk, (bm25 + self._window_weight * bonus) * self.boost(chunk, audience), bm25, passage))
        hits.sort(key=lambda h: -h.score)  # stable: BM25 order breaks ties
        return hits[: top_k or self._top_k]


@lru_cache
def get_index() -> BM25Index:
    return BM25Index(load_chunks(), load_config())


def retrieve(query: str, **kwargs) -> list[Hit]:
    """Module-level shortcut: `retrieve("когда начинает действовать антидепрессант", audience="patient")`."""
    return get_index().search(query, **kwargs)