344 lines
8.1 KiB
Python
344 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
from tools.rag.rag_query_evidence import (
|
|
build_focus_excerpt,
|
|
expand_results_with_query_evidence,
|
|
find_best_query_evidence_chunk,
|
|
)
|
|
|
|
|
|
def test_profile_header_does_not_hide_later_application_evidence():
|
|
text = (
|
|
"Dokument: Anna Nová\n\nAnna Nová\n"
|
|
+ "\n".join(["Všeobecná poznámka."] * 25)
|
|
+ "\nState:\n- Proposed own Flask application\n- Created a FastAPI application with a data model.\n"
|
|
)
|
|
excerpt = build_focus_excerpt("Aký webový framework použila Anna Nová pre aplikáciu?", text)
|
|
assert "Created a FastAPI application" in excerpt
|
|
assert "Proposed own Flask application" in excerpt
|
|
|
|
|
|
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]
|