ElebobaandClaude Opus 5.5
6a77831603
CI / ruff · pytest (pull_request) Successful in 1m30s
- docs/source/rag: PDF sources (KR PTSD 2023, KR depression 2024, meds module spec), pdf->txt converter (tables as markdown, page markers, hand-transcribed figures), word-level check against pdftotext; fixes HAM-D/MADRS fake headers and split rows - config/rag: corpus txt + sources.yaml (documents, audiences, retrieval params, synonyms) - app/llm/rag: Snowball stemmer, section-aware chunker with page citations, BM25 + synonym expansion + densest-passage reranking, extractive answer() with coverage gate, patient dose/prescribing filter, optional LLM rephrasing rejected when it adds content; CLI python -m app.llm.rag - tests: gold retrieval set (hit@3 >= 15/18), off-topic, safety, faithfulness - ADR-011 Refs #60 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
310 lines
15 KiB
Python
310 lines
15 KiB
Python
"""RAG knowledge base (#60): corpus integrity, retrieval quality, grounded and safe answers."""
|
||
|
||
import re
|
||
|
||
import pytest
|
||
|
||
from app.core.types import ScrubbedText
|
||
from app.llm.gateway import LLMGateway
|
||
from app.llm.mock import MockTasks
|
||
from app.llm.rag import qa as answer_module
|
||
from app.llm.rag.corpus import heading, load_chunks, load_config, rag_dir, split_sentences
|
||
from app.llm.rag.index import get_index, tokens
|
||
from app.llm.rag.qa import LIT_REFS, PRESCRIBING, QUOTE_FILTER, answer, fallback_text, select
|
||
from app.llm.rag.stemmer import stem
|
||
|
||
S = ScrubbedText
|
||
|
||
# question -> (document, pattern that must occur in one of the top-3 chunks). Answers checked by hand
|
||
# against the PDFs; the set doubles as the tuning set, so the threshold below guards against regressions.
|
||
GOLD = [
|
||
("через сколько начинают действовать антидепрессанты", "kr_depression", r"Тимоаналептический эффект обычно"),
|
||
("когда антидепрессант даст эффект", "kr_depression", r"Тимоаналептический эффект обычно"),
|
||
("помогает ли ДПДГ при ПТСР", "kr_ptsd", r"ДПДГ"),
|
||
("какая психотерапия эффективна при посттравматическом расстройстве", "kr_ptsd", r"наибольшей доказательной"),
|
||
("как правильно отменять антидепрессанты", "kr_depression", r"постепенное снижение дозы"),
|
||
("что такое флэшбек", "kr_ptsd", r"Флэшбек \(flashback"),
|
||
("шкала Гамильтона для депрессии", "kr_depression", r"Шкала Гамильтона для оценки депрессии"),
|
||
("помощь военнослужащим после боевой психической травмы", "kr_ptsd", r"военнослужащ"),
|
||
("колумбийская шкала суицидального риска", "kr_depression", r"Колумбийская шкала"),
|
||
("вызывают ли бензодиазепины зависимость", "kr_ptsd", r"формировать зависимость"),
|
||
("что сказать пациенту о ПТСР", "kr_ptsd", r"Что такое посттравматическое"),
|
||
("электросудорожная терапия при резистентной депрессии", "kr_depression", r"лектросудорожн"),
|
||
("скрининговый опросник ПТСР", "kr_ptsd", r"Trauma Screening"),
|
||
("физическая активность при депрессии", "kr_depression", r"физическ"),
|
||
("сертралин при ПТСР", "kr_ptsd", r"ертралин"),
|
||
("противорецидивная терапия депрессии сколько длится", "kr_depression", r"ротиворецидивн"),
|
||
("правило M7 срочные симптомы", "meds_module_spec", r"M7"),
|
||
("через сколько оценивают ответ на лечение", "meds_module_spec", r"4–6 неделе"),
|
||
]
|
||
GOLD_MIN_HITS = 15
|
||
|
||
OFF_TOPIC = [
|
||
"какой рецепт борща",
|
||
"как починить машину",
|
||
"расписание электричек",
|
||
"кто выиграл чемпионат по футболу",
|
||
"курс доллара сегодня",
|
||
"погода в москве завтра",
|
||
"как установить windows",
|
||
"лучший фильм года",
|
||
"как выучить английский",
|
||
]
|
||
|
||
PATIENT_QUESTIONS = [
|
||
"через сколько начинают действовать антидепрессанты",
|
||
"что такое ПТСР",
|
||
"что такое флэшбек",
|
||
"помогает ли психотерапия при ПТСР",
|
||
"как правильно отменять антидепрессанты",
|
||
"что будет если пить таблетки с алкоголем",
|
||
"сертралин при ПТСР",
|
||
"вызывают ли бензодиазепины зависимость",
|
||
]
|
||
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def mock_gateway(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(MockTasks()))
|
||
|
||
|
||
class FixedPhraseTasks(MockTasks):
|
||
"""An LLM that returns a fixed rephrasing."""
|
||
|
||
name = "fake-llm"
|
||
|
||
def __init__(self, reply: str) -> None:
|
||
self.reply = reply
|
||
|
||
async def phrase(self, intent: str, payload: ScrubbedText, base_text: str) -> str | None:
|
||
return self.reply
|
||
|
||
|
||
def _words(text: str) -> str:
|
||
return " ".join(re.findall(r"\w+", LIT_REFS.sub("", text).lower()))
|
||
|
||
|
||
# --- stemmer ------------------------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("word", "expected"),
|
||
[
|
||
# reference outputs of the Snowball Russian stemmer
|
||
("вагоне", "вагон"),
|
||
("важнейшими", "важн"),
|
||
("гостиница", "гостиниц"),
|
||
("сделавшись", "сдела"),
|
||
("посттравматическое", "посттравматическ"),
|
||
("антидепрессантов", "антидепрессант"),
|
||
("пробуждения", "пробужден"),
|
||
("ёлками", "елк"),
|
||
("SSRI", "ssri"),
|
||
],
|
||
)
|
||
def test_stemmer(word: str, expected: str) -> None:
|
||
assert stem(word) == expected
|
||
|
||
|
||
def test_tokens_drop_stop_words_and_split_hyphens() -> None:
|
||
assert tokens("что такое когнитивно-поведенческая терапия") == [
|
||
"когнитивн-поведенческ",
|
||
"когнитивн",
|
||
"поведенческ",
|
||
"терап",
|
||
]
|
||
|
||
|
||
# --- corpus -------------------------------------------------------------------------------------------
|
||
|
||
|
||
def _page_of(text: str, line: str) -> int | None:
|
||
"""Page number of the first line equal to `line`."""
|
||
page = None
|
||
for ln in text.splitlines():
|
||
if m := re.match(r"^--- стр\. (\d+) ---$", ln):
|
||
page = int(m.group(1))
|
||
elif ln == line:
|
||
return page
|
||
return None
|
||
|
||
|
||
def test_every_page_is_chunked_except_bibliography() -> None:
|
||
chunks = load_chunks()
|
||
for doc in load_config().documents:
|
||
text = (rag_dir() / doc.file).read_text(encoding="utf-8")
|
||
n_pages = len(re.findall(r"^--- стр\. \d+ ---$", text, re.M))
|
||
own = [c for c in chunks if c.doc.id == doc.id]
|
||
assert own, doc.id
|
||
assert all(1 <= c.page_start <= c.page_end <= n_pages for c in own)
|
||
missing = set(range(1, n_pages + 1)) - {p for c in own for p in range(c.page_start, c.page_end + 1)}
|
||
lit, appendix = _page_of(text, "Список литературы"), _page_of(text, "Приложение А1")
|
||
allowed = set(range(lit + 1, appendix)) if lit and appendix else set()
|
||
assert missing <= allowed, (doc.id, sorted(missing - allowed))
|
||
|
||
|
||
def test_chunks_are_bounded_and_ids_unique() -> None:
|
||
chunks = load_chunks()
|
||
assert len({c.id for c in chunks}) == len(chunks)
|
||
max_chars = load_config().chunking.max_chars
|
||
# only a single oversized table row may exceed the limit; prose is split by sentences
|
||
for c in chunks:
|
||
if len(c.text) > 2 * max_chars:
|
||
assert any(len(line) > max_chars for line in c.text.splitlines() if line.startswith("|")), c.id
|
||
|
||
|
||
def test_bibliography_and_abbreviations_not_indexed() -> None:
|
||
sections = {c.section for c in load_chunks()}
|
||
assert not any(s.startswith(("Список литературы", "Список сокращений")) for s in sections)
|
||
assert "3.2. Психофармакотерапия" in sections
|
||
assert "3.6 Противорецидивная (профилактическая) терапия" in sections
|
||
|
||
|
||
def test_table_pieces_repeat_header() -> None:
|
||
for c in load_chunks():
|
||
lines = c.text.splitlines()
|
||
for i, line in enumerate(lines):
|
||
if line.startswith("|") and (i == 0 or not lines[i - 1].startswith("|")):
|
||
assert i + 1 < len(lines) and re.match(r"^\|(-+\|)+$", lines[i + 1]), c.id
|
||
|
||
|
||
def test_hamd_items_not_duplicated_or_split() -> None:
|
||
"""Converter regression: item 1 of HAM-D was repeated as a fake header on every page, items 2/17 split."""
|
||
text = (rag_dir() / "corpus" / "kr_depression.txt").read_text(encoding="utf-8")
|
||
assert text.count("| 1. СНИЖЕННОЕ НАСТРОЕНИЕ") == 1
|
||
assert re.search(r"^\| 2\. ЧУВСТВО ВИНЫ 0 □", text, re.M)
|
||
assert re.search(r"^\| 17\. КРИТИКА 0 □", text, re.M)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("line", "is_heading"),
|
||
[
|
||
("2.1 Жалобы и анамнез", True),
|
||
("3.1.2. Купирующая терапия ДЭ средней степени", True),
|
||
("3. Лечение, включая медикаментозную и немедикаментозную терапии", True),
|
||
("Приложение Г5", True),
|
||
("2. Стойкие симптомы повышения психологической чувствительности", False), # list item
|
||
("28.07.2022 N 69435).", False), # date
|
||
("| 2.1 | таблица |", False),
|
||
],
|
||
)
|
||
def test_heading(line: str, is_heading: bool) -> None:
|
||
assert (heading(line) is not None) == is_heading
|
||
|
||
|
||
def test_split_sentences_keeps_abbreviations() -> None:
|
||
assert split_sentences("Алгоритм ведения (см. Приложение Б5). Далее текст.") == [
|
||
"Алгоритм ведения (см. Приложение Б5).",
|
||
"Далее текст.",
|
||
]
|
||
|
||
|
||
# --- retrieval ----------------------------------------------------------------------------------------
|
||
|
||
|
||
def test_gold_questions_hit_at_3() -> None:
|
||
index = get_index()
|
||
misses = [
|
||
q
|
||
for q, doc, pattern in GOLD
|
||
if not any(
|
||
h.chunk.doc.id == doc and re.search(pattern, h.chunk.text) for h in index.search(q, top_k=3, min_score=0)
|
||
)
|
||
]
|
||
assert len(GOLD) - len(misses) >= GOLD_MIN_HITS, misses
|
||
|
||
|
||
def test_synonyms_expand_abbreviations() -> None:
|
||
terms = get_index().query_terms("ПТСР")
|
||
assert terms["птср"] == 1.0
|
||
assert 0 < terms["посттравматическ"] < 1.0
|
||
|
||
|
||
def test_team_documents_never_reach_patients_or_doctors() -> None:
|
||
for audience in ("patient", "doctor"):
|
||
hits = get_index().search("правила M1–M10 модуль лекарства", audience=audience, min_score=0, top_k=20)
|
||
assert all("team" not in h.chunk.doc.audiences for h in hits)
|
||
team = get_index().search("правило M7 срочные симптомы", audience="team")
|
||
assert team and team[0].chunk.doc.id == "meds_module_spec"
|
||
|
||
|
||
def test_patient_information_appendix_boosted() -> None:
|
||
hit = get_index().search("что такое ПТСР", audience="patient")[0]
|
||
assert hit.chunk.section.startswith("Приложение В")
|
||
|
||
|
||
# --- answers ------------------------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize("question", OFF_TOPIC)
|
||
async def test_off_topic_has_no_answer(question: str) -> None:
|
||
assert await answer(S(question), audience="patient") is None
|
||
assert await answer(S(question), audience="doctor") is None
|
||
|
||
|
||
@pytest.mark.parametrize("question", ["можно увеличить дозу сертралина?", "можно пить полтаблетки?"])
|
||
async def test_dose_question_redirected(question: str) -> None:
|
||
assert await answer(S(question), audience="patient") is None
|
||
assert "врач" in fallback_text("patient")
|
||
|
||
|
||
@pytest.mark.parametrize("question", PATIENT_QUESTIONS)
|
||
async def test_patient_answer_is_extractive_and_safe(question: str) -> None:
|
||
reply = await answer(S(question), audience="patient")
|
||
if reply is None: # "ask your doctor" is always a valid outcome
|
||
return
|
||
sources = {c.id: _words(c.text) for c in load_chunks()}
|
||
for sentence in split_sentences(reply.excerpt):
|
||
assert any(_words(sentence) in text for text in sources.values()), sentence # quoted, not generated
|
||
assert not PRESCRIBING.search(reply.excerpt)
|
||
assert not any(p.search(reply.excerpt) for p in QUOTE_FILTER)
|
||
assert "Источник:" in reply.text and "стр." in reply.text
|
||
assert reply.text.rstrip().endswith(answer_module.load_rag_texts()["patient"].doctor_note)
|
||
assert all(c.doc != "meds_module_spec" for c in reply.citations)
|
||
|
||
|
||
async def test_onset_answer_content() -> None:
|
||
reply = await answer(S("через сколько начинают действовать антидепрессанты"), audience="patient")
|
||
assert reply is not None
|
||
assert re.search(r"недел", reply.excerpt)
|
||
assert reply.coverage >= load_config().answer.min_coverage
|
||
|
||
|
||
async def test_definition_question_prefers_definition() -> None:
|
||
reply = await answer(S("что такое флэшбек"), audience="patient")
|
||
assert reply is not None and reply.excerpt.startswith("Флэшбек (flashback")
|
||
assert reply.citations[0].pages == "2"
|
||
|
||
|
||
async def test_doctor_answer_verbatim_with_citations() -> None:
|
||
reply = await answer(S("синдром отмены антидепрессантов"), audience="doctor")
|
||
assert reply is not None and not reply.phrased_by_llm
|
||
assert re.search(r"\[КР Депрессия 2024, стр\. \d+(–\d+)?\]", reply.text)
|
||
assert "СОПА" in reply.excerpt
|
||
|
||
|
||
async def test_doctor_sees_dosing_patient_does_not() -> None:
|
||
q = S("сертралин при ПТСР")
|
||
doctor = await answer(q, audience="doctor")
|
||
assert doctor is not None and re.search(r"сертралин\*\* 25-200 мг", doctor.excerpt)
|
||
patient = await answer(q, audience="patient")
|
||
assert patient is None or "мг" not in patient.excerpt
|
||
|
||
|
||
async def test_llm_rephrasing_that_adds_numbers_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(FixedPhraseTasks("Эффект будет через 2 дня.")))
|
||
reply = await answer(S("через сколько начинают действовать антидепрессанты"), audience="patient")
|
||
assert reply is not None and not reply.phrased_by_llm
|
||
assert "2 дня" not in reply.text
|
||
|
||
|
||
async def test_llm_rephrasing_kept_when_faithful(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
q = "помогает ли психотерапия при ПТСР"
|
||
excerpt = " ".join(p.text for p in select(get_index(), q, "patient"))
|
||
assert "Это помогает" in excerpt # anaphora pulled in the sentence it refers to
|
||
simplified = excerpt.replace("Одновременно с этим, психотерапия", "Кроме того, психотерапия")
|
||
monkeypatch.setattr(answer_module, "get_gateway", lambda: LLMGateway(FixedPhraseTasks(simplified)))
|
||
reply = await answer(S(q), audience="patient")
|
||
assert reply is not None and reply.phrased_by_llm
|
||
assert simplified in reply.text and reply.excerpt == excerpt # the verbatim quote is kept for audit
|