adaptive-oci-chunking 0.1.0__py3-none-any.whl

adaptive_chunking/__init__.py

@@ -0,0 +1,5 @@
+from adaptive_chunking.models import Chunk, ChunkingResult, Document
+from adaptive_chunking.pipeline import AdaptiveChunker
+
+__all__ = ["AdaptiveChunker", "Chunk", "ChunkingResult", "Document"]
+

adaptive_chunking/api.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from adaptive_chunking.pipeline import AdaptiveChunker
+
+try:
+    from fastapi import FastAPI
+except ImportError as exc:  # pragma: no cover
+    raise RuntimeError("Install API support with `pip install -e .[api]`.") from exc
+
+
+class ChunkRequest(BaseModel):
+    text: str = Field(min_length=1)
+    document_id: str = "document"
+
+
+class ChunkResponse(BaseModel):
+    document_id: str
+    strategy_name: str
+    score: float
+    chunks: list[dict]
+    metrics: list[dict]
+
+
+app = FastAPI(title="Adaptive OCI Chunking")
+chunker = AdaptiveChunker()
+
+
+@app.post("/chunk", response_model=ChunkResponse)
+def chunk(request: ChunkRequest) -> dict:
+    result = chunker.chunk(request.text, document_id=request.document_id)
+    return {
+        "document_id": result.document_id,
+        "strategy_name": result.strategy_name,
+        "score": result.score,
+        "chunks": [chunk.__dict__ for chunk in result.chunks],
+        "metrics": [metric.__dict__ for metric in result.metrics],
+    }
+

