hkuds--lightrag/tests/chunker/test_chunker_semantic_vector.py

"""Unit tests for ``chunking_by_semantic_vector`` (process_options=V)."""

import asyncio
import logging

import numpy as np
import pytest

pytest.importorskip("langchain_experimental")

from lightrag.chunker import chunking_by_semantic_vector  # noqa: E402
from lightrag.utils import EmbeddingFunc, Tokenizer, TokenizerInterface  # noqa: E402


class _CharTokenizer(TokenizerInterface):
    """1 char ≈ 1 token."""

    def encode(self, content: str):
        return [ord(ch) for ch in content]

    def decode(self, tokens):
        return "".join(chr(t) for t in tokens)


def _tok() -> Tokenizer:
    return Tokenizer("char-tokenizer", _CharTokenizer())


def _make_deterministic_embedding(dim: int = 8) -> EmbeddingFunc:
    """A toy async embedding func that hashes each input text into a
    stable unit vector — enough to drive SemanticChunker without needing
    a real model."""

    async def _embed(texts, **kwargs):
        rng = np.random.default_rng(seed=0)
        # Use a simple hash → seeded rng to get reproducible vectors per text.
        rows = []
        for text in texts:
            seed = abs(hash(text)) % (2**32)
            rng = np.random.default_rng(seed=seed)
            vec = rng.normal(size=dim).astype(np.float32)
            vec /= np.linalg.norm(vec) or 1.0
            rows.append(vec)
        return np.vstack(rows)

    return EmbeddingFunc(embedding_dim=dim, max_token_size=4096, func=_embed)


@pytest.mark.offline
def test_v_chunker_runs_with_stub_embedding():
    """Async chunker should split a multi-sentence body into ≥1 chunk
    when given a working embedding func."""
    body = (
        "Quantum mechanics describes nature at small scales. "
        "It contradicts classical intuition. "
        "Bread is baked from flour. "
        "Sourdough requires a long fermentation. "
    )

    async def _run():
        chunks = await chunking_by_semantic_vector(
            _tok(),
            body,
            chunk_token_size=200,
            embedding_func=_make_deterministic_embedding(),
        )
        return chunks

    chunks = asyncio.run(_run())

    assert len(chunks) >= 1
    # Each chunk dict has the canonical schema.
    assert all({"tokens", "content", "chunk_order_index"} <= set(c) for c in chunks)
    # chunk_order_index is contiguous starting at 0.
    assert [c["chunk_order_index"] for c in chunks] == list(range(len(chunks)))
    # No empty content rows.
    assert all(c["content"].strip() for c in chunks)
    for chunk in chunks:
        span = chunk["_source_span"]
        assert body[span["start"] : span["end"]] == chunk["content"]


class _ListHandler(logging.Handler):
    def __init__(self) -> None:
        super().__init__()
        self.records: list[logging.LogRecord] = []

    def emit(self, record: logging.LogRecord) -> None:
        self.records.append(record)


@pytest.mark.offline
def test_v_chunker_falls_back_to_recursive_when_no_embedding():
    """When ``embedding_func`` is None, V must log a warning and route
    to chunking_by_recursive_character (R) — V's only differentiator
    is embeddings, so without them R is the closest neighbour."""
    body = "Para A.\n\nPara B for fallback test.\n\nPara C."

    lightrag_logger = logging.getLogger("lightrag")
    handler = _ListHandler()
    handler.setLevel(logging.WARNING)
    lightrag_logger.addHandler(handler)
    try:

        async def _run():
            return await chunking_by_semantic_vector(
                _tok(),
                body,
                chunk_token_size=20,
                embedding_func=None,
            )

        chunks = asyncio.run(_run())
    finally:
        lightrag_logger.removeHandler(handler)

    assert len(chunks) >= 1
    assert any(
        "embedding_func is None" in rec.getMessage()
        for rec in handler.records
        if rec.levelno == logging.WARNING
    )


@pytest.mark.offline
def test_v_chunker_empty_input_returns_empty_list():
    async def _run():
        return await chunking_by_semantic_vector(_tok(), "")

    assert asyncio.run(_run()) == []


@pytest.mark.offline
def test_v_chunker_spans_repeated_sentences_and_number_of_chunks():
    body = (
        "Repeat sentence.\n"
        "Repeat sentence. "
        "A distant topic appears. "
        "Another distant topic appears."
    )

    async def _run():
        return await chunking_by_semantic_vector(
            _tok(),
            body,
            chunk_token_size=400,
            embedding_func=_make_deterministic_embedding(),
            number_of_chunks=2,
        )

    chunks = asyncio.run(_run())

    assert chunks
    for chunk in chunks:
        span = chunk["_source_span"]
        assert body[span["start"] : span["end"]] == chunk["content"]


@pytest.mark.offline
def test_v_oversized_group_resplit_preserves_child_source_spans():
    body = " ".join(f"word{i:02d}" for i in range(30)) + "."

    async def _run():
        return await chunking_by_semantic_vector(
            _tok(),
            body,
            chunk_token_size=35,
            embedding_func=_make_deterministic_embedding(),
        )

    chunks = asyncio.run(_run())

    assert len(chunks) > 1
    for chunk in chunks:
        span = chunk["_source_span"]
        assert body[span["start"] : span["end"]] == chunk["content"]
        assert chunk["tokens"] <= 35


@pytest.mark.offline
def test_semantic_groups_mirror_matches_upstream_split_text():
    """Drift guard: ``_semantic_groups_with_spans`` re-implements the private body
    of ``SemanticChunker.split_text`` to keep verbatim source spans. If an upstream
    langchain-experimental release changes that grouping (or a private member it
    relies on), this catches it — the mirror must yield the same group count and the
    same content modulo whitespace (it keeps the verbatim slice; upstream reflows
    sentences with single spaces)."""
    from langchain_core.embeddings import Embeddings
    from langchain_experimental.text_splitter import SemanticChunker

    from lightrag.chunker.semantic_vector import _semantic_groups_with_spans
    from lightrag.constants import DEFAULT_SENTENCE_SPLIT_REGEX

    class _SyncEmbeddings(Embeddings):
        """Deterministic per-text vectors, mirroring _make_deterministic_embedding
        but synchronous so SemanticChunker can be driven directly."""

        def embed_documents(self, texts):
            rows = []
            for text in texts:
                seed = abs(hash(text)) % (2**32)
                rng = np.random.default_rng(seed=seed)
                vec = rng.normal(size=8).astype(np.float64)
                vec /= np.linalg.norm(vec) or 1.0
                rows.append(vec.tolist())
            return rows

        def embed_query(self, text):
            return self.embed_documents([text])[0]

    text = (
        "Quantum mechanics describes nature at small scales. "
        "It contradicts classical intuition. "
        "Bread is baked from flour. "
        "Sourdough requires a long fermentation."
    )
    splitter = SemanticChunker(
        _SyncEmbeddings(),
        buffer_size=1,
        breakpoint_threshold_type="percentile",
        sentence_split_regex=DEFAULT_SENTENCE_SPLIT_REGEX,
    )

    # Same process → identical deterministic embeddings on both calls → identical
    # distances/breakpoints, so any divergence is a real grouping drift.
    upstream = splitter.split_text(text)
    mirror = _semantic_groups_with_spans(splitter, text)

    assert len(mirror) == len(upstream)
    for (grp_text, start, end), up in zip(mirror, upstream):
        assert grp_text == text[start:end]  # mirror keeps the verbatim source slice
        assert "".join(grp_text.split()) == "".join(up.split())