373 lines
9.9 KiB
Python
373 lines
9.9 KiB
Python
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
|