adaptive_chunking/chunkers.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+
+from adaptive_chunking.models import Chunk
+from adaptive_chunking.text import cosine_bow, normalize_space, sentences
+
+
+class BaseChunker(ABC):
+    name: str
+
+    @abstractmethod
+    def split(self, text: str) -> list[Chunk]:
+        raise NotImplementedError
+
+    def _build_chunks(self, spans: list[tuple[int, int]], text: str) -> list[Chunk]:
+        chunks: list[Chunk] = []
+        for start, end in spans:
+            chunk_text = normalize_space(text[start:end])
+            if chunk_text:
+                chunks.append(
+                    Chunk(
+                        text=chunk_text,
+                        index=len(chunks),
+                        start_char=start,
+                        end_char=end,
+                    )
+                )
+        return chunks
+
+
+class FixedWindowChunker(BaseChunker):
+    name = "fixed-window"
+
+    def __init__(self, chunk_size: int = 1200, overlap: int = 120) -> None:
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+        if overlap < 0 or overlap >= chunk_size:
+            raise ValueError("overlap must be non-negative and smaller than chunk_size")
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def split(self, text: str) -> list[Chunk]:
+        spans: list[tuple[int, int]] = []
+        step = self.chunk_size - self.overlap
+        for start in range(0, len(text), step):
+            end = min(start + self.chunk_size, len(text))
+            spans.append((start, end))
+            if end == len(text):
+                break
+        return self._build_chunks(spans, text)
+
+
+class SingleChunker(BaseChunker):
+    name = "single"
+
+    def split(self, text: str) -> list[Chunk]:
+        return self._build_chunks([(0, len(text))], text)
+
+
+class DelimiterChunker(BaseChunker):
+    name = "delimiter"
+
+    def __init__(
+        self,
+        delimiter: str = "\n---\n",
+        keep_delimiter: bool = False,
+        max_size: int = 1800,
+    ) -> None:
+        self.delimiter = delimiter
+        self.keep_delimiter = keep_delimiter
+        self.max_size = max_size
+        self.fallback = RecursiveChunker(chunk_size=max_size)
+
+    def split(self, text: str) -> list[Chunk]:
+        if not self.delimiter or self.delimiter not in text:
+            return self.fallback.split(text)
+        spans: list[tuple[int, int]] = []
+        cursor = 0
+        while cursor < len(text):
+            split_at = text.find(self.delimiter, cursor)
+            if split_at < 0:
+                spans.append((cursor, len(text)))
+                break
+            end = split_at + len(self.delimiter) if self.keep_delimiter else split_at
+            spans.append((cursor, end))
+            cursor = split_at + len(self.delimiter)
+        return self._split_oversized_spans(spans, text)
+
+    def _split_oversized_spans(self, spans: list[tuple[int, int]], text: str) -> list[Chunk]:
+        chunks: list[Chunk] = []
+        for start, end in spans:
+            if end - start <= self.max_size:
+                chunks.extend(self._build_chunks([(start, end)], text))
+                continue
+            for chunk in self.fallback.split(text[start:end]):
+                chunks.append(
+                    Chunk(chunk.text, len(chunks), start + chunk.start_char, start + chunk.end_char)
+                )
+        return chunks
+
+
+class PageChunker(DelimiterChunker):
+    name = "page"
+
+    def __init__(self, page_delimiter: str = "\f", max_size: int = 2200) -> None:
+        super().__init__(delimiter=page_delimiter, keep_delimiter=False, max_size=max_size)
+
+
+class RecursiveChunker(BaseChunker):
+    name = "recursive"
+
+    def __init__(self, chunk_size: int = 1200, separators: tuple[str, ...] | None = None) -> None:
+        self.chunk_size = chunk_size
+        self.separators = separators or ("\n\n", "\n", ". ", " ")
+
+    def split(self, text: str) -> list[Chunk]:
+        spans = self._split_span(text, 0, len(text), 0)
+        return self._build_chunks(spans, text)
+
+    def _split_span(
+        self,
+        text: str,
+        start: int,
+        end: int,
+        separator_index: int,
+    ) -> list[tuple[int, int]]:
+        if end - start <= self.chunk_size:
+            return [(start, end)]
+        if separator_index >= len(self.separators):
+            return [
+                (cursor, min(cursor + self.chunk_size, end))
+                for cursor in range(start, end, self.chunk_size)
+            ]
+
+        separator = self.separators[separator_index]
+        pieces: list[tuple[int, int]] = []
+        cursor = start
+        while cursor < end:
+            limit = min(cursor + self.chunk_size, end)
+            split_at = text.rfind(separator, cursor, limit)
+            if split_at <= cursor:
+                pieces.extend(self._split_span(text, cursor, limit, separator_index + 1))
+                cursor = limit
+            else:
+                split_end = split_at + len(separator)
+                pieces.append((cursor, split_end))
+                cursor = split_end
+        return pieces
+
+
+class SplitThenMergeChunker(BaseChunker):
+    name = "split-then-merge"
+
+    def __init__(self, min_size: int = 600, max_size: int = 1400) -> None:
+        if min_size <= 0 or max_size <= min_size:
+            raise ValueError("expected 0 < min_size < max_size")
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def split(self, text: str) -> list[Chunk]:
+        raw_spans = _paragraph_spans(text)
+        if not raw_spans:
+            return []
+        merged: list[tuple[int, int]] = []
+        start, end = raw_spans[0]
+        for next_start, next_end in raw_spans[1:]:
+            proposed_size = next_end - start
+            if proposed_size <= self.max_size or end - start < self.min_size:
+                end = next_end
+            else:
+                merged.append((start, end))
+                start, end = next_start, next_end
+        merged.append((start, end))
+        return self._build_chunks(merged, text)
+
+
+class SectionAwareChunker(BaseChunker):
+    name = "section-aware"
+
+    def __init__(self, min_size: int = 500, max_size: int = 1800) -> None:
+        self.min_size = min_size
+        self.max_size = max_size
+        self.fallback = SplitThenMergeChunker(min_size=min_size, max_size=max_size)
+        self.heading_pattern = re.compile(
+            r"(?m)^(?:#{1,6}\s+.+|\d+(?:\.\d+)*\s+[A-Z].+|[A-Z][A-Za-z0-9 ,:;&()/-]{3,80})$"
+        )
+
+    def split(self, text: str) -> list[Chunk]:
+        starts = sorted(
+            {0, len(text), *(match.start() for match in self.heading_pattern.finditer(text))}
+        )
+        spans = [(starts[index], starts[index + 1]) for index in range(len(starts) - 1)]
+        merged: list[tuple[int, int]] = []
+        current_start: int | None = None
+        current_end: int | None = None
+        for start, end in spans:
+            if not text[start:end].strip():
+                continue
+            if current_start is None:
+                current_start, current_end = start, end
+                continue
+            assert current_end is not None
+            proposed = end - current_start
+            if proposed <= self.max_size or current_end - current_start < self.min_size:
+                current_end = end
+            else:
+                merged.append((current_start, current_end))
+                current_start, current_end = start, end
+        if current_start is not None and current_end is not None:
+            merged.append((current_start, current_end))
+        chunks: list[Chunk] = []
+        for start, end in merged:
+            if end - start <= self.max_size:
+                chunks.extend(self._build_chunks([(start, end)], text))
+            else:
+                for chunk in self.fallback.split(text[start:end]):
+                    chunks.append(
+                        Chunk(
+                            chunk.text,
+                            len(chunks),
+                            start + chunk.start_char,
+                            start + chunk.end_char,
+                        )
+                    )
+        return chunks
+
+
+class SemanticChunker(BaseChunker):
+    name = "semantic"
+
+    def __init__(
+        self,
+        max_size: int = 1400,
+        min_size: int = 350,
+        similarity_threshold: float = 0.10,
+    ) -> None:
+        self.max_size = max_size
+        self.min_size = min_size
+        self.similarity_threshold = similarity_threshold
+
+    def split(self, text: str) -> list[Chunk]:
+        sentence_spans = _sentence_spans(text)
+        if not sentence_spans:
+            return []
+        chunks: list[tuple[int, int]] = []
+        start, end = sentence_spans[0]
+        previous_text = text[start:end]
+        for next_start, next_end in sentence_spans[1:]:
+            next_text = text[next_start:next_end]
+            proposed_size = next_end - start
+            similarity = cosine_bow(previous_text, next_text)
+            should_break = (
+                end - start >= self.min_size
+                and (proposed_size > self.max_size or similarity < self.similarity_threshold)
+            )
+            if should_break:
+                chunks.append((start, end))
+                start, end = next_start, next_end
+            else:
+                end = next_end
+            previous_text = next_text
+        chunks.append((start, end))
+        return self._build_chunks(chunks, text)
+
+
+class RegexSectionChunker(BaseChunker):
+    name = "regex-section"
+
+    def __init__(self, max_size: int = 1800, heading_pattern: str | None = None) -> None:
+        self.max_size = max_size
+        self.heading_pattern = re.compile(
+            heading_pattern or r"(?m)^(?:#{1,6}\s+.+|\d+(?:\.\d+)*\s+[A-Z].+)$"
+        )
+        self.fallback = RecursiveChunker(chunk_size=max_size)
+
+    def split(self, text: str) -> list[Chunk]:
+        starts = [match.start() for match in self.heading_pattern.finditer(text)]
+        if not starts:
+            return self.fallback.split(text)
+        starts = sorted(set([0, *starts, len(text)]))
+        spans = [(starts[index], starts[index + 1]) for index in range(len(starts) - 1)]
+        chunks: list[Chunk] = []
+        for start, end in spans:
+            section = text[start:end]
+            if len(section) <= self.max_size:
+                chunks.extend(self._build_chunks([(start, end)], text))
+            else:
+                for chunk in self.fallback.split(section):
+                    chunks.append(
+                        Chunk(
+                            text=chunk.text,
+                            index=len(chunks),
+                            start_char=start + chunk.start_char,
+                            end_char=start + chunk.end_char,
+                        )
+                    )
+        return [
+            Chunk(c.text, i, c.start_char, c.end_char, c.metadata)
+            for i, c in enumerate(chunks)
+        ]
+
+
+def default_chunkers() -> list[BaseChunker]:
+    return [
+        SingleChunker(),
+        FixedWindowChunker(),
+        RecursiveChunker(),
+        SplitThenMergeChunker(),
+        SectionAwareChunker(),
+        DelimiterChunker(),
+        PageChunker(),
+        SemanticChunker(),
+        RegexSectionChunker(),
+    ]
+
+
+def _paragraph_spans(text: str) -> list[tuple[int, int]]:
+    spans: list[tuple[int, int]] = []
+    cursor = 0
+    for match in re.finditer(r"\n\s*\n", text):
+        end = match.end()
+        if text[cursor:end].strip():
+            spans.append((cursor, end))
+        cursor = end
+    if text[cursor:].strip():
+        spans.append((cursor, len(text)))
+    return spans or ([(0, len(text))] if text.strip() else [])
+
+
+def _sentence_spans(text: str) -> list[tuple[int, int]]:
+    spans: list[tuple[int, int]] = []
+    cursor = 0
+    for sentence in sentences(text):
+        start = text.find(sentence, cursor)
+        if start < 0:
+            continue
+        end = start + len(sentence)
+        spans.append((start, end))
+        cursor = end
+    return spans

