dot-parser 2.0.0__py3-none-any.whl

dot_parser/chunking.py

@@ -0,0 +1,257 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+import re
+import statistics
+
+import semchunk
+
+from dot_parser.models import Chunk
+from dot_parser.tokens import count_tokens
+
+_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
+
+
+def _split_by_headings(markdown: str) -> list[dict]:
+    """Split markdown into sections by headings, tracking hierarchy via a stack."""
+    sections: list[dict] = []
+    stack: list[str] = []  # heading lines currently in scope
+
+    last_end = 0
+    for match in _HEADING_RE.finditer(markdown):
+        # Flush text before this heading as body of the previous section
+        if last_end < match.start():
+            body = markdown[last_end : match.start()].strip()
+            if body:
+                if sections:
+                    sections[-1]["body"] = body
+                else:
+                    # Text before any heading → root section
+                    sections.append({"heading": None, "section_path": list(stack), "body": body})
+
+        level = len(match.group(1))
+        heading_line = match.group(0).strip()
+
+        # Pop headings of level >= current
+        while stack and _heading_level(stack[-1]) >= level:
+            stack.pop()
+
+        stack.append(heading_line)
+
+        sections.append({"heading": heading_line, "section_path": list(stack), "body": ""})
+        last_end = match.end()
+
+    # Remaining text after the last heading
+    trailing = markdown[last_end:].strip()
+    if trailing:
+        if sections:
+            sections[-1]["body"] = trailing
+        else:
+            sections.append({"heading": None, "section_path": list(stack), "body": trailing})
+
+    return sections
+
+
+def _heading_level(heading_line: str) -> int:
+    m = re.match(r"^(#{1,6})\s", heading_line)
+    return len(m.group(1)) if m else 0
+
+
+def _split_oversized(section: dict, max_tokens: int) -> list[dict]:
+    """Split a section that exceeds max_tokens."""
+    heading = section["heading"]
+    body = section["body"]
+    path = section["section_path"]
+
+    full_text = f"{heading}\n\n{body}" if heading else body
+    if count_tokens(full_text) <= max_tokens:
+        return [section]
+
+    # Try paragraph splitting first
+    paragraphs = re.split(r"\n\n+", body)
+    if len(paragraphs) > 1:
+        sub_chunks = _group_paragraphs(paragraphs, max_tokens, heading)
+    else:
+        # Semantic splitting via semchunk
+        sub_chunks = semchunk.chunk(body, chunk_size=max_tokens, token_counter=count_tokens)
+
+    results = []
+    for i, text in enumerate(sub_chunks):
+        if i == 0 and heading:
+            results.append(
+                {
+                    "heading": heading,
+                    "section_path": list(path),
+                    "body": text,
+                }
+            )
+        else:
+            results.append(
+                {
+                    "heading": None,
+                    "section_path": list(path),
+                    "body": text,
+                }
+            )
+    return results
+
+
+def _group_paragraphs(paragraphs: list[str], max_tokens: int, heading: str | None) -> list[str]:
+    """Group paragraphs into chunks that fit within max_tokens."""
+    groups: list[str] = []
+    current: list[str] = []
+    current_tokens = 0
+
+    # Account for heading in first group
+    heading_tokens = count_tokens(heading + "\n\n") if heading else 0
+    first_group = True
+
+    for para in paragraphs:
+        para = para.strip()
+        if not para:
+            continue
+        para_tokens = count_tokens(para)
+        overhead = heading_tokens if first_group and not current else 0
+
+        if current and current_tokens + overhead + para_tokens > max_tokens:
+            groups.append("\n\n".join(current))
+            current = [para]
+            current_tokens = para_tokens
+            first_group = False
+        else:
+            current.append(para)
+            current_tokens += para_tokens
+
+    if current:
+        groups.append("\n\n".join(current))
+
+    return groups
+
+
+def _merge_small(sections: list[dict], max_tokens: int) -> list[dict]:
+    """Merge adjacent sections sharing the same parent if they fit."""
+    if not sections:
+        return sections
+
+    merged: list[dict] = [sections[0]]
+
+    for section in sections[1:]:
+        prev = merged[-1]
+
+        # Heading-only sections always start a new group
+        if section["heading"] and not section["body"]:
+            merged.append(section)
+            continue
+
+        # Check same parent (all but last element of section_path)
+        prev_parent = prev["section_path"][:-1] if prev["section_path"] else []
+        curr_parent = section["section_path"][:-1] if section["section_path"] else []
+
+        if prev_parent == curr_parent:
+            prev_text = _section_text(prev)
+            curr_text = _section_text(section)
+            combined_tokens = count_tokens(prev_text + "\n\n" + curr_text)
+
+            if combined_tokens <= max_tokens:
+                # Merge into previous
+                merged[-1] = {
+                    "heading": prev["heading"],
+                    "section_path": prev["section_path"],
+                    "body": (prev_text + "\n\n" + curr_text).strip(),
+                    "merged": True,
+                }
+                continue
+
+        merged.append(section)
+
+    return merged
+
+
+def _absorb_undersized(sections: list[dict], max_tokens: int) -> list[dict]:
+    """Absorb chunks smaller than 15% of max_tokens into neighbors."""
+    threshold = max_tokens * 0.15
+    result: list[dict] = []
+
+    for section in sections:
+        text = _section_text(section)
+        tokens = count_tokens(text)
+
+        if tokens < threshold and result:
+            # Append to previous
+            prev = result[-1]
+            prev_text = _section_text(prev)
+            result[-1] = {
+                "heading": prev["heading"],
+                "section_path": prev["section_path"],
+                "body": (prev_text + "\n\n" + text).strip(),
+                "merged": True,
+            }
+        else:
+            result.append(section)
+
+    # Check if first chunk is undersized and can be prepended to next
+    if len(result) >= 2:
+        first_text = _section_text(result[0])
+        if count_tokens(first_text) < threshold:
+            second = result[1]
+            second_text = _section_text(second)
+            result[1] = {
+                "heading": second["heading"],
+                "section_path": second["section_path"],
+                "body": (first_text + "\n\n" + second_text).strip(),
+                "merged": True,
+            }
+            result.pop(0)
+
+    return result
+
+
+def _section_text(section: dict) -> str:
+    if section.get("merged"):
+        return section["body"]
+    heading = section["heading"]
+    body = section["body"]
+    if heading and body:
+        return f"{heading}\n\n{body}"
+    return heading or body
+
+
+def _auto_max_tokens(sections: list[dict]) -> int:
+    """Compute max_tokens from 95th percentile of section sizes, clamped 64–512."""
+    sizes = [count_tokens(_section_text(s)) for s in sections if _section_text(s)]
+    if not sizes:
+        return 256
+    p95 = statistics.quantiles(sizes, n=20)[18] if len(sizes) >= 2 else sizes[0]
+    return max(64, min(512, int(p95)))
+
+
+def chunk(markdown: str, max_tokens: int | None = None, merge: bool = True) -> list[Chunk]:
+    if not markdown.strip():
+        return []
+
+    # Step 1: Split by headings
+    sections = _split_by_headings(markdown)
+
+    # Auto-compute max_tokens if not provided
+    if max_tokens is None:
+        max_tokens = _auto_max_tokens(sections)
+
+    # Step 2: Split oversized sections
+    split_sections: list[dict] = []
+    for section in sections:
+        split_sections.extend(_split_oversized(section, max_tokens))
+
+    # Steps 3 & 4: Merge and absorb (when merge=True)
+    if merge:
+        split_sections = _merge_small(split_sections, max_tokens)
+        split_sections = _absorb_undersized(split_sections, max_tokens)
+
+    # Build Chunk objects
+    return [
+        Chunk(
+            content=_section_text(s),
+            section_path=s["section_path"],
+        )
+        for s in split_sections
+        if _section_text(s)
+    ]