1399 lines
27 KiB
Python
1399 lines
27 KiB
Python
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"
|
|
]
|
|
)
|