wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
253 lines
9.8 KiB
Python
253 lines
9.8 KiB
Python
"""Tests for the recursive / markdown / parent_child / semantic strategies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from application.parser.chunking import Chunker
|
|
from application.parser.chunking_creator import ChunkerCreator
|
|
from application.parser.chunking_strategies import (
|
|
MarkdownChunker,
|
|
ParentChildChunker,
|
|
RecursiveChunker,
|
|
SemanticChunker,
|
|
)
|
|
from application.parser.schema.base import Document
|
|
from application.utils import get_encoding
|
|
|
|
|
|
def _tok(text: str) -> int:
|
|
return len(get_encoding().encode(text))
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRegistration:
|
|
def test_strategies_registered(self):
|
|
# create_chunker self-bootstraps the strategy module.
|
|
ChunkerCreator.create_chunker("recursive")
|
|
for key, cls in (
|
|
("recursive", RecursiveChunker),
|
|
("markdown", MarkdownChunker),
|
|
("parent_child", ParentChildChunker),
|
|
("semantic", SemanticChunker),
|
|
):
|
|
assert ChunkerCreator.chunkers.get(key) is cls
|
|
|
|
def test_worker_kwargs_accepted(self):
|
|
# The worker builds every strategy with the classic kwarg set.
|
|
for strat in ("recursive", "markdown", "parent_child", "semantic"):
|
|
chunker = ChunkerCreator.create_chunker(
|
|
strat,
|
|
chunking_strategy=strat,
|
|
max_tokens=200,
|
|
min_tokens=20,
|
|
duplicate_headers=False,
|
|
)
|
|
assert chunker.max_tokens == 200
|
|
assert chunker.min_tokens == 20
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRecursive:
|
|
def test_caps_at_max_tokens(self):
|
|
chunker = RecursiveChunker(max_tokens=40, min_tokens=5)
|
|
docs = [Document(text="word " * 500, doc_id="d")]
|
|
out = chunker.chunk(docs)
|
|
assert len(out) > 1
|
|
for c in out:
|
|
assert _tok(c.text) <= 40
|
|
assert c.extra_info["token_count"] == _tok(c.text)
|
|
|
|
def test_splits_on_separator_hierarchy(self):
|
|
# Paragraph boundaries should drive the split before token slicing.
|
|
text = "\n\n".join(["para " * 30 for _ in range(5)])
|
|
chunker = RecursiveChunker(max_tokens=60, min_tokens=5)
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) >= 2
|
|
for c in out:
|
|
assert _tok(c.text) <= 60
|
|
|
|
def test_small_doc_single_chunk(self):
|
|
chunker = RecursiveChunker(max_tokens=2000, min_tokens=1)
|
|
out = chunker.chunk([Document(text="short text here", doc_id="d")])
|
|
assert len(out) == 1
|
|
assert out[0].text.strip() == "short text here"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMarkdown:
|
|
def test_splits_on_headings(self):
|
|
text = "# A\nalpha\n\n## B\nbeta\n\n### C\ngamma"
|
|
chunker = MarkdownChunker(max_tokens=2000, min_tokens=1)
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
# One section per heading.
|
|
assert len(out) == 3
|
|
assert out[0].text.startswith("# A")
|
|
assert out[1].text.startswith("## B")
|
|
|
|
def test_oversized_section_token_capped(self):
|
|
text = "# Big\n" + "word " * 400
|
|
chunker = MarkdownChunker(max_tokens=50, min_tokens=5)
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) > 1
|
|
for c in out:
|
|
assert _tok(c.text) <= 50
|
|
|
|
def test_no_heading_falls_back_to_single_or_capped(self):
|
|
chunker = MarkdownChunker(max_tokens=2000, min_tokens=1)
|
|
out = chunker.chunk([Document(text="plain text no heading", doc_id="d")])
|
|
assert len(out) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestParentChild:
|
|
def test_children_smaller_than_parent(self):
|
|
chunker = ParentChildChunker(max_tokens=60, min_tokens=15)
|
|
out = chunker.chunk([Document(text="alpha " * 200, doc_id="d")])
|
|
assert len(out) > 1
|
|
for c in out:
|
|
assert _tok(c.text) <= 15
|
|
assert _tok(c.extra_info["parent_text"]) <= 60
|
|
assert _tok(c.text) <= _tok(c.extra_info["parent_text"])
|
|
|
|
def test_parent_text_reaches_vectorstore_metadata(self):
|
|
chunker = ParentChildChunker(max_tokens=80, min_tokens=20)
|
|
out = chunker.chunk([Document(text="beta " * 150, doc_id="d")])
|
|
lc = out[0].to_langchain_format()
|
|
# parent_text must survive the langchain conversion into metadata.
|
|
assert "parent_text" in lc.metadata
|
|
assert lc.metadata["parent_text"]
|
|
assert lc.page_content == out[0].text
|
|
|
|
def test_child_size_defaults_when_min_zero(self):
|
|
chunker = ParentChildChunker(max_tokens=200, min_tokens=0)
|
|
out = chunker.chunk([Document(text="gamma " * 200, doc_id="d")])
|
|
assert all("parent_text" in c.extra_info for c in out)
|
|
|
|
|
|
_EMB_TARGET = "application.vectorstore.base.EmbeddingsSingleton.get_instance"
|
|
|
|
|
|
class _FakeEmbeddings:
|
|
def __init__(self, vectors):
|
|
self._vectors = vectors
|
|
|
|
def embed_documents(self, sentences):
|
|
return self._vectors
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSemantic:
|
|
def test_breakpoint_forces_split(self):
|
|
# Two topics: sentences 0-1 vs 2-3, orthogonal embeddings between.
|
|
text = "Alpha one. Alpha two. Beta one. Beta two."
|
|
vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
|
|
with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) == 2
|
|
assert "Alpha" in out[0].text and "Beta" not in out[0].text
|
|
assert "Beta" in out[1].text and "Alpha" not in out[1].text
|
|
|
|
def test_no_breakpoint_single_chunk(self):
|
|
# Identical embeddings -> zero distances -> no split.
|
|
text = "Same one. Same two. Same three. Same four."
|
|
vectors = [[1.0, 0.0]] * 4
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
|
|
with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) == 1
|
|
assert out[0].extra_info["token_count"] == _tok(out[0].text)
|
|
|
|
def test_max_tokens_enforced(self):
|
|
# A single semantic group larger than max_tokens is hard-split.
|
|
long_sentence = "word " * 300 + "."
|
|
text = f"{long_sentence} {long_sentence}"
|
|
vectors = [[1.0, 0.0], [1.0, 0.0]]
|
|
chunker = SemanticChunker(max_tokens=40, min_tokens=0)
|
|
with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) > 1
|
|
for c in out:
|
|
assert _tok(c.text) <= 40
|
|
|
|
def test_min_tokens_merges_neighbours(self):
|
|
# Non-uniform distances yield several breakpoints and tiny groups,
|
|
# which must merge until they clear min_tokens.
|
|
text = "A. B. C. D. E. F."
|
|
vectors = [
|
|
[1.0, 0.0],
|
|
[0.0, 1.0],
|
|
[0.0, 1.0],
|
|
[1.0, 0.0],
|
|
[0.0, 1.0],
|
|
[0.0, 1.0],
|
|
]
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=8)
|
|
with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
assert len(out) < 6
|
|
assert _tok(out[0].text) >= 8
|
|
|
|
def test_embeddings_error_falls_back_to_recursive(self):
|
|
text = "First sentence here. Second sentence here. Third one."
|
|
|
|
def _boom(*args, **kwargs):
|
|
raise RuntimeError("model unavailable")
|
|
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
|
|
with patch(_EMB_TARGET, side_effect=_boom):
|
|
out = chunker.chunk([Document(text=text, doc_id="d")])
|
|
recursive = RecursiveChunker(max_tokens=2000, min_tokens=0)
|
|
expected = recursive.chunk([Document(text=text, doc_id="d")])
|
|
assert [c.text for c in out] == [c.text for c in expected]
|
|
|
|
def test_too_few_sentences_falls_back(self):
|
|
# A single sentence cannot be semantically split.
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
|
|
with patch(_EMB_TARGET, side_effect=AssertionError("must not embed")):
|
|
out = chunker.chunk([Document(text="just one sentence", doc_id="d")])
|
|
assert len(out) == 1
|
|
assert out[0].text.strip() == "just one sentence"
|
|
|
|
def test_source_and_extra_info_preserved(self):
|
|
text = "Alpha one. Alpha two. Beta one. Beta two."
|
|
vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]
|
|
doc = Document(
|
|
text=text,
|
|
doc_id="d",
|
|
extra_info={"source": "file.md", "title": "T"},
|
|
)
|
|
chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
|
|
with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)):
|
|
out = chunker.chunk([doc])
|
|
assert len(out) == 2
|
|
for c in out:
|
|
assert c.extra_info["source"] == "file.md"
|
|
assert c.extra_info["title"] == "T"
|
|
assert c.extra_info["token_count"] == _tok(c.text)
|
|
assert c.doc_id.startswith("d-")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestClassicByteIdentical:
|
|
def test_classic_chunk_unchanged(self):
|
|
# The new strategies must not perturb the classic baseline.
|
|
docs = [
|
|
Document(text="A short paragraph.", doc_id="small"),
|
|
Document(text="word " * 4000, doc_id="large"),
|
|
]
|
|
params = dict(max_tokens=1250, min_tokens=150, duplicate_headers=False)
|
|
direct = Chunker(chunking_strategy="classic_chunk", **params).chunk(docs)
|
|
via = ChunkerCreator.create_chunker("classic_chunk", **params).chunk(
|
|
[
|
|
Document(text="A short paragraph.", doc_id="small"),
|
|
Document(text="word " * 4000, doc_id="large"),
|
|
]
|
|
)
|
|
assert [(c.doc_id, c.text, c.extra_info) for c in via] == [
|
|
(c.doc_id, c.text, c.extra_info) for c in direct
|
|
]
|