mdkeychunker 0.3.0__py3-none-any.whl

mdkeychunker/__init__.py

@@ -0,0 +1,7 @@
+"""MDKeyChunker — Markdown chunking with LLM enrichment for RAG."""
+__version__ = "0.3.0"
+from .pipeline import Pipeline
+from .models import Chunk
+from .config import Config
+
+__all__ = ["Pipeline", "Chunk", "Config"]

mdkeychunker/chunker.py

@@ -0,0 +1,402 @@
+"""
+Markdown structural chunker.
+
+Splits Markdown documents following structure (headers, lists, tables, code blocks)
+without breaking semantic units.
+"""
+
+import re
+from typing import List, Tuple, Optional
+from dataclasses import dataclass
+
+from .models import Chunk
+from .config import Config
+
+
+_THEMATIC_BREAK = re.compile(r"^ {0,3}(?:(?:\*[ \t]*){3,}|(?:-[ \t]*){3,}|(?:_[ \t]*){3,})$")
+_FENCE_OPEN = re.compile(r"^ {0,3}(`{3,}|~{3,})(.*)$")
+
+
+def _is_thematic_break(line: str) -> bool:
+    return bool(_THEMATIC_BREAK.match(line))
+
+
+def _fence_open(line: str) -> Optional[re.Match]:
+    """A fence opener; a backtick fence's info string may not contain backticks."""
+    m = _FENCE_OPEN.match(line)
+    if m and m.group(1)[0] == "`" and "`" in m.group(2):
+        return None
+    return m
+
+
+@dataclass
+class _Block:
+    """Internal structural block in Markdown."""
+    type: str  # header, paragraph, list, code, table, blockquote, yaml_front_matter
+    content: str
+    start_line: int
+    end_line: int
+    level: int = 0  # For headers
+    section_title: str = ""
+
+
+class MarkdownChunker:
+    """
+    Chunks Markdown documents following structural boundaries.
+
+    Never splits:
+    - Code blocks
+    - Tables
+    - YAML front matter
+    - List items mid-item
+    - Blockquotes mid-paragraph
+    """
+
+    def __init__(self, config: Config):
+        self.config = config
+        self.min_chunk_size = config.min_chunk_size
+        self.max_chunk_size = config.max_chunk_size
+        self._header_stack: List[Tuple[int, str]] = []
+
+    def chunk(self, markdown_text: str) -> List[Chunk]:
+        """Chunk a Markdown document into structural units."""
+        blocks = self._parse_blocks(markdown_text)
+        chunks = self._group_blocks(blocks)
+        return chunks
+
+    # ------------------------------------------------------------------ #
+    # Parsing                                                              #
+    # ------------------------------------------------------------------ #
+
+    def _parse_blocks(self, text: str) -> List[_Block]:
+        """Parse Markdown into structural blocks."""
+        blocks: List[_Block] = []
+        lines = text.replace("\r\n", "\n").replace("\r", "\n").split("\n")
+        i = 0
+        current_section = ""
+        self._header_stack = []
+
+        while i < len(lines):
+            line = lines[i]
+
+            # YAML front matter
+            if i == 0 and line.lstrip("\ufeff").strip() == "---":
+                block, consumed = self._parse_yaml_front_matter(lines, i)
+                if block:
+                    blocks.append(block)
+                    i += consumed
+                    continue
+
+            # Fenced code blocks (a backtick info string may not contain backticks)
+            if _fence_open(line):
+                block, consumed = self._parse_code_block(lines, i)
+                block.section_title = current_section
+                blocks.append(block)
+                i += consumed
+                continue
+
+            # Headers
+            header_match = re.match(r"^(#{1,6})\s+(.+?)(?:\s+#+)?\s*$", line)
+            setext_match = (
+                not header_match
+                and line.strip()
+                and i + 1 < len(lines)
+                and re.match(r"^ {0,3}(=+|-+)\s*$", lines[i + 1])
+                and not _is_thematic_break(line)
+                and not self._is_list_line(line)
+                and not line.strip().startswith(">")
+            )
+            if header_match or setext_match:
+                if header_match:
+                    level = len(header_match.group(1))
+                    title = header_match.group(2).strip()
+                    content, end = line, i
+                else:
+                    level = 1 if lines[i + 1].strip().startswith("=") else 2
+                    title = line.strip()
+                    content, end = line + "\n" + lines[i + 1], i + 1
+                self._update_header_stack(level, title)
+                current_section = self._get_current_section()
+                blocks.append(_Block(
+                    type="header",
+                    content=content,
+                    start_line=i,
+                    end_line=end,
+                    level=level,
+                    section_title=current_section,
+                ))
+                i = end + 1
+                continue
+
+            # Tables
+            if self._is_table_line(line, lines, i):
+                block, consumed = self._parse_table(lines, i)
+                block.section_title = current_section
+                blocks.append(block)
+                i += consumed
+                continue
+
+            # Lists
+            if self._is_list_line(line):
+                block, consumed = self._parse_list(lines, i)
+                block.section_title = current_section
+                blocks.append(block)
+                i += consumed
+                continue
+
+            # Blockquotes
+            if line.strip().startswith(">"):
+                block, consumed = self._parse_blockquote(lines, i)
+                block.section_title = current_section
+                blocks.append(block)
+                i += consumed
+                continue
+
+            # Horizontal rules
+            if _is_thematic_break(line):
+                blocks.append(_Block(
+                    type="horizontal_rule",
+                    content=line,
+                    start_line=i,
+                    end_line=i,
+                    section_title=current_section,
+                ))
+                i += 1
+                continue
+
+            # Paragraphs
+            if line.strip():
+                block, consumed = self._parse_paragraph(lines, i)
+                block.section_title = current_section
+                blocks.append(block)
+                i += consumed
+                continue
+
+            i += 1  # empty line
+
+        return blocks
+
+    # ------------------------------------------------------------------ #
+    # Block parsers                                                        #
+    # ------------------------------------------------------------------ #
+
+    def _parse_yaml_front_matter(
+        self, lines: List[str], start: int
+    ) -> Tuple[Optional[_Block], int]:
+        end = start + 1
+        while end < len(lines) and lines[end].strip() != "---":
+            end += 1
+        if end >= len(lines):
+            return None, 0
+        content = "\n".join(lines[start: end + 1])
+        return _Block(
+            type="yaml_front_matter",
+            content=content,
+            start_line=start,
+            end_line=end,
+        ), end - start + 1
+
+    def _parse_code_block(self, lines: List[str], start: int) -> Tuple[_Block, int]:
+        opener = lines[start].strip()
+        fence_char = opener[0]
+        fence_len = len(opener) - len(opener.lstrip(fence_char))
+        # CommonMark: closes only on a line of the same char, at least as long,
+        # indented at most 3 spaces, with nothing else on the line
+        closer = re.compile(rf"^ {{0,3}}{re.escape(fence_char)}{{{fence_len},}}[ \t]*$")
+        end = start + 1
+        while end < len(lines):
+            if closer.match(lines[end]):
+                end += 1
+                break
+            end += 1
+        content = "\n".join(lines[start:end])
+        return _Block(
+            type="code", content=content, start_line=start, end_line=end - 1
+        ), end - start
+
+    def _is_table_line(self, line: str, lines: list[str], idx: int) -> bool:
+        """Check if this line starts a table (header row followed by separator)."""
+        if "|" not in line:
+            return False
+        # Must have a separator line immediately after — only detect from the header row
+        if idx + 1 < len(lines):
+            return bool(re.match(r"^\s*\|?[\s:-]+\|[\s|:-]+\|?\s*$", lines[idx + 1]))
+        return False
+
+    def _parse_table(self, lines: List[str], start: int) -> Tuple[_Block, int]:
+        end = start
+        while end < len(lines) and ("|" in lines[end] and lines[end].strip()):
+            end += 1
+        content = "\n".join(lines[start:end])
+        return _Block(
+            type="table", content=content, start_line=start, end_line=end - 1
+        ), end - start
+
+    def _is_list_line(self, line: str) -> bool:
+        stripped = line.strip()
+        return bool(
+            re.match(r"^[\*\-\+]\s+",
+                     stripped) or re.match(r"^\d+[.)]\s+", stripped)
+        )
+
+    def _parse_list(self, lines: List[str], start: int) -> Tuple[_Block, int]:
+        end = start
+        base_indent = len(lines[start]) - len(lines[start].lstrip())
+        end += 1
+        while end < len(lines):
+            line = lines[end]
+            if not line.strip():
+                peek = end + 1
+                while peek < len(lines) and not lines[peek].strip():
+                    peek += 1
+                if peek < len(lines) and (
+                    self._is_list_line(lines[peek])
+                    or len(lines[peek]) - len(lines[peek].lstrip()) > base_indent
+                ):
+                    end += 1
+                    continue
+                break
+            indent = len(line) - len(line.lstrip())
+            if self._is_list_line(line) or indent > base_indent:
+                end += 1
+            else:
+                break
+        content = "\n".join(lines[start:end])
+        return _Block(
+            type="list", content=content, start_line=start, end_line=end - 1
+        ), end - start
+
+    def _parse_blockquote(self, lines: List[str], start: int) -> Tuple[_Block, int]:
+        end = start
+        while end < len(lines) and lines[end].strip().startswith(">"):
+            end += 1
+        content = "\n".join(lines[start:end])
+        return _Block(
+            type="blockquote", content=content, start_line=start, end_line=end - 1
+        ), end - start
+
+    def _parse_paragraph(self, lines: List[str], start: int) -> Tuple[_Block, int]:
+        end = start
+        while end < len(lines):
+            line = lines[end]
+            if not line.strip():
+                break
+            if end > start and (
+                line.strip().startswith("#")
+                or _fence_open(line)
+                or self._is_list_line(line)
+                or line.strip().startswith(">")
+                or "|" in line
+            ):
+                break
+            end += 1
+        content = "\n".join(lines[start:end])
+        return _Block(
+            type="paragraph", content=content, start_line=start, end_line=end - 1
+        ), max(end - start, 1)
+
+    # ------------------------------------------------------------------ #
+    # Header stack                                                         #
+    # ------------------------------------------------------------------ #
+
+    def _update_header_stack(self, level: int, title: str) -> None:
+        self._header_stack = [(l, t)
+                              for l, t in self._header_stack if l < level]
+        self._header_stack.append((level, title))
+
+    def _get_current_section(self) -> str:
+        return " > ".join(t for _, t in self._header_stack)
+
+    # ------------------------------------------------------------------ #
+    # Grouping                                                             #
+    # ------------------------------------------------------------------ #
+
+    def _group_blocks(self, blocks: List[_Block]) -> List[Chunk]:
+        """Group blocks into chunks respecting structure and size.
+
+        Rules: every header starts a new chunk once the current chunk has body
+        text; a chunk holding only headers is never flushed for size, so headers
+        always travel with their content; atomic blocks are never split.
+        """
+        chunks: List[Chunk] = []
+        current: List[_Block] = []
+        current_size = 0
+        has_body = False
+        ATOMIC = {"code", "table", "yaml_front_matter", "blockquote"}
+
+        def flush() -> None:
+            nonlocal current, current_size, has_body
+            chunks.append(self._make_chunk(current))
+            current, current_size, has_body = [], 0, False
+
+        for block in blocks:
+            bsize = len(block.content)
+            if block.type == "header":
+                if has_body:
+                    flush()
+            elif has_body and current_size + bsize > self.max_chunk_size:
+                flush()
+
+            current.append(block)
+            current_size += bsize
+            if block.type != "header":
+                has_body = True
+                if block.type in ATOMIC and current_size >= self.max_chunk_size:
+                    flush()
+
+        if current:
+            tail = self._make_chunk(current)
+            if not has_body and chunks:
+                # Trailing headers with no content (e.g. an empty last section)
+                chunks[-1] = self._join(chunks[-1], tail)
+            else:
+                chunks.append(tail)
+
+        return self._merge_small_chunks(chunks)
+
+    def _make_chunk(self, blocks: List[_Block]) -> Chunk:
+        text = "\n\n".join(b.content for b in blocks)
+        # Section of the content; preamble or front matter falls back to the first header's
+        section_title = next(
+            (b.section_title for b in blocks if b.type != "header" and b.section_title),
+            next((b.section_title for b in blocks if b.section_title), ""),
+        )
+        content_types = list(set(b.type for b in blocks))
+        return Chunk(
+            text=text,
+            section_title=section_title,
+            start_line=blocks[0].start_line,
+            end_line=blocks[-1].end_line,
+            content_types=content_types,
+        )
+
+    @staticmethod
+    def _join(a: Chunk, b: Chunk) -> Chunk:
+        """Concatenate two adjacent chunks; the larger part names the section."""
+        major = a if len(a.text) >= len(b.text) else b
+        return Chunk(
+            text=a.text + "\n\n" + b.text,
+            section_title=major.section_title or a.section_title or b.section_title,
+            start_line=a.start_line,
+            end_line=b.end_line,
+            content_types=list(set(a.content_types + b.content_types)),
+        )
+
+    def _merge_small_chunks(self, chunks: List[Chunk]) -> List[Chunk]:
+        """Merge chunks below min_chunk_size forward (repeatedly), capped at 2x max."""
+        cap = self.max_chunk_size * 2
+        merged: List[Chunk] = []
+        for c in chunks:
+            if not c.text.strip():
+                continue
+            if merged and len(merged[-1].text) < self.min_chunk_size \
+                    and len(merged[-1].text) + len(c.text) <= cap:
+                merged[-1] = self._join(merged[-1], c)
+            else:
+                merged.append(c)
+        # A too-small final chunk joins the previous one, within the same cap
+        if len(merged) >= 2 and len(merged[-1].text) < self.min_chunk_size \
+                and len(merged[-2].text) + len(merged[-1].text) <= cap:
+            last = merged.pop()
+            merged[-1] = self._join(merged[-1], last)
+        return merged

