contextual-engine 0.1.0__py3-none-any.whl

contextual/docs/chunker.py

@@ -0,0 +1,550 @@
+"""Heading-aware markdown/MDX/RST document chunker for Contextual Phase 2.
+
+Splits documents into semantically coherent chunks by heading hierarchy.
+Strategy:
+- Each H2 section is the primary chunk unit.
+- H3+ sub-sections are nested under their H2 parent.
+- Every chunk's ``embedding_text`` carries a breadcrumb prefix:
+  ``[File: path | Section: H1 > H2 > H3]\\nchunk content``
+- Code blocks <80 tokens stay inline; 80-600 become a parallel ``docs_code``
+  chunk; >600 become standalone with a placeholder in the parent chunk.
+
+Supported formats:
+  .md   — CommonMark + GFM, YAML/TOML frontmatter stripped
+  .mdx  — MDX: JSX/import/export lines stripped before markdown parse
+  .rst  — reStructuredText: docutils → pseudo-markdown heading tree
+  .txt  — Plain text: split by double-newline paragraphs
+
+Token accounting uses a fast 4-chars-per-token approximation consistent with
+jina-embeddings-v2-base-code (768-dim) which the embedder uses.
+"""
+from __future__ import annotations
+
+import re
+import unicodedata
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Generator
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# Token budget constants (tune here only)
+# ---------------------------------------------------------------------------
+CHUNK_TARGET_TOKENS: int = 450
+CHUNK_MAX_TOKENS: int = 800
+CHUNK_MIN_TOKENS: int = 120
+OVERLAP_TOKENS: int = 60          # trailing overlap appended to next chunk
+CODE_BLOCK_INLINE_MAX: int = 80   # tokens — keep inline in parent chunk
+CODE_BLOCK_SPLIT_MAX: int = 600   # tokens — emit parallel docs_code chunk; >600 → standalone
+
+
+# ---------------------------------------------------------------------------
+# Public data model
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class DocChunk:
+    """A single semantically coherent chunk from a document.
+
+    Fields that map directly to the ``chunks`` table columns:
+    - ``content``       raw text sent to BM25/FTS5 index
+    - ``embedding_text`` breadcrumb-prefixed text sent to the embedder
+    - ``heading_path``  e.g. "Installation > macOS > Homebrew"
+    - ``heading_level`` depth of the deepest heading in this chunk (1-6)
+    - ``chunk_type``    "docs" or "docs_code"
+    - ``start_line``    1-based line in source file
+    - ``end_line``      1-based line in source file (inclusive)
+    - ``chunk_index``   sequential position within the file (0-based)
+    - ``token_estimate`` fast approximation; not exact token count
+    """
+
+    file_path: str                          # repo-relative or absolute
+    content: str
+    embedding_text: str
+    heading_path: str                       # ">"-joined breadcrumb
+    heading_level: int                      # 1-6
+    chunk_type: str                         # "docs" | "docs_code"
+    start_line: int
+    end_line: int
+    chunk_index: int
+    token_estimate: int
+    frontmatter: dict = field(default_factory=dict)   # parsed YAML metadata
+
+
+# ---------------------------------------------------------------------------
+# Token estimation (fast, no tokeniser dependency)
+# ---------------------------------------------------------------------------
+
+def _est_tokens(text: str) -> int:
+    """Estimate token count at ~4 chars/token (standard LLM approximation)."""
+    return max(1, len(text) // 4)
+
+
+# ---------------------------------------------------------------------------
+# MDX pre-processor
+# ---------------------------------------------------------------------------
+# Patterns that are valid in MDX but crash markdown-it-py
+_MDX_IMPORT_EXPORT_RE = re.compile(
+    r"^(import|export)\s.*$",
+    re.MULTILINE,
+)
+_JSX_COMPONENT_RE = re.compile(
+    r"<([A-Z][A-Za-z0-9]*)(?:\s[^>]*)?\s*/?>.*?(?:</\1>)?",
+    re.DOTALL,
+)
+_MDX_EXPRESSION_RE = re.compile(r"\{[^}]*\}", re.DOTALL)
+
+
+def _strip_mdx(source: str) -> str:
+    """Strip MDX-specific syntax, leaving valid CommonMark behind."""
+    source = _MDX_IMPORT_EXPORT_RE.sub("", source)
+    source = _JSX_COMPONENT_RE.sub("", source)
+    source = _MDX_EXPRESSION_RE.sub("", source)
+    # Collapse multiple blank lines left by stripping
+    source = re.sub(r"\n{3,}", "\n\n", source)
+    return source
+
+
+# ---------------------------------------------------------------------------
+# RST → pseudo-markdown converter
+# ---------------------------------------------------------------------------
+# RST heading adornments in typical decoration order
+_RST_ADORNMENTS = r"[=\-`:'\"~^_*+#<>]"
+_RST_HEADING_RE = re.compile(
+    rf"^({_RST_ADORNMENTS}){{3,}}\n(.+)\n({_RST_ADORNMENTS}){{3,}}$|"
+    rf"^(.+)\n({_RST_ADORNMENTS}){{3,}}$",
+    re.MULTILINE,
+)
+
+# Map adornment char to heading level using order of first encounter
+def _rst_to_markdown(source: str) -> str:
+    """Convert RST heading syntax to markdown ATX headings (#, ##, ###, ...).
+
+    We do a two-pass: first collect all adornment characters in order of
+    appearance (determines hierarchy), then substitute.  Non-heading RST
+    directives are left as-is — they become fenced code blocks or plain text
+    when parsed by markdown-it-py.
+    """
+    adornment_order: list[str] = []
+    lines = source.splitlines(keepends=True)
+    out: list[str] = []
+    i = 0
+    while i < len(lines):
+        line = lines[i].rstrip("\n")
+        # Detect underline-only (section title above)
+        if i > 0 and re.match(rf"^{_RST_ADORNMENTS}{{3,}}$", line):
+            prev = lines[i - 1].rstrip("\n")
+            char = line[0]
+            if char not in adornment_order:
+                adornment_order.append(char)
+            level = min(adornment_order.index(char) + 1, 6)
+            # Replace the already-pushed previous line with an ATX heading
+            if out:
+                out[-1] = "#" * level + " " + prev + "\n"
+            i += 1
+            continue
+        out.append(lines[i])
+        i += 1
+    return "".join(out)
+
+
+# ---------------------------------------------------------------------------
+# Frontmatter stripper
+# ---------------------------------------------------------------------------
+_YAML_FM_RE = re.compile(r"^---\s*\n(.*?\n?)---\s*\n", re.DOTALL)
+_TOML_FM_RE = re.compile(r"^\+\+\+\s*\n(.*?\n?)\+\+\+\s*\n", re.DOTALL)
+
+
+def _strip_frontmatter(source: str) -> tuple[str, dict]:
+    """Strip YAML or TOML frontmatter and return (body, metadata_dict)."""
+    meta: dict = {}
+    for pattern in (_YAML_FM_RE, _TOML_FM_RE):
+        m = pattern.match(source)
+        if m:
+            raw = m.group(1)
+            try:
+                import yaml  # pyyaml ships with python-frontmatter
+                meta = yaml.safe_load(raw) or {}
+            except Exception:  # noqa: BLE001
+                pass
+            source = source[m.end():]
+            break
+    return source, meta
+
+
+# ---------------------------------------------------------------------------
+# Heading tree builder
+# ---------------------------------------------------------------------------
+
+@dataclass
+class _HeadingNode:
+    level: int                              # 1–6
+    title: str
+    start_line: int
+    end_line: int                           # set after next heading found
+    body_lines: list[str] = field(default_factory=list)
+    children: list["_HeadingNode"] = field(default_factory=list)
+
+
+_ATX_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*?)(?:\s+#+)?$")
+
+
+def _build_heading_tree(lines: list[str]) -> _HeadingNode:
+    """Parse ATX headings and build a heading tree from markdown lines.
+
+    Returns a virtual root node whose children are the top-level headings.
+    Lines before the first heading are collected under root.body_lines.
+    """
+    root = _HeadingNode(level=0, title="__root__", start_line=1, end_line=len(lines))
+    stack: list[_HeadingNode] = [root]
+
+    for lineno, raw in enumerate(lines, start=1):
+        m = _ATX_HEADING_RE.match(raw.rstrip("\n"))
+        if m:
+            level = len(m.group(1))
+            title = m.group(2).strip()
+            node = _HeadingNode(level=level, title=title, start_line=lineno, end_line=len(lines))
+
+            # Close nodes deeper than current level
+            while len(stack) > 1 and stack[-1].level >= level:
+                closed = stack.pop()
+                closed.end_line = lineno - 1
+
+            # Attach to parent
+            stack[-1].children.append(node)
+            stack.append(node)
+        else:
+            stack[-1].body_lines.append(raw)
+
+    return root
+
+
+def _node_breadcrumb(ancestors: list[_HeadingNode], node: _HeadingNode) -> str:
+    """Build a ' > '-joined breadcrumb from ancestor titles + this node."""
+    parts = [n.title for n in ancestors if n.level > 0] + [node.title]
+    return " > ".join(parts)
+
+
+def _collect_body(node: _HeadingNode, lines: list[str]) -> str:
+    """Collect the raw text body of a node (excluding child headings)."""
+    # Body is lines from node start to (first child start - 1) or node end
+    if node.children:
+        body_end = node.children[0].start_line - 1
+    else:
+        body_end = node.end_line
+    start = node.start_line        # start_line is the heading line itself
+    body_lines = lines[start:body_end]  # skip heading line (index start-1)
+    return "".join(body_lines).strip()
+
+
+# ---------------------------------------------------------------------------
+# Code block extractor
+# ---------------------------------------------------------------------------
+_FENCED_CODE_RE = re.compile(
+    r"(```(?:[^\n]*)?\n)(.*?)(```\s*$)",
+    re.DOTALL | re.MULTILINE,
+)
+
+
+def _split_code_blocks(
+    text: str,
+) -> tuple[str, list[tuple[str, str]]]:
+    """Extract large code blocks from text.
+
+    Returns:
+        (text_with_placeholders, [(placeholder, code_block_text), ...])
+    Large = > CODE_BLOCK_INLINE_MAX estimated tokens.
+    """
+    extracted: list[tuple[str, str]] = []
+    counter = [0]
+
+    def replacer(m: re.Match) -> str:
+        full_block = m.group(0)
+        tok = _est_tokens(full_block)
+        if tok <= CODE_BLOCK_INLINE_MAX:
+            return full_block   # keep inline
+        placeholder = f"[[CODE_BLOCK_{counter[0]}]]"
+        extracted.append((placeholder, full_block))
+        counter[0] += 1
+        return placeholder
+
+    result = _FENCED_CODE_RE.sub(replacer, text)
+    return result, extracted
+
+
+# ---------------------------------------------------------------------------
+# Paragraph splitter (for overlong sections)
+# ---------------------------------------------------------------------------
+
+def _split_by_paragraphs(
+    text: str,
+    max_tokens: int = CHUNK_MAX_TOKENS,
+    overlap_tokens: int = OVERLAP_TOKENS,
+) -> list[str]:
+    """Split overlong text into paragraph-aligned sub-chunks with overlap."""
+    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
+    if not paragraphs:
+        return [text] if text.strip() else []
+
+    chunks: list[str] = []
+    current_parts: list[str] = []
+    current_tokens = 0
+    overlap_tail = ""
+
+    for para in paragraphs:
+        ptok = _est_tokens(para)
+        if current_tokens + ptok > max_tokens and current_parts:
+            chunk_text = overlap_tail + "\n\n".join(current_parts)
+            chunks.append(chunk_text)
+            # Carry last paragraph as overlap seed
+            overlap_tail = current_parts[-1][:overlap_tokens * 4] + "\n\n"
+            current_parts = []
+            current_tokens = _est_tokens(overlap_tail)
+        current_parts.append(para)
+        current_tokens += ptok
+
+    if current_parts:
+        chunks.append(overlap_tail + "\n\n".join(current_parts))
+
+    return chunks if chunks else [text]
+
+
+# ---------------------------------------------------------------------------
+# Main chunker
+# ---------------------------------------------------------------------------
+
+class DocsChunker:
+    """Chunk a document file into ``DocChunk`` instances.
+
+    Usage::
+
+        chunker = DocsChunker()
+        chunks = chunker.chunk_file(Path("docs/README.md"))
+    """
+
+    SUPPORTED_EXTENSIONS = frozenset({".md", ".mdx", ".rst", ".txt"})
+
+    def __init__(
+        self,
+        target_tokens: int = CHUNK_TARGET_TOKENS,
+        max_tokens: int = CHUNK_MAX_TOKENS,
+        min_tokens: int = CHUNK_MIN_TOKENS,
+    ) -> None:
+        self._target = target_tokens
+        self._max = max_tokens
+        self._min = min_tokens
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def chunk_file(self, file_path: Path, repo_root: Path | None = None) -> list[DocChunk]:
+        """Parse and chunk a single document file.
+
+        Args:
+            file_path: Absolute path to the document.
+            repo_root: If provided, ``DocChunk.file_path`` is repo-relative.
+
+        Returns:
+            List of ``DocChunk`` objects (may be empty for blank files).
+        """
+        suffix = file_path.suffix.lower()
+        if suffix not in self.SUPPORTED_EXTENSIONS:
+            logger.debug("Skipping unsupported extension", path=str(file_path), suffix=suffix)
+            return []
+
+        try:
+            source = file_path.read_text(encoding="utf-8", errors="replace")
+        except OSError as exc:
+            logger.warning("Cannot read file", path=str(file_path), error=str(exc))
+            return []
+
+        if not source.strip():
+            return []
+
+        rel_path = str(file_path.relative_to(repo_root)) if repo_root else str(file_path)
+        return list(self._chunk_source(source, suffix, rel_path))
+
+    # ------------------------------------------------------------------
+    # Internal dispatch
+    # ------------------------------------------------------------------
+
+    def _chunk_source(
+        self, source: str, suffix: str, rel_path: str
+    ) -> Generator[DocChunk, None, None]:
+        """Dispatch to format-specific handler and yield DocChunk objects."""
+        if suffix == ".rst":
+            source = _rst_to_markdown(source)
+
+        if suffix == ".txt":
+            yield from self._chunk_plain_text(source, rel_path)
+            return
+
+        # Strip frontmatter for .md / .mdx / (converted .rst)
+        source, frontmatter = _strip_frontmatter(source)
+
+        if suffix == ".mdx":
+            source = _strip_mdx(source)
+
+        yield from self._chunk_markdown(source, rel_path, frontmatter)
+
+    # ------------------------------------------------------------------
+    # Markdown chunker
+    # ------------------------------------------------------------------
+
+    def _chunk_markdown(
+        self,
+        source: str,
+        rel_path: str,
+        frontmatter: dict,
+    ) -> Generator[DocChunk, None, None]:
+        """Main heading-aware chunker for CommonMark/MDX/RST (post-convert)."""
+        lines = source.splitlines(keepends=True)
+        root = _build_heading_tree(lines)
+
+        chunk_index = 0
+        file_label = Path(rel_path).name
+
+        # Pre-heading body (before any heading in the file)
+        if root.body_lines:
+            preamble = "".join(root.body_lines).strip()
+            if _est_tokens(preamble) >= self._min:
+                emb = f"[File: {file_label}]\n{preamble}"
+                yield DocChunk(
+                    file_path=rel_path,
+                    content=preamble,
+                    embedding_text=emb,
+                    heading_path="",
+                    heading_level=0,
+                    chunk_type="docs",
+                    start_line=1,
+                    end_line=root.children[0].start_line - 1 if root.children else len(lines),
+                    chunk_index=chunk_index,
+                    token_estimate=_est_tokens(preamble),
+                    frontmatter=frontmatter,
+                )
+                chunk_index += 1
+
+        # Walk heading tree depth-first
+        def _walk(
+            node: _HeadingNode, ancestors: list[_HeadingNode]
+        ) -> Generator[DocChunk, None, None]:
+            nonlocal chunk_index
+
+            breadcrumb = _node_breadcrumb(ancestors, node)
+            body = _collect_body(node, lines)
+
+            # --- Extract oversize code blocks from body ---
+            body_no_code, code_blocks = _split_code_blocks(body)
+
+            # --- Emit code chunks for large blocks ---
+            for placeholder, code_text in code_blocks:
+                tok = _est_tokens(code_text)
+                if tok <= CODE_BLOCK_SPLIT_MAX:
+                    # Parallel docs_code chunk
+                    emb = f"[File: {file_label} | Section: {breadcrumb}]\n{code_text}"
+                    yield DocChunk(
+                        file_path=rel_path,
+                        content=code_text,
+                        embedding_text=emb,
+                        heading_path=breadcrumb,
+                        heading_level=node.level,
+                        chunk_type="docs_code",
+                        start_line=node.start_line,
+                        end_line=node.end_line,
+                        chunk_index=chunk_index,
+                        token_estimate=tok,
+                        frontmatter=frontmatter,
+                    )
+                    chunk_index += 1
+                # else: standalone — placeholder stays in body_no_code; body
+                # will reference it but we don't emit the raw code separately
+                # (too large; user sees [CODE_BLOCK_N] marker in content)
+
+            # --- Split body_no_code into sub-chunks if overlong ---
+            if body_no_code.strip():
+                sub_texts = (
+                    _split_by_paragraphs(body_no_code, self._max)
+                    if _est_tokens(body_no_code) > self._max
+                    else [body_no_code]
+                )
+                for idx, sub in enumerate(sub_texts):
+                    sub = sub.strip()
+                    if not sub:
+                        continue
+                    tok = _est_tokens(sub)
+                    if tok < self._min and len(sub_texts) > 1:
+                        # Too small sub-chunk from a split — merge into prev
+                        # by not yielding; caller will see gap in chunk_index
+                        continue
+
+                    section_label = f"{breadcrumb} (part {idx+1})" if len(sub_texts) > 1 else breadcrumb
+                    emb = f"[File: {file_label} | Section: {section_label}]\n{sub}"
+                    yield DocChunk(
+                        file_path=rel_path,
+                        content=sub,
+                        embedding_text=emb,
+                        heading_path=breadcrumb,
+                        heading_level=node.level,
+                        chunk_type="docs",
+                        start_line=node.start_line,
+                        end_line=node.end_line,
+                        chunk_index=chunk_index,
+                        token_estimate=tok,
+                        frontmatter=frontmatter,
+                    )
+                    chunk_index += 1
+
+            # Recurse into children
+            for child in node.children:
+                yield from _walk(child, ancestors + [node])
+
+        for top_node in root.children:
+            yield from _walk(top_node, [])
+
+    # ------------------------------------------------------------------
+    # Plain text chunker
+    # ------------------------------------------------------------------
+
+    def _chunk_plain_text(
+        self, source: str, rel_path: str
+    ) -> Generator[DocChunk, None, None]:
+        """Paragraph-based chunker for .txt files (no heading structure)."""
+        file_label = Path(rel_path).name
+        sub_texts = _split_by_paragraphs(source, self._max)
+        for idx, sub in enumerate(sub_texts):
+            sub = sub.strip()
+            if not sub:
+                continue
+            tok = _est_tokens(sub)
+            if tok < self._min and len(sub_texts) > 1:
+                continue
+            emb = f"[File: {file_label}]\n{sub}"
+            yield DocChunk(
+                file_path=rel_path,
+                content=sub,
+                embedding_text=emb,
+                heading_path="",
+                heading_level=0,
+                chunk_type="docs",
+                start_line=1,
+                end_line=source.count("\n") + 1,
+                chunk_index=idx,
+                token_estimate=tok,
+                frontmatter={},
+            )
+
+
+# ---------------------------------------------------------------------------
+# Module-level convenience
+# ---------------------------------------------------------------------------
+
+def chunk_doc_file(
+    file_path: Path,
+    repo_root: Path | None = None,
+) -> list[DocChunk]:
+    """Convenience wrapper — chunk a single doc file with default settings."""
+    return DocsChunker().chunk_file(file_path, repo_root=repo_root)