PTSRaiAgent/app/llm/rag/corpus.py

"""Knowledge-base corpus: `config/rag/sources.yaml` + txt files -> citable chunks.

The txt files come from `docs/source/rag/pdf_to_txt.py`: paragraphs are separated by blank lines, pages by
`--- стр. N ---` markers, tables are markdown, text transcribed from figures follows an
`[Изображение: ...]` line. Chunking rules:

* a chunk never crosses a section boundary (numbered headings "2.1 ...", "Приложение Г5", ...), so its
  section title is reliable context; the title is stored separately and indexed with extra weight;
* "Комментарии: ..." after a recommendation stays with the recommendation (the comment explains it);
* a line starting in lowercase continues the previous paragraph (cut by a page break the converter missed);
* tables are split by rows, every table piece repeats the table header, so a chunk is readable alone;
* each chunk keeps the page range it came from, for the citation "[КР ПТСР 2023, стр. 45]".
"""

import re
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path

import yaml
from pydantic import BaseModel, Field

from app.config import get_settings

Audience = str  # "doctor" | "patient" | "team"

PAGE_MARKER = re.compile(r"^--- стр\. (\d+) ---$")
TABLE_DIVIDER = re.compile(r"^\|(-+\|)+$")
# "2.1 Жалобы и анамнез", "3.1.2. Купирующая терапия ..." — a numbered subsection title.
SUBSECTION = re.compile(r"^(\d{1,2}(?:\.\d{1,2})+)\.?\s+([А-ЯЁA-Z].*)$")  # not dates "28.07.2022 N ..."
# Top-level KR sections: "1. Краткая информация ...", "3. Лечение, включая ...". Plain "1. ..." lines are
# list items, so a top-level heading must start with one of the standard KR section names.
TOPLEVEL = re.compile(
    r"^([1-9])\.\s+(Краткая информация|Диагностика|Лечение|Медицинская реабилитация|Профилактика"
    r"|Организация оказания|Дополнительная информация)"
)
NAMED = re.compile(
    r"^(Приложение [А-ЯЁ]\d*(?:-[А-ЯЁ]+\d*)?\b.*|Критерии оценки качества медицинской помощи|Список литературы"
    r"|Список сокращений|Термины и определения)$"
)
# Unnumbered short title lines (team specs): capitalised, short, no final punctuation, not a list item.
SHORT_TITLE = re.compile(r"^[А-ЯЁ][^.:;!?|]{2,60}$")
COMMENT = re.compile(r"^Комментари[йи]\s*:")
SENTENCE_SPLIT = re.compile(r"(?<=[.!?…])\s+(?=[А-ЯЁA-Z0-9«\"(])")
# Abbreviations after which a capital letter does not start a new sentence: "(см. Приложение Б5)".
ABBREVIATION_END = re.compile(
    r"(?:\b(?:см|т\.\s?е|т\.\s?д|т\.\s?п|др|напр|рис|табл|стр|г|гг|им|ср|англ|лат|греч)|\b[А-ЯЁA-Z])\.$"
)
MAX_TITLE_CHARS = 300


class DocumentMeta(BaseModel):
    id: str
    file: str
    title: str
    short: str
    issuer: str
    year: int
    icd: list[str] = Field(default_factory=list)
    audiences: list[Audience]


class ChunkingConfig(BaseModel):
    max_chars: int = 1200
    min_chars: int = 200
    skip_sections: list[str] = Field(default_factory=list)


class RetrievalConfig(BaseModel):
    top_k: int = 5
    bm25_k1: float = 1.5
    bm25_b: float = 0.75
    title_weight: int = 2
    min_score: float = 3.0
    rerank_candidates: int = 30
    window_sentences: int = 2
    window_weight: float = 0.5
    # section title prefix -> score multiplier, per audience ({"patient": {"Приложение В": 1.5}})
    section_boost: dict[str, dict[str, float]] = Field(default_factory=dict)


class AnswerConfig(BaseModel):
    max_chunks: int = 3
    max_passages: int = 2
    max_sentences: int = 3
    max_chars: int = 600
    min_coverage: float = 0.5
    second_passage_ratio: float = 0.8
    neighbor_ratio: float = 0.5


class RagConfig(BaseModel):
    documents: list[DocumentMeta]
    chunking: ChunkingConfig = ChunkingConfig()
    retrieval: RetrievalConfig = RetrievalConfig()
    answer: AnswerConfig = AnswerConfig()
    synonyms: list[list[str]] = Field(default_factory=list)


@dataclass(frozen=True)
class Chunk:
    id: str  # "<doc>#<n>"
    doc: DocumentMeta
    section: str  # nearest heading ("" before the first one)
    page_start: int
    page_end: int
    text: str  # display text: paragraphs and markdown table rows

    @property
    def pages(self) -> str:
        return str(self.page_start) if self.page_start == self.page_end else f"{self.page_start}–{self.page_end}"

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


@dataclass
class _Unit:
    """Indivisible piece of a chunk: a paragraph (with its comment) or one table row."""

    text: str
    page: int
    table_header: str | None = None  # "| a | b |\n|---|---|" for table rows


@dataclass
class _Section:
    title: str
    units: list[_Unit] = field(default_factory=list)


def rag_dir() -> Path:
    return get_settings().config_dir / "rag"


@lru_cache
def load_config(path: Path | None = None) -> RagConfig:
    path = path or rag_dir() / "sources.yaml"
    return RagConfig.model_validate(yaml.safe_load(path.read_text(encoding="utf-8")))


