pridanie_testov
This commit is contained in:
parent
67b51683cb
commit
07134b7681
828
test/test_rag.py
828
test/test_rag.py
@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@ -1835,3 +1837,829 @@ def test_rag_endpoint_live_context_has_source_boundaries(
|
||||
"context"
|
||||
]
|
||||
)
|
||||
|
||||
def create_section_lead_test_db(
|
||||
tmp_path: Path,
|
||||
) -> Path:
|
||||
db_path = (
|
||||
tmp_path
|
||||
/ "section_lead.sqlite"
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE chunks (
|
||||
id INTEGER PRIMARY KEY,
|
||||
chunk_id TEXT NOT NULL,
|
||||
document_path TEXT NOT NULL,
|
||||
chunk_index INTEGER NOT NULL,
|
||||
heading_paths_json TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
published INTEGER NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
return db_path
|
||||
|
||||
|
||||
def insert_section_lead_chunk(
|
||||
db_path: Path,
|
||||
*,
|
||||
chunk_id: str,
|
||||
document_path: 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,
|
||||
chunk_index,
|
||||
heading_paths_json,
|
||||
text,
|
||||
published
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
chunk_id,
|
||||
document_path,
|
||||
chunk_index,
|
||||
json.dumps(
|
||||
heading_paths,
|
||||
ensure_ascii=False,
|
||||
),
|
||||
text,
|
||||
(
|
||||
1
|
||||
if published
|
||||
else 0
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def section_lead_primary_result() -> dict[str, Any]:
|
||||
return {
|
||||
"chunk_id": (
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-3"
|
||||
),
|
||||
"chunk_index": 3,
|
||||
"document_path": (
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md"
|
||||
),
|
||||
"title": "Ján Holp",
|
||||
"author": "Daniel Hladek",
|
||||
"published": True,
|
||||
"heading_paths": [
|
||||
[
|
||||
"Ján Holp",
|
||||
"Diplomová práca 2021",
|
||||
],
|
||||
],
|
||||
"text": (
|
||||
"Najrelevantnejší úsek "
|
||||
"z neskoršej časti sekcie."
|
||||
),
|
||||
"source_url": (
|
||||
"https://zp.kemt.fei.tuke.sk/"
|
||||
"students/2016/jan_holp"
|
||||
),
|
||||
"match_strategy": "vector",
|
||||
"fts_rank": None,
|
||||
"vector_rank": 1,
|
||||
"vector_score": 0.91,
|
||||
"hybrid_score": 0.03,
|
||||
}
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_uses_earliest_chunk_in_same_section(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
heading_paths = result[
|
||||
"heading_paths"
|
||||
]
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=heading_paths,
|
||||
text=(
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií."
|
||||
),
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-1"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=1,
|
||||
heading_paths=heading_paths,
|
||||
text=(
|
||||
"Druhý chunk rovnakej sekcie."
|
||||
),
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert lead is not None
|
||||
|
||||
assert (
|
||||
lead["chunk_index"]
|
||||
== 0
|
||||
)
|
||||
|
||||
assert (
|
||||
lead["chunk_id"]
|
||||
== (
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
lead["text"]
|
||||
== (
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií."
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_ignores_other_sections(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=[
|
||||
[
|
||||
"Ján Holp",
|
||||
"Bakalárska práca 2019",
|
||||
],
|
||||
],
|
||||
text=(
|
||||
"Toto je iná sekcia."
|
||||
),
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-2"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=2,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=(
|
||||
"Správny začiatok "
|
||||
"diplomovej sekcie."
|
||||
),
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert lead is not None
|
||||
|
||||
assert (
|
||||
lead["chunk_index"]
|
||||
== 2
|
||||
)
|
||||
|
||||
assert (
|
||||
lead["text"]
|
||||
== (
|
||||
"Správny začiatok "
|
||||
"diplomovej sekcie."
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_ignores_other_documents(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2017/"
|
||||
"other/README.md::chunk-0"
|
||||
),
|
||||
document_path=(
|
||||
"pages/students/2017/"
|
||||
"other/README.md"
|
||||
),
|
||||
chunk_index=0,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=(
|
||||
"Rovnaká sekcia, "
|
||||
"ale iný dokument."
|
||||
),
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert lead is None
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_returns_none_when_primary_is_section_start(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
result[
|
||||
"chunk_id"
|
||||
] = (
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
)
|
||||
|
||||
result[
|
||||
"chunk_index"
|
||||
] = 0
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=result[
|
||||
"chunk_id"
|
||||
],
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=result[
|
||||
"text"
|
||||
],
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert lead is None
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_returns_none_without_heading_paths(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
result[
|
||||
"heading_paths"
|
||||
] = []
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert lead is None
|
||||
|
||||
|
||||
def test_load_section_lead_chunk_respects_published_only(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=(
|
||||
"Nepublikovaný začiatok "
|
||||
"sekcie."
|
||||
),
|
||||
published=False,
|
||||
)
|
||||
|
||||
with sqlite3.connect(
|
||||
db_path
|
||||
) as conn:
|
||||
conn.row_factory = (
|
||||
sqlite3.Row
|
||||
)
|
||||
|
||||
published_lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=True,
|
||||
)
|
||||
)
|
||||
|
||||
unrestricted_lead = (
|
||||
rag_utils.load_section_lead_chunk(
|
||||
conn,
|
||||
result,
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
published_lead
|
||||
is None
|
||||
)
|
||||
|
||||
assert (
|
||||
unrestricted_lead
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
def test_expand_results_with_section_leads_preserves_primary_retrieval_data(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=(
|
||||
"Začiatok relevantnej sekcie."
|
||||
),
|
||||
)
|
||||
|
||||
expanded = (
|
||||
rag_utils.expand_results_with_section_leads(
|
||||
db_path,
|
||||
[
|
||||
result
|
||||
],
|
||||
published_only=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
len(
|
||||
expanded
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
item = expanded[
|
||||
0
|
||||
]
|
||||
|
||||
assert (
|
||||
item[
|
||||
"chunk_id"
|
||||
]
|
||||
== result[
|
||||
"chunk_id"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"hybrid_score"
|
||||
]
|
||||
== result[
|
||||
"hybrid_score"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"vector_rank"
|
||||
]
|
||||
== result[
|
||||
"vector_rank"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"section_lead_text"
|
||||
]
|
||||
== (
|
||||
"Začiatok relevantnej sekcie."
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"context_expansion"
|
||||
][
|
||||
"applied"
|
||||
]
|
||||
is True
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"context_expansion"
|
||||
][
|
||||
"primary_chunk_id"
|
||||
]
|
||||
== result[
|
||||
"chunk_id"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
item[
|
||||
"context_expansion"
|
||||
][
|
||||
"lead_chunk_index"
|
||||
]
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
def test_build_source_text_combines_section_lead_and_primary_text() -> None:
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
result[
|
||||
"section_lead_text"
|
||||
] = (
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií."
|
||||
)
|
||||
|
||||
text = (
|
||||
rag_utils.build_source_text(
|
||||
result
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
"ZAČIATOK RELEVANTNEJ SEKCIE"
|
||||
in text
|
||||
)
|
||||
|
||||
assert (
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií."
|
||||
in text
|
||||
)
|
||||
|
||||
assert (
|
||||
"NAJRELEVANTNEJŠÍ NÁJDENÝ ÚSEK"
|
||||
in text
|
||||
)
|
||||
|
||||
assert (
|
||||
result[
|
||||
"text"
|
||||
]
|
||||
in text
|
||||
)
|
||||
|
||||
|
||||
def test_build_source_text_does_not_duplicate_identical_text() -> None:
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
result[
|
||||
"section_lead_text"
|
||||
] = result[
|
||||
"text"
|
||||
]
|
||||
|
||||
text = (
|
||||
rag_utils.build_source_text(
|
||||
result
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
text
|
||||
== result[
|
||||
"text"
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_build_rag_context_expands_section_lead_end_to_end(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
db_path = (
|
||||
create_section_lead_test_db(
|
||||
tmp_path
|
||||
)
|
||||
)
|
||||
|
||||
result = (
|
||||
section_lead_primary_result()
|
||||
)
|
||||
|
||||
insert_section_lead_chunk(
|
||||
db_path,
|
||||
chunk_id=(
|
||||
"pages/students/2016/"
|
||||
"jan_holp/README.md::chunk-0"
|
||||
),
|
||||
document_path=result[
|
||||
"document_path"
|
||||
],
|
||||
chunk_index=0,
|
||||
heading_paths=result[
|
||||
"heading_paths"
|
||||
],
|
||||
text=(
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií "
|
||||
"v slovenskom jazyku."
|
||||
),
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
assert (
|
||||
max_per_document
|
||||
== 1
|
||||
)
|
||||
|
||||
return {
|
||||
"engine": (
|
||||
"hybrid_fts5_embeddings"
|
||||
),
|
||||
"strategies": [
|
||||
"vector"
|
||||
],
|
||||
"results": [
|
||||
result
|
||||
],
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
rag_utils,
|
||||
"search_database",
|
||||
fake_search_database,
|
||||
)
|
||||
|
||||
response = (
|
||||
build_rag_context(
|
||||
db_path,
|
||||
(
|
||||
"Aká je téma diplomovej "
|
||||
"práce Jána Holpa?"
|
||||
),
|
||||
limit=5,
|
||||
max_per_document=1,
|
||||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
response[
|
||||
"source_count"
|
||||
]
|
||||
== 1
|
||||
)
|
||||
|
||||
assert (
|
||||
len(
|
||||
response[
|
||||
"sources"
|
||||
]
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
source = response[
|
||||
"sources"
|
||||
][0]
|
||||
|
||||
assert (
|
||||
source[
|
||||
"retrieval"
|
||||
][
|
||||
"hybrid_score"
|
||||
]
|
||||
== result[
|
||||
"hybrid_score"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií "
|
||||
"v slovenskom jazyku."
|
||||
in source[
|
||||
"text"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
result[
|
||||
"text"
|
||||
]
|
||||
in source[
|
||||
"text"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
"Názov diplomovej práce: "
|
||||
"Systém získavania informácií "
|
||||
"v slovenskom jazyku."
|
||||
in response[
|
||||
"context"
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
result[
|
||||
"text"
|
||||
]
|
||||
in response[
|
||||
"context"
|
||||
]
|
||||
)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user