144 lines
3.9 KiB
Python
144 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
import scripts.build_chunks as chunking
|
|
|
|
|
|
def configure_small_chunks(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
maximum: int = 90,
|
|
overlap: int = 18,
|
|
minimum: int = 25,
|
|
) -> None:
|
|
monkeypatch.setattr(chunking, "MAX_TOKENS", maximum)
|
|
monkeypatch.setattr(chunking, "OVERLAP_TOKENS", overlap)
|
|
monkeypatch.setattr(chunking, "MIN_CHUNK_TOKENS", minimum)
|
|
|
|
|
|
def body_without_context(text: str) -> str:
|
|
parts = text.split("\n\n", maxsplit=1)
|
|
return parts[1] if len(parts) == 2 else ""
|
|
|
|
|
|
def test_chunks_respect_token_limit_and_keep_context(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
configure_small_chunks(monkeypatch)
|
|
|
|
sentences = [
|
|
f"Veta {index} obsahuje dostatočne dlhý skúšobný text."
|
|
for index in range(1, 30)
|
|
]
|
|
markdown = "## Obsah\n\n" + " ".join(sentences)
|
|
|
|
drafts = chunking.chunk_markdown(
|
|
markdown,
|
|
document_title="Testovací dokument",
|
|
)
|
|
|
|
assert len(drafts) > 1
|
|
|
|
for draft in drafts:
|
|
assert draft.text.startswith("Dokument: Testovací dokument")
|
|
assert "Sekcia: Obsah" in draft.text
|
|
assert chunking.TOKEN_COUNTER.count(draft.text) <= 90
|
|
assert body_without_context(draft.text).startswith("Veta ")
|
|
|
|
|
|
def test_overlap_never_starts_in_middle_of_word_or_sentence(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
configure_small_chunks(monkeypatch, maximum=75, overlap=20, minimum=20)
|
|
|
|
markdown = "## Sekcia\n\n" + " ".join(
|
|
f"Presná veta číslo {index} sa končí bodkou."
|
|
for index in range(1, 24)
|
|
)
|
|
|
|
drafts = chunking.chunk_markdown(markdown, document_title="Dokument")
|
|
|
|
assert len(drafts) > 1
|
|
|
|
for draft in drafts:
|
|
body = body_without_context(draft.text)
|
|
assert body.startswith("Presná veta číslo ")
|
|
assert not body.startswith(("t is done", "he databases", "te down"))
|
|
|
|
|
|
def test_short_sections_are_merged_or_balanced(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
configure_small_chunks(monkeypatch, maximum=170, overlap=20, minimum=50)
|
|
|
|
markdown = """
|
|
## Úlohy
|
|
|
|
Krátke zadanie.
|
|
|
|
## Výsledky
|
|
|
|
Táto sekcia obsahuje dlhší text s výsledkami a ďalším vysvetlením. Druhá veta dopĺňa kontext a umožní bezpečné spojenie krátkej časti.
|
|
"""
|
|
|
|
drafts = chunking.chunk_markdown(markdown, document_title="Test")
|
|
|
|
assert drafts
|
|
assert any("Krátke zadanie." in draft.text for draft in drafts)
|
|
assert all(draft.units for draft in drafts)
|
|
assert all(
|
|
draft.text.strip() not in {
|
|
"Dokument: Test\nSekcia: Úlohy",
|
|
"Dokument: Test\n\nSekcia: Úlohy",
|
|
}
|
|
for draft in drafts
|
|
)
|
|
|
|
|
|
def test_empty_markdown_produces_no_chunks() -> None:
|
|
assert chunking.chunk_markdown("", document_title="Prázdny") == []
|
|
assert chunking.chunk_markdown("\n\n", document_title="Prázdny") == []
|
|
|
|
|
|
def test_heading_inside_code_fence_is_not_section(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
configure_small_chunks(monkeypatch, maximum=220, overlap=20, minimum=20)
|
|
|
|
markdown = """
|
|
# Skutočná sekcia
|
|
|
|
```markdown
|
|
## Toto nie je nadpis
|
|
print("ahoj")
|
|
```
|
|
|
|
Text po bloku kódu.
|
|
"""
|
|
|
|
drafts = chunking.chunk_markdown(markdown, document_title="Kód")
|
|
|
|
assert drafts
|
|
assert any("## Toto nie je nadpis" in draft.text for draft in drafts)
|
|
assert all(
|
|
"Toto nie je nadpis" not in path
|
|
for draft in drafts
|
|
for path in draft.heading_paths
|
|
)
|
|
|
|
|
|
def test_no_chunk_exceeds_configured_limit(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
configure_small_chunks(monkeypatch, maximum=60, overlap=10, minimum=10)
|
|
|
|
markdown = "# Dlhý text\n\n" + " ".join(
|
|
f"slovo{index}" for index in range(1, 250)
|
|
)
|
|
|
|
drafts = chunking.chunk_markdown(markdown, document_title="Limit")
|
|
|
|
assert drafts
|
|
assert max(chunking.TOKEN_COUNTER.count(draft.text) for draft in drafts) <= 60
|