code-finder 0.1.0__py3-none-any.whl

claude_context/markdown_chunker.py

@@ -0,0 +1,421 @@
+"""
+Markdown Chunker for Claude Context
+
+Section-aware chunking of markdown files (README, docs) that:
+1. Preserves header hierarchy (H1 > H2 > H3) as scope chains
+2. Extracts code blocks with language tags
+3. Tracks links and references
+4. Creates semantically meaningful chunks for retrieval
+
+This complements the AST chunker by handling documentation content
+that explains "what/why/how-to-use" rather than implementation details.
+"""
+
+import re
+import logging
+from pathlib import Path
+from dataclasses import dataclass, field
+from typing import List, Dict, Any, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MarkdownChunk:
+    """Represents a chunk from a markdown file."""
+    text: str                          # Raw markdown content
+    contextualized_text: str           # Text with header context prepended
+    chunk_type: str                    # 'section', 'code_block', 'list', 'paragraph'
+    name: str                          # Section title or code language
+
+    # Location
+    line_range: Tuple[int, int]        # (start, end) 0-indexed
+    file_path: Optional[Path] = None
+
+    # Hierarchy
+    scope: List[str] = field(default_factory=list)  # Header chain: ["Getting Started", "Installation"]
+    header_level: int = 0              # 1 for H1, 2 for H2, etc.
+
+    # Code block metadata
+    code_language: Optional[str] = None
+    code_content: Optional[str] = None
+
+    # Links found in this section
+    links: List[Dict[str, str]] = field(default_factory=list)  # [{"text": "...", "url": "..."}]
+
+    @property
+    def size_chars(self) -> int:
+        return len(self.text)
+
+
+class MarkdownChunker:
+    """
+    Markdown-aware chunker that creates semantically meaningful chunks.
+
+    Unlike generic text splitting, this:
+    - Keeps sections together as logical units
+    - Tracks header hierarchy for scope chains
+    - Extracts code blocks as separate, searchable chunks
+    - Preserves context through header prepending
+    """
+
+    def __init__(
+        self,
+        max_chunk_size: int = 2000,
+        min_chunk_size: int = 100,
+        extract_code_blocks: bool = True,
+        include_header_context: bool = True
+    ):
+        """
+        Initialize the markdown chunker.
+
+        Args:
+            max_chunk_size: Maximum characters per chunk
+            min_chunk_size: Minimum characters to create a chunk
+            extract_code_blocks: Whether to extract code blocks as separate chunks
+            include_header_context: Whether to prepend header hierarchy to chunks
+        """
+        self.max_chunk_size = max_chunk_size
+        self.min_chunk_size = min_chunk_size
+        self.extract_code_blocks = extract_code_blocks
+        self.include_header_context = include_header_context
+
+        # Regex patterns
+        self.header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
+        self.code_block_pattern = re.compile(
+            r'^```(\w*)\n(.*?)^```',
+            re.MULTILINE | re.DOTALL
+        )
+        self.link_pattern = re.compile(r'\[([^\]]+)\]\(([^)]+)\)')
+        self.image_pattern = re.compile(r'!\[([^\]]*)\]\(([^)]+)\)')
+
+    def chunk_file(self, file_path: Path, content: Optional[str] = None) -> List[MarkdownChunk]:
+        """
+        Chunk a markdown file into semantic units.
+
+        Args:
+            file_path: Path to the markdown file
+            content: Optional content string (if not provided, reads from file)
+
+        Returns:
+            List of MarkdownChunk objects
+        """
+        if content is None:
+            try:
+                content = file_path.read_text(encoding='utf-8')
+            except Exception as e:
+                logger.error(f"Failed to read {file_path}: {e}")
+                return []
+
+        if not content.strip():
+            return []
+
+        chunks = []
+
+        # First, extract and replace code blocks with placeholders
+        code_blocks = []
+        if self.extract_code_blocks:
+            content, code_blocks = self._extract_code_blocks(content)
+
+        # Parse sections based on headers
+        sections = self._parse_sections(content, file_path)
+
+        # Create chunks from sections
+        for section in sections:
+            section_chunks = self._chunk_section(section, file_path)
+            chunks.extend(section_chunks)
+
+        # Add code blocks as separate chunks with their section context
+        for cb in code_blocks:
+            chunk = self._create_code_block_chunk(cb, file_path)
+            if chunk:
+                chunks.append(chunk)
+
+        logger.debug(f"Created {len(chunks)} chunks from {file_path.name}")
+        return chunks
+
+    def _extract_code_blocks(self, content: str) -> Tuple[str, List[Dict]]:
+        """Extract code blocks and replace with placeholders."""
+        code_blocks = []
+        placeholder_content = content
+
+        for i, match in enumerate(self.code_block_pattern.finditer(content)):
+            language = match.group(1) or 'text'
+            code = match.group(2).strip()
+
+            # Find line numbers
+            pre_content = content[:match.start()]
+            start_line = pre_content.count('\n')
+            end_line = start_line + code.count('\n') + 2  # +2 for ``` lines
+
+            # Find the section context (last header before this code block)
+            scope = self._find_scope_at_position(content, match.start())
+
+            code_blocks.append({
+                'language': language,
+                'code': code,
+                'start_line': start_line,
+                'end_line': end_line,
+                'scope': scope,
+                'full_match': match.group(0)
+            })
+
+            # Replace with placeholder to preserve line numbers
+            placeholder = f"\n[CODE_BLOCK_{i}]\n"
+            placeholder_content = placeholder_content.replace(match.group(0), placeholder, 1)
+
+        return placeholder_content, code_blocks
+
+    def _find_scope_at_position(self, content: str, position: int) -> List[str]:
+        """Find the header hierarchy at a given position in the content."""
+        pre_content = content[:position]
+
+        scope = []
+        current_level = 0
+
+        for match in self.header_pattern.finditer(pre_content):
+            level = len(match.group(1))
+            title = match.group(2).strip()
+
+            if level <= current_level:
+                # Pop scope back to parent level
+                while scope and len(scope) >= level:
+                    scope.pop()
+
+            scope.append(title)
+            current_level = level
+
+        return scope
+
+    def _parse_sections(self, content: str, file_path: Path) -> List[Dict]:
+        """Parse markdown into sections based on headers."""
+        lines = content.split('\n')
+        sections = []
+        current_section = {
+            'title': file_path.stem,  # Use filename as default title
+            'level': 0,
+            'content_lines': [],
+            'start_line': 0,
+            'scope': []
+        }
+
+        scope_stack = []  # [(level, title), ...]
+
+        for i, line in enumerate(lines):
+            header_match = self.header_pattern.match(line)
+
+            if header_match:
+                # Save current section if it has content
+                if current_section['content_lines']:
+                    current_section['end_line'] = i - 1
+                    sections.append(current_section)
+
+                level = len(header_match.group(1))
+                title = header_match.group(2).strip()
+
+                # Update scope stack
+                while scope_stack and scope_stack[-1][0] >= level:
+                    scope_stack.pop()
+                scope_stack.append((level, title))
+
+                # Create new section
+                current_section = {
+                    'title': title,
+                    'level': level,
+                    'content_lines': [],
+                    'start_line': i,
+                    'scope': [s[1] for s in scope_stack]
+                }
+            else:
+                current_section['content_lines'].append(line)
+
+        # Don't forget the last section
+        if current_section['content_lines']:
+            current_section['end_line'] = len(lines) - 1
+            sections.append(current_section)
+
+        return sections
+
+    def _chunk_section(self, section: Dict, file_path: Path) -> List[MarkdownChunk]:
+        """Create chunks from a section, splitting if too large."""
+        content = '\n'.join(section['content_lines']).strip()
+
+        if not content or len(content) < self.min_chunk_size:
+            # Skip very small sections or merge with title only
+            if section['title'] and section['level'] > 0:
+                content = f"# {section['title']}\n\n{content}" if content else f"# {section['title']}"
+                if len(content) < self.min_chunk_size:
+                    return []
+
+        chunks = []
+        scope = section['scope']
+
+        # Build contextualized text
+        if self.include_header_context and scope:
+            context_prefix = " > ".join(scope) + "\n\n"
+        else:
+            context_prefix = ""
+
+        # Extract links from content
+        links = self._extract_links(content)
+
+        if len(content) <= self.max_chunk_size:
+            # Single chunk for this section
+            chunk = MarkdownChunk(
+                text=content,
+                contextualized_text=context_prefix + content,
+                chunk_type='section',
+                name=section['title'],
+                line_range=(section['start_line'], section.get('end_line', section['start_line'])),
+                file_path=file_path,
+                scope=scope[:-1] if scope else [],  # Parent scope (excluding self)
+                header_level=section['level'],
+                links=links
+            )
+            chunks.append(chunk)
+        else:
+            # Split large sections by paragraphs
+            paragraphs = self._split_by_paragraphs(content)
+            current_chunk_lines = []
+            current_size = 0
+            chunk_start = section['start_line']
+
+            for para in paragraphs:
+                para_size = len(para)
+
+                if current_size + para_size > self.max_chunk_size and current_chunk_lines:
+                    # Emit current chunk
+                    chunk_content = '\n\n'.join(current_chunk_lines)
+                    chunk = MarkdownChunk(
+                        text=chunk_content,
+                        contextualized_text=context_prefix + chunk_content,
+                        chunk_type='section',
+                        name=section['title'],
+                        line_range=(chunk_start, chunk_start + chunk_content.count('\n')),
+                        file_path=file_path,
+                        scope=scope[:-1] if scope else [],
+                        header_level=section['level'],
+                        links=self._extract_links(chunk_content)
+                    )
+                    chunks.append(chunk)
+
+                    current_chunk_lines = []
+                    current_size = 0
+                    chunk_start = chunk_start + chunk_content.count('\n') + 1
+
+                current_chunk_lines.append(para)
+                current_size += para_size
+
+            # Don't forget the last chunk
+            if current_chunk_lines:
+                chunk_content = '\n\n'.join(current_chunk_lines)
+                chunk = MarkdownChunk(
+                    text=chunk_content,
+                    contextualized_text=context_prefix + chunk_content,
+                    chunk_type='section',
+                    name=section['title'],
+                    line_range=(chunk_start, section.get('end_line', chunk_start)),
+                    file_path=file_path,
+                    scope=scope[:-1] if scope else [],
+                    header_level=section['level'],
+                    links=self._extract_links(chunk_content)
+                )
+                chunks.append(chunk)
+
+        return chunks
+
+    def _create_code_block_chunk(self, cb: Dict, file_path: Path) -> Optional[MarkdownChunk]:
+        """Create a chunk from an extracted code block."""
+        code = cb['code']
+        if len(code) < self.min_chunk_size // 2:  # Allow smaller code blocks
+            return None
+
+        language = cb['language']
+        scope = cb['scope']
+
+        # Build contextualized text
+        if self.include_header_context and scope:
+            context_prefix = " > ".join(scope) + f"\n\nCode example ({language}):\n\n"
+        else:
+            context_prefix = f"Code example ({language}):\n\n"
+
+        return MarkdownChunk(
+            text=f"```{language}\n{code}\n```",
+            contextualized_text=context_prefix + code,
+            chunk_type='code_block',
+            name=f"{language} example",
+            line_range=(cb['start_line'], cb['end_line']),
+            file_path=file_path,
+            scope=scope,
+            header_level=0,
+            code_language=language,
+            code_content=code
+        )
+
+    def _split_by_paragraphs(self, content: str) -> List[str]:
+        """Split content into paragraphs."""
+        # Split on double newlines
+        paragraphs = re.split(r'\n\s*\n', content)
+        return [p.strip() for p in paragraphs if p.strip()]
+
+    def _extract_links(self, content: str) -> List[Dict[str, str]]:
+        """Extract markdown links from content."""
+        links = []
+
+        for match in self.link_pattern.finditer(content):
+            # Skip images
+            if content[match.start()-1:match.start()] == '!':
+                continue
+            links.append({
+                'text': match.group(1),
+                'url': match.group(2)
+            })
+
+        return links
+
+
+# Convenience function
+def chunk_markdown_file(file_path: Path, **kwargs) -> List[MarkdownChunk]:
+    """
+    Chunk a markdown file using default settings.
+
+    Args:
+        file_path: Path to the markdown file
+        **kwargs: Additional arguments for MarkdownChunker
+
+    Returns:
+        List of MarkdownChunk objects
+    """
+    chunker = MarkdownChunker(**kwargs)
+    return chunker.chunk_file(file_path)
+
+
+if __name__ == "__main__":
+    # Test with a sample README
+    import sys
+
+    if len(sys.argv) > 1:
+        test_file = Path(sys.argv[1])
+    else:
+        test_file = Path("README.md")
+
+    if test_file.exists():
+        chunker = MarkdownChunker()
+        chunks = chunker.chunk_file(test_file)
+
+        print(f"\n📄 Chunked {test_file.name}: {len(chunks)} chunks\n")
+
+        for i, chunk in enumerate(chunks, 1):
+            print(f"{'='*60}")
+            print(f"Chunk {i}: {chunk.name} ({chunk.chunk_type})")
+            print(f"  Scope: {' > '.join(chunk.scope) if chunk.scope else '(root)'}")
+            print(f"  Lines: {chunk.line_range[0]}-{chunk.line_range[1]}")
+            print(f"  Size: {chunk.size_chars} chars")
+            if chunk.code_language:
+                print(f"  Language: {chunk.code_language}")
+            if chunk.links:
+                print(f"  Links: {len(chunk.links)}")
+            print(f"\n  Preview: {chunk.text[:150]}...")
+            print()
+    else:
+        print(f"File not found: {test_file}")