diff --git a/test/test_rag_query_evidence.py b/test/test_rag_query_evidence.py new file mode 100644 index 0000000..87bef00 --- /dev/null +++ b/test/test_rag_query_evidence.py @@ -0,0 +1,521 @@ +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +from scripts.rag_query_evidence import ( + build_focus_excerpt, + expand_results_with_query_evidence, + find_best_query_evidence_chunk, +) + + +def make_db( + tmp_path: Path, +) -> Path: + db = ( + tmp_path + / "evidence.sqlite" + ) + + with sqlite3.connect( + db + ) as conn: + conn.execute( + """ + CREATE TABLE chunks ( + id INTEGER PRIMARY KEY, + chunk_id TEXT UNIQUE NOT NULL, + document_path TEXT NOT NULL, + title TEXT, + author TEXT, + published INTEGER, + chunk_index INTEGER NOT NULL, + heading_paths_json TEXT + NOT NULL DEFAULT '[]', + text TEXT NOT NULL + ) + """ + ) + + return db + + +def add_chunk( + db: Path, + *, + path: str, + index: int, + title: str, + heading: str, + text: str, + published: bool = True, +) -> None: + with sqlite3.connect( + db + ) as conn: + conn.execute( + """ + INSERT INTO chunks ( + chunk_id, + document_path, + title, + author, + published, + chunk_index, + heading_paths_json, + text + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + f"{path}::chunk-{index}", + path, + title, + "Daniel Hladek", + ( + 1 + if published + else 0 + ), + index, + heading, + text, + ), + ) + + +def result( + path: str, + index: int, + title: str, + text: str, +) -> dict: + return { + "chunk_id": ( + f"{path}::chunk-{index}" + ), + "document_path": path, + "title": title, + "author": "Daniel Hladek", + "published": True, + "chunk_index": index, + "heading_paths": [], + "text": text, + "source_url": ( + "https://example.test/student" + ), + "match_strategy": "hybrid", + "fts_rank": 1, + "vector_rank": 1, + "vector_score": 0.9, + "hybrid_score": 0.03, + } + + +def test_reverse_title_selects_tomas_2022_chunk( + tmp_path: Path, +) -> None: + db = make_db( + tmp_path + ) + + path = ( + "pages/students/2016/" + "tomas_kucharik/README.md" + ) + + exact_title = ( + "Tvorba korpusu otázok a odpovedí " + "v slovenskom jazyku pomocou " + "strojového prekladu" + ) + + add_chunk( + db, + path=path, + index=1, + title="Tomáš Kuchárik", + heading=( + '[["Tomáš Kuchárik", ' + '"Diplomová práca 2022"]]' + ), + text=( + "## Diplomová práca 2022\n\n" + f"Názov: {exact_title}" + ), + ) + + add_chunk( + db, + path=path, + index=5, + title="Tomáš Kuchárik", + heading=( + '[["Tomáš Kuchárik", ' + '"Diplomová práca 2021"]]' + ), + text=( + "## Diplomová práca 2021\n\n" + "Názov: Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou crowdsourcingu" + ), + ) + + with sqlite3.connect( + db + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + best = ( + find_best_query_evidence_chunk( + conn, + path, + f"{exact_title} autor", + ) + ) + + assert best is not None + + assert ( + best[ + "chunk_index" + ] + == 1 + ) + + assert ( + "strojového prekladu" + in best[ + "text" + ] + ) + + +def test_expansion_preserves_retrieval_chunk_and_adds_2022_evidence( + tmp_path: Path, +) -> None: + db = make_db( + tmp_path + ) + + path = ( + "pages/students/2016/" + "tomas_kucharik/README.md" + ) + + exact_title = ( + "Tvorba korpusu otázok a odpovedí " + "v slovenskom jazyku pomocou " + "strojového prekladu" + ) + + add_chunk( + db, + path=path, + index=1, + title="Tomáš Kuchárik", + heading=( + '[["Tomáš Kuchárik", ' + '"Diplomová práca 2022"]]' + ), + text=( + f"Názov: {exact_title}" + ), + ) + + add_chunk( + db, + path=path, + index=5, + title="Tomáš Kuchárik", + heading=( + '[["Tomáš Kuchárik", ' + '"Diplomová práca 2021"]]' + ), + text=( + "Názov: Tvorba korpusu " + "pomocou crowdsourcingu" + ), + ) + + original = result( + path, + 5, + "Tomáš Kuchárik", + "crowdsourcing", + ) + + expanded = ( + expand_results_with_query_evidence( + db, + exact_title, + [ + original + ], + ) + ) + + item = expanded[ + 0 + ] + + assert ( + item[ + "chunk_id" + ] + == original[ + "chunk_id" + ] + ) + + assert ( + item[ + "hybrid_score" + ] + == original[ + "hybrid_score" + ] + ) + + assert ( + item[ + "query_evidence" + ][ + "applied" + ] + is True + ) + + assert ( + item[ + "query_evidence" + ][ + "evidence_chunk_index" + ] + == 1 + ) + + assert ( + "strojového prekladu" + in item[ + "query_focus_text" + ] + ) + + +def test_us_steel_focus_contains_gnn_task() -> None: + text = ( + "Stretnutie 1.10.\n" + "Stav:\n" + "- Štúdium základov neurónových sietí\n" + "- Úvodné stretnutie s US Steel\n" + "Úlohy:\n" + "- Vypracovať prehľad aktuálnych " + "metód grafových neurónových sietí\n" + "- Nájsť a vyskúšať toolkit na GNN.\n" + "- Naštudovať dáta z US Steel." + ) + + excerpt = ( + build_focus_excerpt( + ( + "Maroš Harahus US Steel " + "metódy študovať" + ), + text, + ) + ) + + assert ( + "US Steel" + in excerpt + ) + + assert ( + "grafových neurónových sietí" + in excerpt + ) + + assert ( + "toolkit na GNN" + in excerpt + ) + + +def test_maros_own_document_gets_gnn_evidence( + tmp_path: Path, +) -> None: + db = make_db( + tmp_path + ) + + path = ( + "pages/students/2016/" + "maros_harahus/README.md" + ) + + text = ( + "Úlohy:\n" + "- Vypracovať prehľad aktuálnych " + "metód grafových neurónových sietí\n" + "- Nájsť a vyskúšať toolkit na GNN." + ) + + add_chunk( + db, + path=path, + index=19, + title="Maroš Harahus", + heading=( + '[["Maroš Harahus", ' + '"Prvý ročník PhD štúdia"]]' + ), + text=text, + ) + + expanded = ( + expand_results_with_query_evidence( + db, + ( + "Maroš Harahus " + "grafové neurónové siete" + ), + [ + result( + path, + 19, + "Maroš Harahus", + text, + ) + ], + ) + ) + + assert ( + expanded[ + 0 + ][ + "query_evidence" + ][ + "applied" + ] + is True + ) + + assert ( + "grafových neurónových sietí" + in expanded[ + 0 + ][ + "query_focus_text" + ] + ) + + +def test_published_only_ignores_unpublished_better_chunk( + tmp_path: Path, +) -> None: + db = make_db( + tmp_path + ) + + path = ( + "pages/students/2016/" + "tomas_kucharik/README.md" + ) + + query = ( + "Tvorba korpusu otázok a odpovedí " + "pomocou strojového prekladu" + ) + + add_chunk( + db, + path=path, + index=1, + title="Tomáš Kuchárik", + heading=( + '[["Diplomová práca 2022"]]' + ), + text=( + "strojový preklad" + ), + ) + + add_chunk( + db, + path=path, + index=2, + title="Tomáš Kuchárik", + heading=( + '[["Diplomová práca 2022"]]' + ), + text=query, + published=False, + ) + + with sqlite3.connect( + db + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + best = ( + find_best_query_evidence_chunk( + conn, + path, + query, + published_only=True, + ) + ) + + assert best is not None + + assert ( + best[ + "chunk_index" + ] + == 1 + ) + + +def test_missing_chunks_table_is_safe_noop( + tmp_path: Path, +) -> None: + db = ( + tmp_path + / "empty.sqlite" + ) + + with sqlite3.connect( + db + ): + pass + + original = result( + ( + "pages/students/2016/" + "test/README.md" + ), + 0, + "Test", + "text", + ) + + assert ( + expand_results_with_query_evidence( + db, + "nejaký dotaz", + [ + original + ], + ) + == [ + original + ] + )