contextual-engine 0.1.0__py3-none-any.whl

contextual/indexing/chunker.py

@@ -0,0 +1,691 @@
+"""Tree-sitter AST-based code chunking engine.
+
+Implements the Sweep split-merge algorithm (validated as cAST in arXiv:2506.15655)
+with deterministic structural headers for embedding. Every chunk gets context
+about its file location, parent scope, imports, and decorators - reducing retrieval
+errors by ~35% per Anthropic's contextual retrieval research.
+
+Key invariants:
+- Bytes not chars (tree-sitter indexes in bytes)
+- 1,500 byte target / 2,000 max / 50 min
+- Content-hash keyed (SHA-256) for incremental re-embedding
+- Thread-safe parser pool (Parser objects are NOT thread-safe)
+- Never split mid-function (preserve AST structural boundaries)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import threading
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+from tree_sitter_language_pack import SupportedLanguage, get_parser
+
+from contextual.core.errors import ErrorCode, IndexingError
+from contextual.core.models import Chunk, ChunkType, Language, ModelType
+from contextual.observability import get_logger
+
+
+if TYPE_CHECKING:
+    from tree_sitter import Node, Parser, Tree
+
+logger = get_logger(__name__)
+
+# Chunking parameters from TECHNICAL_SPEC §6
+TARGET_CHUNK_BYTES = 1500
+MAX_CHUNK_BYTES = 2000
+MIN_CHUNK_BYTES = 50
+MAX_FILE_BYTES = 2 * 1024 * 1024  # 2 MB
+FALLBACK_THRESHOLD_BYTES = 500 * 1024  # 500 KB
+
+# Language mappings (extension → tree-sitter language name)
+LANGUAGE_MAP = {
+    # Code languages
+    ".py": "python",
+    ".ts": "typescript",
+    ".tsx": "tsx",
+    ".js": "javascript",
+    ".jsx": "javascript",  # JSX uses javascript parser
+    ".go": "go",
+    ".java": "java",
+    ".rs": "rust",
+    ".cs": "c_sharp",
+    # Config formats
+    ".json": "json",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".toml": "toml",
+    "Dockerfile": "dockerfile",
+    ".md": "markdown",
+}
+
+# Top-level node types per language (from TECHNICAL_SPEC §6.1)
+TOP_LEVEL_NODES = {
+    "python": {"function_definition", "class_definition", "decorated_definition"},
+    "typescript": {
+        "function_declaration",
+        "class_declaration",
+        "interface_declaration",
+        "export_statement",
+        "lexical_declaration",
+    },
+    "tsx": {
+        "function_declaration",
+        "class_declaration",
+        "interface_declaration",
+        "export_statement",
+        "lexical_declaration",
+    },
+    "javascript": {
+        "function_declaration",
+        "class_declaration",
+        "export_statement",
+        "lexical_declaration",
+    },
+    "go": {"function_declaration", "type_declaration", "method_declaration"},
+    "java": {"class_declaration", "method_declaration", "interface_declaration"},
+    "rust": {
+        "function_item",
+        "impl_item",
+        "struct_item",
+        "enum_item",
+        "trait_item",
+    },
+    "c_sharp": {"class_declaration", "method_declaration", "interface_declaration"},
+}
+
+
+@dataclass
+class ParsedNode:
+    """Wrapper for tree-sitter nodes with byte ranges."""
+
+    node: Node
+    start_byte: int
+    end_byte: int
+    node_type: str
+    text: bytes
+
+    @property
+    def size_bytes(self) -> int:
+        """Size in bytes."""
+        return self.end_byte - self.start_byte
+
+
+class ParserPool:
+    """Thread-safe parser pool.
+    
+    tree_sitter.Parser is NOT thread-safe. We maintain a pool of parsers
+    per language and hand them out with thread-local storage for safety.
+    """
+
+    def __init__(self) -> None:
+        self._parsers: dict[str, list[Parser]] = defaultdict(list)
+        self._lock = threading.Lock()
+        self._local = threading.local()
+
+    def get_parser(self, language: str) -> Parser:
+        """Get or create a parser for this language in this thread."""
+        # Check thread-local cache first
+        cache_key = f"parser_{language}"
+        if hasattr(self._local, cache_key):
+            return getattr(self._local, cache_key)
+
+        # Try to get from pool
+        with self._lock:
+            if self._parsers[language]:
+                parser = self._parsers[language].pop()
+            else:
+                # Create new parser
+                parser = get_parser(cast("SupportedLanguage", language))
+
+        # Cache in thread-local
+        setattr(self._local, cache_key, parser)
+        return parser
+
+    def return_parser(self, language: str, parser: Parser) -> None:
+        """Return parser to pool (for cross-thread scenarios)."""
+        with self._lock:
+            self._parsers[language].append(parser)
+
+
+# Global parser pool
+_parser_pool = ParserPool()
+
+
+class ASTChunker:
+    """AST-based code chunker using tree-sitter.
+    
+    Implements split-merge algorithm:
+    1. Parse AST and extract top-level nodes
+    2. Merge small adjacent nodes until hitting target size
+    3. Split large nodes recursively at structural boundaries
+    4. Attach structural headers (file, parent, imports, decorators)
+    """
+
+    def __init__(self, file_path: Path, content: bytes, language: str) -> None:
+        """Initialize chunker.
+        
+        Args:
+            file_path: Relative path from project root
+            content: File content as bytes
+            language: Language identifier (python, typescript, etc.)
+            
+        Raises:
+            IndexingError: If file too large or parsing fails
+        """
+        self.file_path = file_path
+        self.content = content
+        self.language = language
+        self.relative_path = str(file_path)
+
+        # Size check
+        if len(content) > MAX_FILE_BYTES:
+            raise IndexingError(
+                ErrorCode.INDEXING_CHUNKER_SPLIT_FAILED,
+                f"File too large: {len(content)} bytes (max {MAX_FILE_BYTES})",
+            )
+
+        # Parse AST
+        self.parser = _parser_pool.get_parser(language)
+        self.tree: Tree = self.parser.parse(content)
+        self.root = self.tree.root_node
+
+        # Extract structural context
+        self.imports = self._extract_imports()
+        self.decorators: dict[int, list[str]] = {}  # node_id → decorator names
+
+    def chunk(self) -> list[Chunk]:
+        """Chunk the file into AST-aware segments.
+        
+        Returns:
+            List of Chunk objects with headers, bodies, and metadata
+        """
+        # Extract top-level nodes
+        top_level = self._extract_top_level_nodes()
+
+        if not top_level:
+            # Fallback: treat entire file as one chunk if no structure
+            logger.warning(
+                "no_top_level_nodes",
+                file=self.relative_path,
+                language=self.language,
+            )
+            return self._create_fallback_chunk()
+
+        # Split-merge algorithm
+        chunks = self._split_merge(top_level)
+
+        # Build Chunk objects with headers
+        return [self._build_chunk(parsed_node) for parsed_node in chunks]
+
+    def _extract_top_level_nodes(self) -> list[ParsedNode]:
+        """Extract top-level structural nodes (functions, classes, etc.)."""
+        if self.language not in TOP_LEVEL_NODES:
+            # Config files: split on reasonable boundaries or whole file
+            return self._extract_config_chunks()
+
+        target_types = TOP_LEVEL_NODES[self.language]
+        nodes = []
+
+        def visit(node: Node) -> None:
+            """Recursively visit nodes, collecting top-level targets."""
+            if node.type in target_types:
+                nodes.append(
+                    ParsedNode(
+                        node=node,
+                        start_byte=node.start_byte,
+                        end_byte=node.end_byte,
+                        node_type=node.type,
+                        text=self.content[node.start_byte : node.end_byte],
+                    )
+                )
+                # Don't recurse into matched nodes
+                return
+
+            # Recurse into children
+            for child in node.children:
+                visit(child)
+
+        visit(self.root)
+        return nodes
+
+    def _extract_config_chunks(self) -> list[ParsedNode]:
+        """Extract chunks from config files (JSON, YAML, TOML, etc.).
+        
+        For config files, we split on reasonable boundaries like top-level
+        keys or sections, with sensible fallback to whole-file.
+        """
+        # Simple strategy: whole file if small, line-based split if large
+        if len(self.content) <= MAX_CHUNK_BYTES:
+            return [
+                ParsedNode(
+                    node=self.root,
+                    start_byte=0,
+                    end_byte=len(self.content),
+                    node_type="document",
+                    text=self.content,
+                )
+            ]
+
+        # Line-based fallback
+        return self._line_based_chunks()
+
+    def _line_based_chunks(self) -> list[ParsedNode]:
+        """Fallback: split on line boundaries when AST strategy fails."""
+        lines = self.content.split(b"\n")
+        chunks = []
+        current_chunk = []
+        current_bytes = 0
+
+        for line in lines:
+            line_bytes = len(line) + 1  # +1 for newline
+            if current_bytes + line_bytes > TARGET_CHUNK_BYTES and current_chunk:
+                # Flush current chunk
+                chunk_text = b"\n".join(current_chunk) + b"\n"
+                chunks.append(
+                    ParsedNode(
+                        node=self.root,  # Placeholder
+                        start_byte=0,  # Will fix in post-processing
+                        end_byte=len(chunk_text),
+                        node_type="line_chunk",
+                        text=chunk_text,
+                    )
+                )
+                current_chunk = [line]
+                current_bytes = line_bytes
+            else:
+                current_chunk.append(line)
+                current_bytes += line_bytes
+
+        # Final chunk
+        if current_chunk:
+            chunk_text = b"\n".join(current_chunk) + b"\n"
+            chunks.append(
+                ParsedNode(
+                    node=self.root,
+                    start_byte=0,
+                    end_byte=len(chunk_text),
+                    node_type="line_chunk",
+                    text=chunk_text,
+                )
+            )
+
+        return chunks
+
+    def _split_merge(self, nodes: list[ParsedNode]) -> list[ParsedNode]:
+        """Split-merge algorithm.
+        
+        1. Merge adjacent small nodes until target size
+        2. Split large nodes at structural boundaries
+        3. Keep complete functions together (never split mid-function)
+        """
+        result = []
+        i = 0
+
+        while i < len(nodes):
+            current = nodes[i]
+
+            # If node is under max, try to merge with next nodes
+            if current.size_bytes < MAX_CHUNK_BYTES:
+                merged = self._merge_forward(nodes, i)
+                result.append(merged)
+                i += merged.end_byte  # Skip merged nodes (placeholder logic)
+            else:
+                # Node too large, try to split
+                split_nodes = self._split_node(current)
+                result.extend(split_nodes)
+
+            i += 1
+
+        return result
+
+    def _merge_forward(self, nodes: list[ParsedNode], start_idx: int) -> ParsedNode:
+        """Merge nodes forward from start_idx until hitting target or max size.
+        
+        Returns:
+            Merged ParsedNode
+        """
+        merged_start = nodes[start_idx].start_byte
+        merged_end = nodes[start_idx].end_byte
+        merged_type = nodes[start_idx].node_type
+        last_idx = start_idx
+
+        for i in range(start_idx + 1, len(nodes)):
+            candidate_end = nodes[i].end_byte
+            candidate_size = candidate_end - merged_start
+
+            if candidate_size > MAX_CHUNK_BYTES:
+                break
+
+            merged_end = candidate_end
+            last_idx = i
+
+            if candidate_size >= TARGET_CHUNK_BYTES:
+                break
+
+        merged_text = self.content[merged_start:merged_end]
+        result = ParsedNode(
+            node=nodes[start_idx].node,  # Use first node as representative
+            start_byte=merged_start,
+            end_byte=merged_end,
+            node_type=merged_type,
+            text=merged_text,
+        )
+        result.end_byte = last_idx  # Hack: store how many nodes we merged
+        return result
+
+    def _split_node(self, node: ParsedNode) -> list[ParsedNode]:
+        """Split a large node at structural boundaries.
+        
+        Strategy: Find child nodes and split recursively, or fall back to
+        byte-based splitting at safe boundaries (newlines).
+        """
+        # Try to split at child boundaries
+        if node.node.child_count > 1:
+            children = []
+            for child in node.node.children:
+                if child.end_byte - child.start_byte >= MIN_CHUNK_BYTES:
+                    children.append(
+                        ParsedNode(
+                            node=child,
+                            start_byte=child.start_byte,
+                            end_byte=child.end_byte,
+                            node_type=child.type,
+                            text=self.content[child.start_byte : child.end_byte],
+                        )
+                    )
+
+            if children:
+                # Recursively split-merge children
+                return self._split_merge(children)
+
+        # Fallback: split at newline boundaries
+        return self._split_at_newlines(node)
+
+    def _split_at_newlines(self, node: ParsedNode) -> list[ParsedNode]:
+        """Split large node at newline boundaries as last resort."""
+        text = node.text
+        chunks = []
+        current_start = 0
+
+        lines = text.split(b"\n")
+        current_chunk = []
+        current_bytes = 0
+
+        for line in lines:
+            line_bytes = len(line) + 1
+            if current_bytes + line_bytes > TARGET_CHUNK_BYTES and current_chunk:
+                chunk_text = b"\n".join(current_chunk) + b"\n"
+                chunks.append(chunk_text)
+                current_chunk = [line]
+                current_bytes = line_bytes
+            else:
+                current_chunk.append(line)
+                current_bytes += line_bytes
+
+        if current_chunk:
+            chunk_text = b"\n".join(current_chunk)
+            chunks.append(chunk_text)
+
+        # Convert to ParsedNode objects
+        return [
+            ParsedNode(
+                node=node.node,
+                start_byte=0,
+                end_byte=len(chunk),
+                node_type=node.node_type,
+                text=chunk,
+            )
+            for chunk in chunks
+            if len(chunk) >= MIN_CHUNK_BYTES
+        ]
+
+    def _extract_imports(self) -> list[str]:
+        """Extract import statements from the file."""
+        imports = []
+
+        # Language-specific import extraction
+        if self.language == "python":
+            imports = self._extract_python_imports()
+        elif self.language in {"typescript", "tsx", "javascript"}:
+            imports = self._extract_js_imports()
+        # Add more languages as needed
+
+        return imports
+
+    def _extract_python_imports(self) -> list[str]:
+        """Extract Python import statements."""
+        imports = []
+
+        def visit(node: Node) -> None:
+            if node.type in {"import_statement", "import_from_statement"}:
+                import_text = self.content[node.start_byte : node.end_byte].decode(
+                    "utf-8", errors="replace"
+                )
+                imports.append(import_text.strip())
+            for child in node.children:
+                visit(child)
+
+        visit(self.root)
+        return imports
+
+    def _extract_js_imports(self) -> list[str]:
+        """Extract JavaScript/TypeScript import statements."""
+        imports = []
+
+        def visit(node: Node) -> None:
+            if node.type in {"import_statement", "import_declaration"}:
+                import_text = self.content[node.start_byte : node.end_byte].decode(
+                    "utf-8", errors="replace"
+                )
+                imports.append(import_text.strip())
+            for child in node.children:
+                visit(child)
+
+        visit(self.root)
+        return imports
+
+    def _get_parent_scope(self, node: Node) -> str | None:
+        """Get parent class or module for a node."""
+        current = node.parent
+        while current:
+            if current.type in {"class_definition", "class_declaration"}:
+                # Find class name
+                for child in current.children:
+                    if child.type in {"identifier", "type_identifier"}:
+                        return self.content[child.start_byte : child.end_byte].decode(
+                            "utf-8", errors="replace"
+                        )
+            current = current.parent
+        return None
+
+    def _extract_symbol_name(self, node: Node) -> str | None:
+        """Extract function/class/method name from node."""
+        # Look for identifier child nodes
+        for child in node.children:
+            if child.type in {"identifier", "type_identifier", "property_identifier"}:
+                return self.content[child.start_byte : child.end_byte].decode(
+                    "utf-8", errors="replace"
+                )
+        return None
+
+    def _build_chunk(self, parsed_node: ParsedNode) -> Chunk:
+        """Build a Chunk object with structural header.
+        
+        Header format (TECHNICAL_SPEC §6.2):
+        FILE: <relative_path>
+        LANG: <language>
+        PARENT: <parent_class_or_module>
+        IMPORTS:
+        <imports_or_none>
+        DECORATORS:
+        <decorators_or_none>
+        -----
+        <chunk_body>
+        """
+        # Extract parent scope
+        parent_scope = self._get_parent_scope(parsed_node.node)
+
+        # Build header
+        header_lines = [
+            f"FILE: {self.relative_path}",
+            f"LANG: {self.language}",
+            f"PARENT: {parent_scope or 'none'}",
+            "IMPORTS:",
+        ]
+
+        if self.imports:
+            header_lines.extend(self.imports)
+        else:
+            header_lines.append("none")
+
+        header_lines.append("DECORATORS:")
+        # TODO: Extract decorators per node
+        header_lines.append("none")
+        header_lines.append("-----")
+
+        header = "\n".join(header_lines)
+
+        # Body is the actual code
+        body = parsed_node.text.decode("utf-8", errors="replace")
+
+        # Full text for hashing
+        full_text = f"{header}\n{body}"
+
+        # Content hash (SHA-256)
+        content_hash = hashlib.sha256(full_text.encode("utf-8")).hexdigest()
+
+        # Extract symbol name from node
+        symbol_name = self._extract_symbol_name(parsed_node.node)
+
+        # Map language to enum
+        try:
+            language_enum = Language(self.language)
+        except ValueError:
+            language_enum = Language.PYTHON  # Fallback
+
+        # Map chunk type to enum
+        chunk_type_map = {
+            "function_definition": ChunkType.FUNCTION,
+            "function_declaration": ChunkType.FUNCTION,
+            "method_declaration": ChunkType.METHOD,
+            "class_definition": ChunkType.CLASS,
+            "class_declaration": ChunkType.CLASS,
+            "interface_declaration": ChunkType.INTERFACE,
+            "module": ChunkType.MODULE,
+            "file": ChunkType.MODULE,
+            "line_chunk": ChunkType.MODULE,
+        }
+        chunk_type_enum = chunk_type_map.get(
+            parsed_node.node_type, ChunkType.MODULE
+        )
+
+        # Extract start and end lines
+        start_line = parsed_node.node.start_point[0]
+        end_line = parsed_node.node.end_point[0]
+
+        return Chunk(
+            vector=[0.0] * 768,  # Empty, filled during embedding phase
+            path=self.relative_path,
+            language=language_enum,
+            symbol_name=symbol_name,
+            chunk_type=chunk_type_enum,
+            content_hash=content_hash,
+            header=header,
+            body=body,
+            start_line=start_line,
+            end_line=end_line,
+            model_type=ModelType.JINA_CODE_V2, # Default model
+        )
+
+    def _create_fallback_chunk(self) -> list[Chunk]:
+        """Create single chunk when no structure found."""
+        header_lines = [
+            f"FILE: {self.relative_path}",
+            f"LANG: {self.language}",
+            "PARENT: none",
+            "IMPORTS:",
+            "none",
+            "DECORATORS:",
+            "none",
+            "-----",
+        ]
+        header = "\n".join(header_lines)
+        body = self.content.decode("utf-8", errors="replace")
+        full_text = f"{header}\n{body}"
+        content_hash = hashlib.sha256(full_text.encode("utf-8")).hexdigest()
+
+        # Map language to enum
+        try:
+            language_enum = Language(self.language)
+        except ValueError:
+            language_enum = Language.PYTHON  # Fallback
+
+        return [
+            Chunk(
+                vector=[0.0] * 768,
+                path=self.relative_path,
+                language=language_enum,
+                symbol_name=None,
+                chunk_type=ChunkType.MODULE,
+                content_hash=content_hash,
+                header=header,
+                body=body,
+                start_line=0,
+                end_line=len(body.splitlines()),
+                model_type=ModelType.JINA_CODE_V2,
+            )
+        ]
+
+
+def detect_language(file_path: Path) -> str | None:
+    """Detect language from file extension.
+    
+    Args:
+        file_path: Path to file
+        
+    Returns:
+        Language identifier or None if unsupported
+    """
+    suffix = file_path.suffix.lower()
+    if suffix in LANGUAGE_MAP:
+        return LANGUAGE_MAP[suffix]
+
+    # Check for Dockerfile (no extension)
+    if file_path.name in {"Dockerfile", "Dockerfile.dev", "Dockerfile.prod"}:
+        return "dockerfile"
+
+    return None
+
+
+def chunk_file(file_path: Path, content: bytes, language: str) -> list[Chunk]:
+    """Chunk a file using tree-sitter AST parsing.
+    
+    Args:
+        file_path: Relative path from project root
+        content: File content as bytes
+        language: Language identifier
+        
+    Returns:
+        List of Chunk objects
+        
+    Raises:
+        IndexingError: If chunking fails
+    """
+    try:
+        chunker = ASTChunker(file_path, content, language)
+        return chunker.chunk()
+    except Exception as e:
+        logger.error(
+            "chunking_failed",
+            file=str(file_path),
+            language=language,
+            error=str(e),
+        )
+        raise IndexingError(
+            ErrorCode.INDEXING_CHUNKER_SPLIT_FAILED,
+            f"Failed to chunk {file_path}: {e}",
+        ) from e