diff --git a/test/test_rag_document_expansion.py b/test/test_rag_document_expansion.py new file mode 100644 index 0000000..41351a0 --- /dev/null +++ b/test/test_rag_document_expansion.py @@ -0,0 +1,1398 @@ +from __future__ import annotations + +import json +import sqlite3 +from pathlib import Path +from typing import Any + +import pytest + +import scripts.rag_utils as rag_utils +from scripts.rag_document_expansion import ( + detect_section_intent, + expand_results_with_exact_document_section, + find_exact_student_document, + load_best_document_section_chunk, +) + + +MAROS_PATH = ( + "pages/students/2016/" + "maros_harahus/README.md" +) + +TOMAS_PATH = ( + "pages/students/2016/" + "tomas_kucharik/README.md" +) + +DAVID_PATH = ( + "pages/students/2019/" + "david_stromp/README.md" +) + + +def create_test_db( + tmp_path: Path, +) -> Path: + db_path = ( + tmp_path + / "rag_document_expansion.sqlite" + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.executescript( + """ + CREATE TABLE documents ( + id INTEGER PRIMARY KEY, + path TEXT UNIQUE NOT NULL, + title TEXT, + author TEXT, + published INTEGER, + content_length INTEGER + NOT NULL DEFAULT 0, + metadata_json TEXT + NOT NULL DEFAULT '{}' + ); + + 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, + text_length INTEGER + NOT NULL DEFAULT 0 + ); + """ + ) + + return db_path + + +def insert_document( + db_path: Path, + *, + path: str, + title: str, + author: str = "Daniel Hladek", + published: bool = True, +) -> None: + with sqlite3.connect( + db_path + ) as conn: + conn.execute( + """ + INSERT INTO documents ( + path, + title, + author, + published + ) + VALUES (?, ?, ?, ?) + """, + ( + path, + title, + author, + ( + 1 + if published + else 0 + ), + ), + ) + + +def insert_chunk( + db_path: Path, + *, + chunk_id: str, + document_path: str, + title: str, + author: str, + chunk_index: int, + heading_paths: list[Any], + text: str, + published: bool = True, +) -> None: + with sqlite3.connect( + db_path + ) as conn: + conn.execute( + """ + INSERT INTO chunks ( + chunk_id, + document_path, + title, + author, + published, + chunk_index, + heading_paths_json, + text, + text_length + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + chunk_id, + document_path, + title, + author, + ( + 1 + if published + else 0 + ), + chunk_index, + json.dumps( + heading_paths, + ensure_ascii=False, + ), + text, + len( + text + ), + ), + ) + + +def add_maros_document( + db_path: Path, +) -> None: + insert_document( + db_path, + path=MAROS_PATH, + title="Maroš Harahus", + ) + + insert_chunk( + db_path, + chunk_id=( + f"{MAROS_PATH}::chunk-2" + ), + document_path=( + MAROS_PATH + ), + title="Maroš Harahus", + author="Daniel Hladek", + chunk_index=2, + heading_paths=[ + [ + "Maroš Harahus", + ( + "Druhý rok " + "doktorandského štúdia" + ), + ], + ], + text=( + "Ciele doktorandského štúdia. " + "Obhájiť minimovku a " + "publikovať článok." + ), + ) + + insert_chunk( + db_path, + chunk_id=( + f"{MAROS_PATH}::chunk-20" + ), + document_path=( + MAROS_PATH + ), + title="Maroš Harahus", + author="Daniel Hladek", + chunk_index=20, + heading_paths=[ + [ + "Maroš Harahus", + "Diplomová práca 2021", + ], + ], + text=( + "Dokument: Maroš Harahus\n" + "Sekcia: Diplomová práca 2021\n\n" + "Názov diplomovej práce: " + "Neurónová morfologická anotácia " + "slovenského jazyka\n\n" + "Zadanie diplomovej práce." + ), + ) + + +def add_david_document( + db_path: Path, +) -> None: + insert_document( + db_path, + path=DAVID_PATH, + title="Dávid Stromp", + author="Maroš Harahus", + ) + + insert_chunk( + db_path, + chunk_id=( + f"{DAVID_PATH}::chunk-0" + ), + document_path=( + DAVID_PATH + ), + title="Dávid Stromp", + author="Maroš Harahus", + chunk_index=0, + heading_paths=[ + [ + "Bakalárska práca 2022" + ], + ], + text=( + "Dokument: Dávid Stromp\n" + "Autor dokumentu: Maroš Harahus\n" + "Sekcia: Bakalárska práca 2022." + ), + ) + + +def add_tomas_document( + db_path: Path, +) -> None: + insert_document( + db_path, + path=TOMAS_PATH, + title="Tomáš Kuchárik", + ) + + insert_chunk( + db_path, + chunk_id=( + f"{TOMAS_PATH}::chunk-0" + ), + document_path=( + TOMAS_PATH + ), + title="Tomáš Kuchárik", + author="Daniel Hladek", + chunk_index=0, + heading_paths=[ + [ + "Tomáš Kuchárik" + ], + ], + text=( + "Dokument: Tomáš Kuchárik\n" + "(študent KPI)\n" + "Súvisiace práce." + ), + ) + + insert_chunk( + db_path, + chunk_id=( + f"{TOMAS_PATH}::chunk-1" + ), + document_path=( + TOMAS_PATH + ), + title="Tomáš Kuchárik", + author="Daniel Hladek", + chunk_index=1, + heading_paths=[ + [ + "Tomáš Kuchárik", + "Diplomová práca 2022", + ], + ], + text=( + "Dokument: Tomáš Kuchárik\n" + "Sekcia: Diplomová práca 2022\n\n" + "Názov: Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou strojového prekladu\n\n" + "Zadanie:" + ), + ) + + insert_chunk( + db_path, + chunk_id=( + f"{TOMAS_PATH}::chunk-5" + ), + document_path=( + TOMAS_PATH + ), + title="Tomáš Kuchárik", + author="Daniel Hladek", + chunk_index=5, + heading_paths=[ + [ + "Tomáš Kuchárik", + "Diplomová práca 2021", + ], + ], + text=( + "Dokument: Tomáš Kuchárik\n" + "Sekcia: Diplomová práca 2021\n\n" + "Názov: Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou crowdsourcingu." + ), + ) + + +def retrieval_result( + *, + document_path: str, + chunk_id: str, + chunk_index: int, + title: str, + author: str, + heading_paths: list[Any], + text: str, +) -> dict[str, Any]: + return { + "chunk_id": chunk_id, + "chunk_index": ( + chunk_index + ), + "document_path": ( + document_path + ), + "title": title, + "author": author, + "published": True, + "heading_paths": ( + heading_paths + ), + "text": text, + "source_url": ( + "https://zp.kemt.fei.tuke.sk/" + + document_path.replace( + "pages/", + "", + ).replace( + "/README.md", + "", + ) + ), + "match_strategy": ( + "all_terms" + ), + "fts_rank": 1, + "vector_rank": 1, + "vector_score": 0.90, + "hybrid_score": 0.05, + "context_expansion": { + "strategy": ( + "section_lead" + ), + "applied": False, + "primary_chunk_id": ( + chunk_id + ), + "primary_chunk_index": ( + chunk_index + ), + "lead_chunk_id": None, + "lead_chunk_index": None, + }, + } + + +@pytest.mark.parametrize( + ( + "query", + "expected", + ), + [ + ( + ( + "Maroš Harahus " + "diplomová práca téma" + ), + "diploma_thesis", + ), + ( + ( + "Tomáš Kuchárik " + "bakalárska práca" + ), + "bachelor_thesis", + ), + ( + ( + "Ján Test " + "dizertačná práca" + ), + "doctoral_thesis", + ), + ( + ( + "Ján Test " + "diplomový projekt" + ), + "diploma_project", + ), + ( + "Maroš Harahus", + None, + ), + ], +) +def test_detect_section_intent( + query: str, + expected: str | None, +) -> None: + assert ( + detect_section_intent( + query + ) + == expected + ) + + +def test_find_exact_student_document_uses_title_not_author( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + add_david_document( + db_path + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + document = ( + find_exact_student_document( + conn, + ( + "Maroš Harahus " + "diplomová práca rok" + ), + published_only=True, + ) + ) + + assert ( + document + is not None + ) + + assert ( + document[ + "document_path" + ] + == MAROS_PATH + ) + + assert ( + document[ + "title" + ] + == "Maroš Harahus" + ) + + +def test_load_best_chunk_q0020_selects_maros_diploma_section( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + document = ( + find_exact_student_document( + conn, + ( + "Maroš Harahus " + "diplomová práca téma" + ), + published_only=True, + ) + ) + + assert ( + document + is not None + ) + + chunk = ( + load_best_document_section_chunk( + conn, + query=( + "Maroš Harahus " + "diplomová práca téma" + ), + document=document, + intent="diploma_thesis", + published_only=True, + ) + ) + + assert chunk is not None + + assert ( + chunk[ + "chunk_id" + ] + == ( + f"{MAROS_PATH}::chunk-20" + ) + ) + + assert ( + "Neurónová morfologická " + "anotácia slovenského jazyka" + in chunk[ + "text" + ] + ) + + +def test_q0020_expands_wrong_chunk_in_correct_document( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + primary = retrieval_result( + document_path=( + MAROS_PATH + ), + chunk_id=( + f"{MAROS_PATH}::chunk-2" + ), + chunk_index=2, + title="Maroš Harahus", + author="Daniel Hladek", + heading_paths=[ + [ + "Maroš Harahus", + ( + "Druhý rok " + "doktorandského štúdia" + ), + ], + ], + text=( + "Ciele doktorandského štúdia." + ), + ) + + expanded = ( + expand_results_with_exact_document_section( + db_path, + ( + "Maroš Harahus " + "diplomová práca téma" + ), + [ + primary + ], + published_only=True, + limit=5, + ) + ) + + assert ( + len( + expanded + ) + == 1 + ) + + item = expanded[ + 0 + ] + + assert ( + item[ + "chunk_id" + ] + == primary[ + "chunk_id" + ] + ) + + assert ( + "Neurónová morfologická " + "anotácia slovenského jazyka" + in item[ + "exact_document_text" + ] + ) + + assert ( + item[ + "document_expansion" + ][ + "applied" + ] + is True + ) + + assert ( + item[ + "document_expansion" + ][ + "added_source" + ] + is False + ) + + +def test_q0021_adds_exact_maros_document_before_author_document( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + add_david_document( + db_path + ) + + david_result = ( + retrieval_result( + document_path=( + DAVID_PATH + ), + chunk_id=( + f"{DAVID_PATH}::chunk-0" + ), + chunk_index=0, + title="Dávid Stromp", + author="Maroš Harahus", + heading_paths=[ + [ + "Bakalárska práca 2022" + ], + ], + text=( + "Autor dokumentu: " + "Maroš Harahus." + ), + ) + ) + + expanded = ( + expand_results_with_exact_document_section( + db_path, + ( + "Maroš Harahus " + "diplomová práca rok" + ), + [ + david_result + ], + published_only=True, + limit=5, + ) + ) + + assert ( + len( + expanded + ) + == 2 + ) + + assert ( + expanded[ + 0 + ][ + "document_path" + ] + == MAROS_PATH + ) + + assert ( + expanded[ + 0 + ][ + "title" + ] + == "Maroš Harahus" + ) + + assert ( + "Diplomová práca 2021" + in expanded[ + 0 + ][ + "text" + ] + ) + + assert ( + expanded[ + 0 + ][ + "document_expansion" + ][ + "added_source" + ] + is True + ) + + assert ( + expanded[ + 1 + ][ + "document_path" + ] + == DAVID_PATH + ) + + +def test_q0028_selects_answer_chunk_in_tomas_document( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_tomas_document( + db_path + ) + + primary = retrieval_result( + document_path=( + TOMAS_PATH + ), + chunk_id=( + f"{TOMAS_PATH}::chunk-0" + ), + chunk_index=0, + title="Tomáš Kuchárik", + author="Daniel Hladek", + heading_paths=[ + [ + "Tomáš Kuchárik" + ], + ], + text=( + "Dokument: Tomáš Kuchárik\n" + "Súvisiace práce." + ), + ) + + expanded = ( + expand_results_with_exact_document_section( + db_path, + ( + "Tomáš Kuchárik " + "diplomová práca" + ), + [ + primary + ], + published_only=True, + limit=5, + ) + ) + + assert ( + len( + expanded + ) + == 1 + ) + + assert ( + "Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou strojového prekladu" + in expanded[ + 0 + ][ + "exact_document_text" + ] + ) + + +def test_latest_work_section_is_preferred_without_year( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_tomas_document( + db_path + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + document = ( + find_exact_student_document( + conn, + ( + "Tomáš Kuchárik " + "diplomová práca" + ), + published_only=True, + ) + ) + + assert document is not None + + chunk = ( + load_best_document_section_chunk( + conn, + query=( + "Tomáš Kuchárik " + "diplomová práca" + ), + document=document, + intent="diploma_thesis", + published_only=True, + ) + ) + + assert chunk is not None + + assert ( + chunk[ + "chunk_index" + ] + == 1 + ) + + assert ( + "strojového prekladu" + in chunk[ + "text" + ] + ) + + +def test_explicit_year_selects_requested_work_section( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_tomas_document( + db_path + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + document = ( + find_exact_student_document( + conn, + ( + "Tomáš Kuchárik " + "diplomová práca 2021" + ), + published_only=True, + ) + ) + + assert document is not None + + chunk = ( + load_best_document_section_chunk( + conn, + query=( + "Tomáš Kuchárik " + "diplomová práca 2021" + ), + document=document, + intent="diploma_thesis", + published_only=True, + ) + ) + + assert chunk is not None + + assert ( + chunk[ + "chunk_index" + ] + == 5 + ) + + assert ( + "crowdsourcingu" + in chunk[ + "text" + ] + ) + + +def test_no_exact_title_leaves_results_unchanged( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + original = retrieval_result( + document_path=( + DAVID_PATH + ), + chunk_id=( + f"{DAVID_PATH}::chunk-0" + ), + chunk_index=0, + title="Dávid Stromp", + author="Maroš Harahus", + heading_paths=[ + [ + "Bakalárska práca 2022" + ], + ], + text="Test", + ) + + expanded = ( + expand_results_with_exact_document_section( + db_path, + ( + "Neznámy Študent " + "diplomová práca" + ), + [ + original + ], + published_only=True, + limit=5, + ) + ) + + assert ( + expanded + == [ + original + ] + ) + + +def test_query_without_section_intent_leaves_results_unchanged( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + original = retrieval_result( + document_path=( + MAROS_PATH + ), + chunk_id=( + f"{MAROS_PATH}::chunk-2" + ), + chunk_index=2, + title="Maroš Harahus", + author="Daniel Hladek", + heading_paths=[ + [ + "Maroš Harahus" + ], + ], + text="Test", + ) + + expanded = ( + expand_results_with_exact_document_section( + db_path, + "Maroš Harahus", + [ + original + ], + published_only=True, + limit=5, + ) + ) + + assert ( + expanded + == [ + original + ] + ) + + +def test_published_only_does_not_use_unpublished_document( + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + insert_document( + db_path, + path=MAROS_PATH, + title="Maroš Harahus", + published=False, + ) + + with sqlite3.connect( + db_path + ) as conn: + conn.row_factory = ( + sqlite3.Row + ) + + document = ( + find_exact_student_document( + conn, + ( + "Maroš Harahus " + "diplomová práca" + ), + published_only=True, + ) + ) + + assert ( + document + is None + ) + + +def test_build_source_text_includes_exact_document_section() -> None: + result = retrieval_result( + document_path=( + TOMAS_PATH + ), + chunk_id=( + f"{TOMAS_PATH}::chunk-0" + ), + chunk_index=0, + title="Tomáš Kuchárik", + author="Daniel Hladek", + heading_paths=[ + [ + "Tomáš Kuchárik" + ], + ], + text=( + "Primárny retrieval úsek." + ), + ) + + result[ + "exact_document_text" + ] = ( + "Diplomová práca 2022\n" + "Názov: Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou strojového prekladu" + ) + + text = ( + rag_utils.build_source_text( + result + ) + ) + + assert ( + "RELEVANTNÁ SEKCIA PRESNE " + "ZHODNÉHO DOKUMENTU" + in text + ) + + assert ( + "strojového prekladu" + in text + ) + + assert ( + "Primárny retrieval úsek." + in text + ) + + +def test_build_rag_context_q0028_regression( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_tomas_document( + db_path + ) + + primary = retrieval_result( + document_path=( + TOMAS_PATH + ), + chunk_id=( + f"{TOMAS_PATH}::chunk-0" + ), + chunk_index=0, + title="Tomáš Kuchárik", + author="Daniel Hladek", + heading_paths=[ + [ + "Tomáš Kuchárik" + ], + ], + text=( + "Dokument: Tomáš Kuchárik\n" + "Súvisiace práce." + ), + ) + + def fake_search_database( + db_path_arg: Path, + query: str, + limit: int, + *, + published_only: bool, + max_per_document: int, + ) -> dict[str, Any]: + assert ( + db_path_arg + == db_path + ) + + return { + "engine": ( + "hybrid_fts5_embeddings" + ), + "strategies": [ + "all_terms" + ], + "results": [ + primary + ], + } + + monkeypatch.setattr( + rag_utils, + "search_database", + fake_search_database, + ) + + response = ( + rag_utils.build_rag_context( + db_path, + ( + "Tomáš Kuchárik " + "diplomová práca" + ), + limit=5, + published_only=True, + max_per_document=1, + ) + ) + + assert ( + response[ + "source_count" + ] + == 1 + ) + + assert ( + "Tvorba korpusu otázok " + "a odpovedí v slovenskom jazyku " + "pomocou strojového prekladu" + in response[ + "context" + ] + ) + + assert ( + response[ + "sources" + ][0][ + "source_url" + ] + == ( + "https://zp.kemt.fei.tuke.sk/" + "students/2016/tomas_kucharik" + ) + ) + + +def test_build_rag_context_q0021_regression( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + db_path = ( + create_test_db( + tmp_path + ) + ) + + add_maros_document( + db_path + ) + + add_david_document( + db_path + ) + + david_result = ( + retrieval_result( + document_path=( + DAVID_PATH + ), + chunk_id=( + f"{DAVID_PATH}::chunk-0" + ), + chunk_index=0, + title="Dávid Stromp", + author="Maroš Harahus", + heading_paths=[ + [ + "Bakalárska práca 2022" + ], + ], + text=( + "Dokument: Dávid Stromp\n" + "Autor dokumentu: " + "Maroš Harahus." + ), + ) + ) + + def fake_search_database( + db_path_arg: Path, + query: str, + limit: int, + *, + published_only: bool, + max_per_document: int, + ) -> dict[str, Any]: + return { + "engine": ( + "hybrid_fts5_embeddings" + ), + "strategies": [ + "prefix_terms" + ], + "results": [ + david_result + ], + } + + monkeypatch.setattr( + rag_utils, + "search_database", + fake_search_database, + ) + + response = ( + rag_utils.build_rag_context( + db_path, + ( + "Maroš Harahus " + "diplomová práca rok" + ), + limit=5, + published_only=True, + max_per_document=1, + ) + ) + + assert ( + response[ + "sources" + ][0][ + "title" + ] + == "Maroš Harahus" + ) + + assert ( + "Diplomová práca 2021" + in response[ + "context" + ] + ) + + assert ( + "Neurónová morfologická " + "anotácia slovenského jazyka" + in response[ + "context" + ] + )