def heading(line: str) -> str | None:
    """The line's text if it is a section heading, else `None`."""
    if len(line) > MAX_TITLE_CHARS or line.startswith("|"):
        return None
    if SUBSECTION.match(line) or TOPLEVEL.match(line) or NAMED.match(line):
        return line
    return None


def _is_short_title(block: list[str], next_block: list[str] | None) -> bool:
    """A lone short line followed by a body paragraph (unnumbered headings in team docs)."""
    return (
        len(block) == 1
        and bool(SHORT_TITLE.match(block[0]))
        and next_block is not None
        and not PAGE_MARKER.match(next_block[0])
        and (len(next_block) > 1 or len(next_block[0]) > 80 or next_block[0].startswith("|"))
    )


def _blocks(text: str) -> list[list[str]]:
    """Blank-line separated blocks of non-empty lines."""
    return [[ln for ln in b.splitlines() if ln.strip()] for b in re.split(r"\n\s*\n", text) if b.strip()]


def _sections(text: str, skip: list[str]) -> list[_Section]:
    blocks = _blocks(text)
    sections = [_Section(title="")]
    page = 1
    skipping = False
    for i, block in enumerate(blocks):
        first = block[0]
        if first.startswith("# ") or first.startswith("Источник:"):
            continue  # file header written by the converter
        m = PAGE_MARKER.match(first)
        if m:
            page = int(m.group(1))
            continue
        nxt = next((b for b in blocks[i + 1 :] if not PAGE_MARKER.match(b[0])), None)
        if first.startswith("|"):
            if not skipping:
                sections[-1].units.extend(_table_units(block, page))
            continue
        if first.startswith("[Изображение"):
            if not skipping:  # transcribed figure: one unit, keeps its own line breaks
                sections[-1].units.append(_Unit("\n".join(block[1:]) if len(block) > 1 else first, page))
            continue
        for j, line in enumerate(block):
            title = heading(line)
            if title is None and j == 0 and _is_short_title(block, nxt):
                title = line
            if title is not None:
                skipping = any(title.startswith(s) for s in skip)
                sections.append(_Section(title=title))
                continue
            if skipping:
                continue
            units = sections[-1].units
            if COMMENT.match(line) and units and units[-1].table_header is None:
                units[-1].text += "\n" + line
            elif line[0].islower() and units and units[-1].table_header is None:
                units[-1].text += " " + line  # paragraph continued after a page break or a figure
            else:
                units.append(_Unit(line, page))
    return [s for s in sections if s.units]


def _table_units(block: list[str], page: int) -> list[_Unit]:
    if len(block) >= 2 and TABLE_DIVIDER.match(block[1]):
        header = f"{block[0]}\n{block[1]}"
        rows = block[2:]
    else:  # continuation without a header line (should not happen with the converter, but be safe)
        header, rows = "", block
    return [_Unit(row, page, table_header=header) for row in rows]


def split_sentences(text: str) -> list[str]:
    """Sentences of one line/paragraph; does not split after "см.", "т.е.", initials."""
    out: list[str] = []
    for part in SENTENCE_SPLIT.split(text):
        if out and ABBREVIATION_END.search(out[-1]):
            out[-1] += " " + part
        else:
            out.append(part)
    return [p.strip() for p in out if p.strip()]


def _split_long(unit: _Unit, max_chars: int) -> list[_Unit]:
    if len(unit.text) <= max_chars or unit.table_header is not None:
        return [unit]
    parts: list[str] = []
    for sentence in split_sentences(unit.text):
        if parts and len(parts[-1]) + 1 + len(sentence) <= max_chars:
            parts[-1] += " " + sentence
        else:
            parts.append(sentence)
    return [_Unit(p, unit.page) for p in parts]


def _render(units: list[_Unit]) -> str:
    out: list[str] = []
    header: str | None = None
    for u in units:
        if u.table_header is not None:
            if header != u.table_header:
                header = u.table_header
                if header:
                    out.append(header)
            out.append(u.text)
        else:
            header = None
            out.append(u.text)
    return "\n".join(out)


def _pack(section: _Section, max_chars: int, min_chars: int) -> list[list[_Unit]]:
    units = [p for u in section.units for p in _split_long(u, max_chars)]
    groups: list[list[_Unit]] = [[]]
    size = 0
    for u in units:
        if groups[-1] and size + len(u.text) > max_chars:
            groups.append([])
            size = 0
        groups[-1].append(u)
        size += len(u.text) + 1
    if len(groups) > 1 and sum(len(u.text) for u in groups[-1]) < min_chars:
        groups[-2].extend(groups.pop())
    return [g for g in groups if g]


def chunk_document(doc: DocumentMeta, text: str, cfg: ChunkingConfig) -> list[Chunk]:
    chunks: list[Chunk] = []
    for section in _sections(text, cfg.skip_sections):
        for group in _pack(section, cfg.max_chars, cfg.min_chars):
            chunks.append(
                Chunk(
                    id=f"{doc.id}#{len(chunks)}",
                    doc=doc,
                    section=section.title,
                    page_start=group[0].page,
                    page_end=group[-1].page,
                    text=_render(group),
                )
            )
    return chunks


@lru_cache
def load_chunks(config_path: Path | None = None) -> tuple[Chunk, ...]:
    cfg = load_config(config_path)
    base = (config_path.parent if config_path else rag_dir()).resolve()
    out: list[Chunk] = []
    for doc in cfg.documents:
        out.extend(chunk_document(doc, (base / doc.file).read_text(encoding="utf-8"), cfg.chunking))
    return tuple(out)