837 lines
20 KiB
Python
837 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import tools.rag.rag_utils as rag_utils
|
|
from tools.rag.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"]
|