feat: add chunking module, remove chonkie (#84)

* retry

* simple version

* refactor

* move chunk to chunk

* remove language detection

* remove group

* add tests, rename

* rename function raw -> inner

* comment desired chunk length, move to chunkboundary

Author: stephantul

pyproject.toml

@@ -30,8 +30,9 @@ dependencies = [
     "vicinity>=0.4.4",
     "numpy>=1.24.0",
     "bm25s>=0.2.0",
-    "chonkie[code]!=1.6.3",
     "pathspec>=0.12",
+    "tree-sitter>=0.25",
+    "tree-sitter-language-pack!=1.6.3"
 ]
 
 [project.optional-dependencies]

src/semble/chunking/__init__.py

@@ -0,0 +1,3 @@
+from semble.chunking.chunking import chunk_source
+
+__all__ = ["chunk_source"]

src/semble/chunking/chunking.py

@@ -0,0 +1,36 @@
+import logging
+
+from semble.chunking.core import chunk, chunk_lines, is_supported_language
+from semble.types import Chunk
+
+logger = logging.getLogger(__name__)
+
+# The desired length of chunks in chars.
+# TODO: makes this configurable
+_DESIRED_CHUNK_LENGTH_CHARS = 1500
+
+
+def chunk_source(source: str, file_path: str, language: str | None) -> list[Chunk]:
+    """Chunk pre-read source text."""
+    if not source.strip():
+        return []
+    if language is not None and is_supported_language(language):
+        chunk_boundaries = chunk(source, language, _DESIRED_CHUNK_LENGTH_CHARS)
+    else:
+        chunk_boundaries = chunk_lines(source, _DESIRED_CHUNK_LENGTH_CHARS)
+
+    chunks: list[Chunk] = []
+    for boundary in chunk_boundaries:
+        # Clamp to start_index so zero-length chunks don't produce an off-by-one.
+        end_index = max(boundary.end - 1, boundary.start)
+        text = source[boundary.start : end_index + 1]
+        chunks.append(
+            Chunk(
+                content=text,
+                file_path=file_path,
+                start_line=source[: boundary.start].count("\n") + 1,
+                end_line=source[:end_index].count("\n") + 1,
+                language=language,
+            )
+        )
+    return chunks

src/semble/chunking/core.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cache
+from logging import getLogger
+
+from tree_sitter import Node, Parser
+from tree_sitter_language_pack import SupportedLanguage, get_parser, manifest_languages
+
+logger = getLogger(__name__)
+
+
+_TREE_SITTER_LANGUAGES: frozenset[str] = frozenset(manifest_languages())
+
+
+def is_supported_language(language: str) -> bool:
+    """Check if the language is supported by tree-sitter."""
+    return language in _TREE_SITTER_LANGUAGES
+
+
+@dataclass
+class ChunkBoundary:
+    """The output of the internal chunking algorithm."""
+
+    start: int
+    end: int
+
+
+@cache
+def _cached_get_parser(language: SupportedLanguage) -> Parser:
+    """Gets a parser from tree_sitter."""
+    return get_parser(language)
+
+
+def _merge_adjacent_chunks(
+    chunks: list[ChunkBoundary],
+    desired_length: int,
+) -> list[ChunkBoundary]:
+    """Merge adjacent chunks up to the desired length."""
+    merged = []
+
+    current_start = chunks[0].start
+    current_end = chunks[0].end
+    current_length = current_end - current_start
+
+    for group in chunks[1:]:
+        start, end = group.start, group.end
+        length = end - start
+
+        if current_length + length > desired_length:
+            merged.append(ChunkBoundary(start=current_start, end=current_end))
+            current_start = start
+            current_end = end
+            current_length = length
+            continue
+
+        current_end = end
+        current_length += length
+
+    merged.append(ChunkBoundary(start=current_start, end=current_end))
+
+    return merged
+
+
+def _merge_node_inner(node: Node, desired_length: int) -> list[ChunkBoundary]:
+    """Recursively merge and split nodes."""
+    # If there are no child nodes, the only thing we can do is return the current node.
+    if not node.children:
+        return [ChunkBoundary(node.start_byte, node.end_byte)]
+
+    groups: list[ChunkBoundary] = []
+    children = node.children
+    index = 0
+
+    while index < len(children):
+        child = children[index]
+        start = child.start_byte
+        end = child.end_byte
+        length = child.end_byte - child.start_byte
+
+        # Increment the pointer, as we accessed a child node.
+        index += 1
+        # If this single chunk is longer than the desired length
+        # we try to split it again.
+        if length > desired_length:
+            groups.extend(_merge_node_inner(child, desired_length))
+            continue
+
+        while index < len(children):
+            # Extend the current group with or more children, if they fit.
+            child = children[index]
+            child_length = child.end_byte - child.start_byte
+
+            if length + child_length > desired_length:
+                break
+
+            end = child.end_byte
+            length += child_length
+            index += 1
+
+        groups.append(ChunkBoundary(start, end))
+
+    return groups
+
+
+def _merge_node(node: Node, desired_length: int) -> list[ChunkBoundary]:
+    """Recursively turn nodes into chunks, then merge adjacent chunks."""
+    raw_chunks = _merge_node_inner(node, desired_length)
+    return _merge_adjacent_chunks(raw_chunks, desired_length)
+
+
+def chunk_lines(text: str, desired_length: int) -> list[ChunkBoundary]:
+    """Chunk source code by line."""
+    if not text.strip():
+        return []
+    lines_as_groups = []
+    index = 0
+    for line in text.splitlines(keepends=True):
+        lines_as_groups.append(ChunkBoundary(start=index, end=index + len(line)))
+        index += len(line)
+
+    return _merge_adjacent_chunks(lines_as_groups, desired_length)
+
+
+def chunk(text: str, language: str, desired_length: int) -> list[ChunkBoundary]:
+    """Chunk source code."""
+    if not text.strip():
+        return []
+
+    as_bytes = text.encode("utf-8")
+    parser = _cached_get_parser(language)
+    root = parser.parse(as_bytes).root_node
+
+    chunks = []
+    for chunk_boundary in _merge_node(root, desired_length):
+        start_char = len(as_bytes[: chunk_boundary.start].decode("utf-8"))
+        end_char = len(as_bytes[: chunk_boundary.end].decode("utf-8"))
+        chunks.append(ChunkBoundary(start=start_char, end=end_char))
+
+    return chunks

src/semble/index/chunker.py

@@ -4,7 +4,7 @@
 import bm25s
 from vicinity.backends.basic import BasicArgs
 
-from semble.index.chunker import chunk_source
+from semble.chunking import chunk_source
 from semble.index.dense import SelectableBasicBackend, embed_chunks
 from semble.index.file_walker import filter_extensions, language_for_path, walk_files
 from semble.index.sparse import enrich_for_bm25