adaptive_chunking/cli.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from adaptive_chunking.io import load_text_file
+from adaptive_chunking.pipeline import AdaptiveChunker
+
+app = typer.Typer(help="Adaptive document chunking for RAG.")
+console = Console()
+
+
+@app.command()
+def chunk(
+    path: Annotated[
+        Path,
+        typer.Argument(exists=True, readable=True, help="Text or Markdown file to chunk."),
+    ],
+    document_id: Annotated[
+        str | None,
+        typer.Option(help="Stable document identifier."),
+    ] = None,
+    json_output: Annotated[
+        bool,
+        typer.Option("--json", help="Print machine-readable JSON."),
+    ] = False,
+) -> None:
+    text = load_text_file(path)
+    result = AdaptiveChunker().chunk(text, document_id=document_id or path.stem)
+    if json_output:
+        console.print(
+            json.dumps(
+                {
+                    "document_id": result.document_id,
+                    "strategy_name": result.strategy_name,
+                    "score": result.score,
+                    "chunks": [chunk.__dict__ for chunk in result.chunks],
+                    "metrics": [metric.__dict__ for metric in result.metrics],
+                },
+                indent=2,
+            )
+        )
+        return
+
+    console.print(f"[bold]Strategy:[/bold] {result.strategy_name}")
+    console.print(f"[bold]Score:[/bold] {result.score:.3f}")
+    table = Table("Metric", "Value", "Weight")
+    for metric in result.metrics:
+        table.add_row(metric.name, f"{metric.value:.3f}", f"{metric.weight:.2f}")
+    console.print(table)
+    for chunk_item in result.chunks:
+        console.rule(f"Chunk {chunk_item.index}")
+        console.print(chunk_item.text)
+
+
+if __name__ == "__main__":
+    app()

