pridanie testov supplemental expansion
This commit is contained in:
parent
67d66a9f3e
commit
70e93a7a00
372
test/test_rag_supplemental_expansion.py
Normal file
372
test/test_rag_supplemental_expansion.py
Normal file
@ -0,0 +1,372 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
from scripts.rag_supplemental_expansion import (
|
||||
canonical_student_main_path,
|
||||
expand_results_with_supplemental_sources,
|
||||
is_multi_document_query,
|
||||
)
|
||||
|
||||
|
||||
def create_db(tmp_path: Path) -> Path:
|
||||
db = tmp_path / "supplemental.sqlite"
|
||||
|
||||
with sqlite3.connect(db) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE documents (
|
||||
id INTEGER PRIMARY KEY,
|
||||
path TEXT NOT NULL,
|
||||
title TEXT,
|
||||
author TEXT,
|
||||
published INTEGER
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE chunks (
|
||||
id INTEGER PRIMARY KEY,
|
||||
chunk_id TEXT NOT NULL,
|
||||
document_path TEXT NOT NULL,
|
||||
title TEXT,
|
||||
author TEXT,
|
||||
published INTEGER,
|
||||
chunk_index INTEGER,
|
||||
heading_paths_json TEXT,
|
||||
text TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
return db
|
||||
|
||||
|
||||
def add_document(
|
||||
db: Path,
|
||||
*,
|
||||
path: str,
|
||||
title: str,
|
||||
text: str,
|
||||
chunk_index: int = 0,
|
||||
published: bool = True,
|
||||
) -> None:
|
||||
with sqlite3.connect(db) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO documents(
|
||||
path,
|
||||
title,
|
||||
author,
|
||||
published
|
||||
) VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
path,
|
||||
title,
|
||||
"Daniel Hladek",
|
||||
1 if published else 0,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO chunks(
|
||||
chunk_id,
|
||||
document_path,
|
||||
title,
|
||||
author,
|
||||
published,
|
||||
chunk_index,
|
||||
heading_paths_json,
|
||||
text
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
f"{path}::chunk-{chunk_index}",
|
||||
path,
|
||||
title,
|
||||
"Daniel Hladek",
|
||||
1 if published else 0,
|
||||
chunk_index,
|
||||
"[]",
|
||||
text,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def base_result(path: str, title: str) -> dict:
|
||||
return {
|
||||
"chunk_id": f"{path}::chunk-9",
|
||||
"document_path": path,
|
||||
"title": title,
|
||||
"author": "Daniel Hladek",
|
||||
"published": True,
|
||||
"chunk_index": 9,
|
||||
"heading_paths": [],
|
||||
"text": "menej relevantný úsek",
|
||||
"source_url": "https://example.invalid/base",
|
||||
}
|
||||
|
||||
|
||||
def test_canonical_student_main_path_maps_child_document() -> None:
|
||||
assert canonical_student_main_path(
|
||||
"pages/students/2016/jakub_maruniak/dp2021/README.md"
|
||||
) == "pages/students/2016/jakub_maruniak/README.md"
|
||||
|
||||
|
||||
def test_exact_person_title_anchors_own_student_document(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
path = "pages/students/2018/ondrej_megela/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=path,
|
||||
title="Ondrej Megela",
|
||||
text=(
|
||||
"Ondrej Megela. Neurónové siete a neurónové "
|
||||
"jazykové modelovanie typu BERT."
|
||||
),
|
||||
)
|
||||
wrong_path = "pages/students/2016/patrik_pavlisin/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=wrong_path,
|
||||
title="Patrik Pavlišin",
|
||||
text="Ondrej Megela je spomenutý pri neurónových sieťach.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Ondrej Megela neurónové siete dokument študent",
|
||||
[base_result(wrong_path, "Patrik Pavlišin")],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result[0]["document_path"] == path
|
||||
assert result[0]["title"] == "Ondrej Megela"
|
||||
|
||||
|
||||
def test_student_family_evidence_includes_child_document(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
main = "pages/students/2016/jakub_maruniak/README.md"
|
||||
child = "pages/students/2016/jakub_maruniak/dp2021/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=main,
|
||||
title="Jakub Maruniak",
|
||||
text="Jakub Maruniak. Rok začiatku štúdia: 2016.",
|
||||
)
|
||||
add_document(
|
||||
db,
|
||||
path=child,
|
||||
title="Diplomová práca Jakuba Maruniaka",
|
||||
text=(
|
||||
"Jakub Maruniak. Diplomová práca 2021. "
|
||||
"Názov: Anotácia a rozpoznávanie pomenovaných entít "
|
||||
"v slovenskom jazyku."
|
||||
),
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Jakub Maruniak diplomová práca názov",
|
||||
[base_result(main, "Jakub Maruniak")],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
blocks = result[0].get("family_evidence_blocks") or []
|
||||
assert blocks
|
||||
assert any(block["document_path"] == child for block in blocks)
|
||||
assert any("Anotácia" in block["focus_text"] for block in blocks)
|
||||
|
||||
|
||||
def test_reverse_lookup_adds_document_with_exact_work_phrase(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
expected = "pages/students/2016/patrik_pavlisin/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=expected,
|
||||
title="Patrik Pavlišin",
|
||||
text=(
|
||||
"Bakalárska práca 2020. "
|
||||
"Názov bakalárskej práce: Metódy automatického prekladu."
|
||||
),
|
||||
)
|
||||
add_document(
|
||||
db,
|
||||
path="pages/students/2020/vladyslav_krupko/README.md",
|
||||
title="Vladyslav Krupko",
|
||||
text="Strojový preklad je súčasťou úloh.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Metódy automatického prekladu študent dokument téma",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result[0]["document_path"] == expected
|
||||
|
||||
|
||||
def test_bilingual_reverse_lookup_finds_internship(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
expected = "pages/interns/bogdan_paul_chis/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=expected,
|
||||
title="Bogdan Paul Chiș",
|
||||
text=(
|
||||
"Topic: (multilingual) Triplet extraction from medical data. "
|
||||
"Goal: Construct a knowledge graph from medical package inserts."
|
||||
),
|
||||
)
|
||||
add_document(
|
||||
db,
|
||||
path="pages/topics/otvorene/README.md",
|
||||
title="Otvorené projekty",
|
||||
text="Multilingual extraction from medical documents is an open topic.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"multilinguálna extrakcia trojíc z medicínskych dát",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result[0]["document_path"] == expected
|
||||
|
||||
|
||||
def test_exact_topic_page_title_is_promoted(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
topic = "pages/topics/translation/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=topic,
|
||||
title="Strojový preklad",
|
||||
text="Štatistický strojový preklad. Neurónový strojový preklad.",
|
||||
)
|
||||
add_document(
|
||||
db,
|
||||
path="pages/students/2016/patrik_pavlisin/README.md",
|
||||
title="Patrik Pavlišin",
|
||||
text="Metódy automatického a strojového prekladu.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Strojový preklad projektová stránka informačná stránka",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result[0]["document_path"] == topic
|
||||
|
||||
|
||||
def test_generic_question_prefers_root_topic_page(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
root = "pages/topics/question/README.md"
|
||||
child = "pages/topics/question/validacie/README.md"
|
||||
add_document(
|
||||
db,
|
||||
path=root,
|
||||
title="Question Answering",
|
||||
text=(
|
||||
"Anotácia otázok. Pre každý odsek má vzniknúť 5 otázok. "
|
||||
"Používa sa anotačný nástroj Prodigy."
|
||||
),
|
||||
)
|
||||
add_document(
|
||||
db,
|
||||
path=child,
|
||||
title="Validácia otázok a odpovedí",
|
||||
text="Anotácia otázok a validácia odsekov.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Koľko otázok má vzniknúť pre každý odsek pri anotácii otázok?",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result[0]["document_path"] == root
|
||||
|
||||
|
||||
def test_multi_document_query_adds_multiple_sources(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
paths = [
|
||||
"pages/students/2022/jan_ptak/README.md",
|
||||
"pages/students/2022/jan_malinovsky/README.md",
|
||||
"pages/students/2021/matej_scislak/README.md",
|
||||
]
|
||||
|
||||
for index, path in enumerate(paths):
|
||||
add_document(
|
||||
db,
|
||||
path=path,
|
||||
title=f"Student {index}",
|
||||
text="RAG a GraphRAG, retrieval augmented generation.",
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Nájdi viacero dokumentov, ktoré riešia RAG alebo GraphRAG.",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
assert is_multi_document_query(
|
||||
"Nájdi viacero dokumentov, ktoré riešia RAG alebo GraphRAG."
|
||||
)
|
||||
assert len(result) >= 2
|
||||
assert len({item["document_path"] for item in result}) >= 2
|
||||
|
||||
|
||||
def test_unpublished_supplement_is_excluded(tmp_path: Path) -> None:
|
||||
db = create_db(tmp_path)
|
||||
add_document(
|
||||
db,
|
||||
path="pages/topics/secret/README.md",
|
||||
title="Tajný projekt",
|
||||
text="Tajný projekt a unikátny výraz xyzabc.",
|
||||
published=False,
|
||||
)
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
db,
|
||||
"Tajný projekt xyzabc",
|
||||
[],
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_missing_database_is_safe_noop(tmp_path: Path) -> None:
|
||||
base = [
|
||||
base_result(
|
||||
"pages/students/2016/jan_holp/README.md",
|
||||
"Ján Holp",
|
||||
)
|
||||
]
|
||||
|
||||
result = expand_results_with_supplemental_sources(
|
||||
tmp_path / "missing.sqlite",
|
||||
"Ján Holp",
|
||||
base,
|
||||
published_only=True,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
assert result == base
|
||||
Loading…
Reference in New Issue
Block a user