mdkeychunker/cli.py

@@ -0,0 +1,63 @@
+"""CLI for MDKeyChunker."""
+import argparse
+import sys
+from pathlib import Path
+from .config import Config
+from .pipeline import Pipeline
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(
+        description="MDKeyChunker — Markdown chunking for RAG")
+    p.add_argument("input", help="Input Markdown file")
+    p.add_argument("-o", "--output", help="Output JSONL file")
+    p.add_argument("--env", help="Path to .env file")
+    p.add_argument("--stats", action="store_true", help="Print statistics")
+    p.add_argument("--no-merge", action="store_true",
+                   help="Disable key-based chunk merging")
+    p.add_argument("--summary", help="Output summary text file")
+    p.add_argument("--provider", choices=["openai", "anthropic", "openai_compatible"],
+                   help="LLM provider (overrides .env)")
+    p.add_argument("--model", help="LLM model name (overrides .env)")
+    p.add_argument(
+        "--base-url", help="LLM base URL for Ollama/vLLM (overrides .env)")
+    args = p.parse_args()
+
+    input_path = Path(args.input)
+    if not input_path.exists():
+        print(f"Error: {args.input} not found", file=sys.stderr)
+        return 1
+
+    config = Config.from_env(args.env)
+
+    if args.no_merge:
+        config.merge_by_keys = False
+    if args.provider:
+        config.llm_provider = args.provider
+    if args.model:
+        config.llm_model = args.model
+    if args.base_url:
+        config.llm_base_url = args.base_url
+
+    pipeline = Pipeline(config)
+    chunks = pipeline.process_file(str(input_path))
+
+    out = args.output or str(input_path.with_suffix(".jsonl"))
+    pipeline.save_jsonl(chunks, out)
+    print(f"✓ {len(chunks)} chunks → {out}")
+
+    if args.summary:
+        pipeline.save_summary(chunks, args.summary)
+        print(f"✓ Summary → {args.summary}")
+
+    if args.stats:
+        print(f"  Keys: {len(set(c.key for c in chunks if c.key))}")
+        print(f"  Entities: {sum(len(c.entities) for c in chunks)}")
+        print(
+            f"  Avg tokens: {sum(c.token_count for c in chunks) // max(len(chunks), 1)}"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

mdkeychunker/config.py

@@ -0,0 +1,40 @@
+"""Configuration for MDKeyChunker."""
+import os
+from dataclasses import dataclass
+from dotenv import load_dotenv
+
+
+@dataclass
+class Config:
+    # LLM
+    llm_provider: str = "openai"           # openai | anthropic | openai_compatible
+    llm_api_key: str = ""
+    llm_base_url: str = ""                 # For Ollama/vLLM: http://localhost:11434/v1
+    llm_model: str = "gpt-4o-mini"
+    # Chunking
+    min_chunk_size: int = 100              # Minimum chars per chunk
+    max_chunk_size: int = 1500             # Soft max chars per chunk
+    # Restructuring
+    merge_by_keys: bool = True             # Enable key-based chunk merging
+    max_merged_size: int = 3000            # Max size after merging
+    # Below this, keyless chunks get context enrichment
+    min_orphan_size: int = 200
+    # General
+    log_level: str = "INFO"
+
+    @classmethod
+    def from_env(cls, env_file: str | None = None) -> "Config":
+        load_dotenv(env_file or ".env")
+        return cls(
+            llm_provider=os.getenv("LLM_PROVIDER", "openai"),
+            llm_api_key=os.getenv("LLM_API_KEY", ""),
+            llm_base_url=os.getenv("LLM_BASE_URL", ""),
+            llm_model=os.getenv("LLM_MODEL", "gpt-4o-mini"),
+            min_chunk_size=int(os.getenv("MIN_CHUNK_SIZE", "100")),
+            max_chunk_size=int(os.getenv("MAX_CHUNK_SIZE", "1500")),
+            merge_by_keys=os.getenv(
+                "MERGE_BY_KEYS", "true").lower() in ("true", "1", "yes"),
+            max_merged_size=int(os.getenv("MAX_MERGED_SIZE", "3000")),
+            min_orphan_size=int(os.getenv("MIN_ORPHAN_SIZE", "200")),
+            log_level=os.getenv("LOG_LEVEL", "INFO"),
+        )