adaptive_chunking/io.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+SUPPORTED_TEXT_SUFFIXES = {".txt", ".md", ".markdown", ".rst"}
+
+
+def load_text_file(path: str | Path) -> str:
+    file_path = Path(path)
+    if file_path.suffix.lower() not in SUPPORTED_TEXT_SUFFIXES:
+        raise ValueError(f"unsupported file type: {file_path.suffix}")
+    return file_path.read_text(encoding="utf-8")
+
+
+def discover_text_files(path: str | Path) -> list[Path]:
+    root = Path(path)
+    if root.is_file():
+        return [root]
+    return sorted(
+        child
+        for child in root.rglob("*")
+        if child.is_file() and child.suffix.lower() in SUPPORTED_TEXT_SUFFIXES
+    )
+

adaptive_chunking/langchain.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Any
+
+from adaptive_chunking.pipeline import AdaptiveChunker
+
+try:
+    from langchain_text_splitters import TextSplitter
+except ImportError:  # pragma: no cover
+    TextSplitter = None  # type: ignore[assignment]
+
+_BaseTextSplitter = TextSplitter if TextSplitter is not None else object
+
+
+class LangChainAdaptiveTextSplitter(_BaseTextSplitter):  # type: ignore[misc,valid-type]
+    """LangChain TextSplitter backed by AdaptiveChunker."""
+
+    def __init__(
+        self,
+        chunker: AdaptiveChunker | None = None,
+        keep_separator: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        if TextSplitter is None:  # pragma: no cover
+            raise RuntimeError(
+                "Install LangChain support with `pip install -e .[langchain]`."
+            )
+        super().__init__(keep_separator=keep_separator, **kwargs)
+        self.chunker = chunker or AdaptiveChunker()
+
+    def split_text(self, text: str) -> list[str]:
+        result = self.chunker.chunk(text)
+        return [chunk.text for chunk in result.chunks]

adaptive_chunking/llama_index.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import Any
+
+from adaptive_chunking.models import Chunk, ChunkingResult
+from adaptive_chunking.pipeline import AdaptiveChunker
+
+
+def chunks_to_llama_nodes(
+    chunks: list[Chunk],
+    *,
+    document_id: str = "document",
+    extra_metadata: dict[str, Any] | None = None,
+) -> list[Any]:
+    try:
+        from llama_index.core.schema import TextNode
+    except ImportError as exc:  # pragma: no cover
+        raise RuntimeError(
+            "Install LlamaIndex support with `pip install -e .[llama-index]`."
+        ) from exc
+
+    metadata = extra_metadata or {}
+    return [
+        TextNode(
+            text=chunk.text,
+            id_=f"{document_id}:{chunk.index}",
+            metadata={
+                **metadata,
+                **chunk.metadata,
+                "document_id": document_id,
+                "chunk_index": chunk.index,
+                "start_char": chunk.start_char,
+                "end_char": chunk.end_char,
+            },
+        )
+        for chunk in chunks
+    ]
+
+
+def result_to_llama_nodes(result: ChunkingResult) -> list[Any]:
+    return chunks_to_llama_nodes(
+        result.chunks,
+        document_id=result.document_id,
+        extra_metadata={"strategy_name": result.strategy_name, "adaptive_score": result.score},
+    )
+
+
+class LlamaIndexAdaptiveParser:
+    """Small adapter with the same practical behavior as a LlamaIndex node parser."""
+
+    def __init__(self, chunker: AdaptiveChunker | None = None) -> None:
+        self.chunker = chunker or AdaptiveChunker()
+
+    def get_nodes_from_documents(self, documents: list[Any], **_: Any) -> list[Any]:
+        nodes: list[Any] = []
+        for document_index, document in enumerate(documents):
+            text = getattr(document, "text", None) or getattr(document, "get_content", lambda: "")()
+            metadata = dict(getattr(document, "metadata", {}) or {})
+            document_id = str(
+                metadata.get("document_id") or getattr(document, "id_", document_index)
+            )
+            result = self.chunker.chunk(text, document_id=document_id)
+            nodes.extend(
+                chunks_to_llama_nodes(
+                    result.chunks,
+                    document_id=document_id,
+                    extra_metadata={
+                        **metadata,
+                        "strategy_name": result.strategy_name,
+                        "adaptive_score": result.score,
+                    },
+                )
+            )
+        return nodes