autochunks 0.0.8__py3-none-any.whl

autochunk/chunkers/parent_child.py

@@ -0,0 +1,172 @@
+
+from __future__ import annotations
+from typing import List, Dict, Any, Optional
+from .base import BaseChunker, Chunk
+from .recursive_character import RecursiveCharacterChunker
+from ..utils.text import count_tokens
+
+class ParentChildChunker(BaseChunker):
+    """
+     Hierarchical (Small-to-Big) Parent-Child Chunker.
+    
+    BEST-OF-BREED FEATURES:
+    1. N-Level Hierarchy: Configurable depth (Document > Section > Paragraph > Sentence).
+    2. Sibling References: Tracks prev_node_id and next_node_id for traversal.
+    3. Parent Context: Stores parent text in metadata for rich LLM context.
+    4. Child Overlap: Optional overlap between children for context continuity.
+    """
+    name = "parent_child"
+
+    def __init__(self, 
+                 chunk_sizes: List[int] = None,
+                 overlap: int = 0,
+                 track_siblings: bool = True):
+        """
+        Initialize the chunker.
+        
+        Args:
+            chunk_sizes: List of sizes for each hierarchy level.
+                         Default [2048, 512, 128] creates 3 levels: large -> medium -> small
+            overlap: Token overlap between sibling chunks at each level
+            track_siblings: If True, add prev/next references to metadata
+        """
+        self.chunk_sizes = chunk_sizes or [2048, 512, 128]
+        self.overlap = overlap
+        self.track_siblings = track_siblings
+
+    def chunk(self, 
+              doc_id: str, 
+              text: str,
+              parent_size: int = None,
+              child_size: int = None,
+              overlap: int = None,
+              return_all_levels: bool = False,
+              **params) -> List[Chunk]:
+        """
+        Create hierarchical chunks with parent-child relationships.
+        
+        Args:
+            doc_id: Document identifier
+            text: Input text
+            parent_size: Override first level size (for backward compatibility)
+            child_size: Override last level size (for backward compatibility)
+            overlap: Override overlap setting
+            return_all_levels: If True, return chunks from all levels, not just leaves
+        
+        Returns:
+            List of Chunk objects (leaf nodes by default, or all nodes if return_all_levels=True)
+        """
+        # Handle legacy 2-level params
+        if parent_size and child_size:
+            chunk_sizes = [parent_size, child_size]
+        else:
+            chunk_sizes = self.chunk_sizes
+        
+        if overlap is None:
+            overlap = self.overlap
+
+        base_chunker = RecursiveCharacterChunker()
+        
+        all_chunks = []
+        
+        def _build_hierarchy(input_text: str, 
+                            level: int, 
+                            parent_info: Dict[str, Any],
+                            node_path: str) -> List[Chunk]:
+            """
+            Recursively build chunk hierarchy.
+            
+            Args:
+                input_text: Text to chunk
+                level: Current hierarchy level (0 = root)
+                parent_info: Info about parent chunk
+                node_path: Path identifier for this node
+            
+            Returns:
+                List of chunks at this level (and below if return_all_levels)
+            """
+            if level >= len(chunk_sizes):
+                return []
+            
+            current_size = chunk_sizes[level]
+            is_leaf = (level == len(chunk_sizes) - 1)
+            
+            # Create chunks at this level
+            level_chunks = base_chunker.chunk(
+                doc_id=f"{doc_id}_L{level}",
+                text=input_text,
+                base_token_size=current_size,
+                overlap=overlap
+            )
+            
+            result_chunks = []
+            
+            for idx, chunk in enumerate(level_chunks):
+                chunk_id = f"{node_path}#L{level}#{idx}"
+                
+                # Build metadata with parent info
+                meta = {
+                    "chunk_index": idx,
+                    "level": level,
+                    "is_leaf": is_leaf,
+                    "strategy": "parent_child",
+                    "token_count": count_tokens(chunk.text)
+                }
+                
+                # Add parent references
+                if parent_info:
+                    meta["parent_id"] = parent_info.get("id")
+                    meta["parent_text"] = parent_info.get("text", "")[:500]  # Truncate for efficiency
+                
+                # Add sibling references
+                if self.track_siblings:
+                    if idx > 0:
+                        meta["prev_sibling_id"] = f"{node_path}#L{level}#{idx - 1}"
+                    if idx < len(level_chunks) - 1:
+                        meta["next_sibling_id"] = f"{node_path}#L{level}#{idx + 1}"
+                
+                node = Chunk(
+                    id=chunk_id,
+                    doc_id=doc_id,
+                    text=chunk.text,
+                    meta=meta
+                )
+                
+                # Add to results based on return_all_levels setting
+                if return_all_levels or is_leaf:
+                    result_chunks.append(node)
+                
+                # Recurse to children if not at leaf level
+                if not is_leaf:
+                    child_parent_info = {
+                        "id": chunk_id,
+                        "text": chunk.text
+                    }
+                    children = _build_hierarchy(
+                        chunk.text,
+                        level + 1,
+                        child_parent_info,
+                        chunk_id
+                    )
+                    
+                    # Update parent with child references
+                    if children and return_all_levels:
+                        node.meta["child_ids"] = [c.id for c in children]
+                    
+                    result_chunks.extend(children)
+            
+            return result_chunks
+        
+        # Build from root
+        root_parent_info = {
+            "id": doc_id,
+            "text": text[:500]  # Document context
+        }
+        
+        all_chunks = _build_hierarchy(text, 0, root_parent_info, doc_id)
+        
+        # Re-index final chunks sequentially
+        for i, chunk in enumerate(all_chunks):
+            chunk.meta["global_index"] = i
+        
+        return all_chunks

autochunk/chunkers/proposition.py

@@ -0,0 +1,175 @@
+
+from __future__ import annotations
+from typing import List, Callable, Optional
+from .base import BaseChunker, Chunk
+from ..utils.text import count_tokens, split_sentences
+
+class PropositionChunker(BaseChunker):
+    """
+    Proposition-Based Chunker for Atomic Fact Extraction.
+    
+    Instead of arbitrary text splits, this chunker extracts atomic propositions
+    (self-contained facts) from the text. Each chunk is a single, verifiable statement.
+    
+    BEST-OF-BREED FEATURES:
+    1. Fact-Level Granularity: Each chunk is one atomic fact.
+    2. Self-Contained: Every proposition is understandable without context.
+    3. Decontextualized: Pronouns and references are resolved.
+    4. LLM-Powered: Uses language model for accurate extraction.
+    
+    Reference: "Dense X Retrieval" paper, Greg Kamradt's proposition chunker.
+    """
+    name = "proposition"
+
+    DEFAULT_SYSTEM_PROMPT = """You are an expert at extracting atomic propositions from text.
+
+An atomic proposition is:
+- A single, self-contained fact
+- Expressed in a complete sentence
+- Understandable WITHOUT any additional context
+- Has all pronouns replaced with their referents
+- Contains no dependent references (like "this", "that", "the above")
+
+For example:
+Original: "John went to the store. He bought milk there."
+Propositions:
+1. John went to the store.
+2. John bought milk at the store.
+
+Note how "He" became "John" and "there" became "at the store"."""
+
+    DEFAULT_USER_TEMPLATE = """Extract all atomic propositions from the following text. Each proposition should be:
+1. A complete, self-contained sentence
+2. Understandable without additional context
+3. Have all pronouns resolved to their referents
+
+TEXT:
+{text}
+
+Output each proposition on a new line, numbered. Only output the propositions, no other text.
+
+1."""
+
+    def __init__(self,
+                 llm_fn: Callable[[str, str], str] = None,
+                 system_prompt: str = None,
+                 user_template: str = None,
+                 max_tokens_per_call: int = 2000):
+        """
+        Initialize the proposition chunker.
+        
+        Args:
+            llm_fn: Function that takes (system_prompt, user_message) and returns LLM response.
+            system_prompt: Custom system prompt
+            user_template: Custom user message template (must include {text})
+            max_tokens_per_call: Max tokens to process in one LLM call
+        """
+        self.llm_fn = llm_fn
+        self.system_prompt = system_prompt or self.DEFAULT_SYSTEM_PROMPT
+        self.user_template = user_template or self.DEFAULT_USER_TEMPLATE
+        self.max_tokens_per_call = max_tokens_per_call
+
+    def chunk(self,
+              doc_id: str,
+              text: str,
+              **params) -> List[Chunk]:
+        """
+        Extract atomic propositions from text.
+        
+        Args:
+            doc_id: Document identifier
+            text: Input text
+        
+        Returns:
+            List of Chunk objects, each containing one proposition
+        """
+        if not text.strip():
+            return []
+        
+        # Fallback if no LLM function
+        if self.llm_fn is None:
+            # Use sentence splitting as basic fallback
+            sentences = split_sentences(text)
+            return [
+                Chunk(
+                    id=f"{doc_id}#prop#{i}",
+                    doc_id=doc_id,
+                    text=s.strip(),
+                    meta={
+                        "chunk_index": i,
+                        "strategy": "proposition_fallback",
+                        "is_atomic": False
+                    }
+                ) for i, s in enumerate(sentences) if s.strip()
+            ]
+        
+        # Process text in chunks if too long
+        propositions = []
+        sentences = split_sentences(text)
+        
+        current_batch = []
+        current_tokens = 0
+        
+        for sentence in sentences:
+            sent_tokens = count_tokens(sentence)
+            
+            if current_tokens + sent_tokens > self.max_tokens_per_call and current_batch:
+                # Process current batch
+                batch_text = " ".join(current_batch)
+                batch_props = self._extract_propositions(batch_text)
+                propositions.extend(batch_props)
+                current_batch = []
+                current_tokens = 0
+            
+            current_batch.append(sentence)
+            current_tokens += sent_tokens
+        
+        # Process final batch
+        if current_batch:
+            batch_text = " ".join(current_batch)
+            batch_props = self._extract_propositions(batch_text)
+            propositions.extend(batch_props)
+        
+        # Create chunk objects
+        chunks = []
+        for i, prop in enumerate(propositions):
+            if prop.strip():
+                chunks.append(Chunk(
+                    id=f"{doc_id}#prop#{i}",
+                    doc_id=doc_id,
+                    text=prop.strip(),
+                    meta={
+                        "chunk_index": i,
+                        "strategy": "proposition",
+                        "is_atomic": True,
+                        "token_count": count_tokens(prop)
+                    }
+                ))
+        
+        return chunks
+
+    def _extract_propositions(self, text: str) -> List[str]:
+        """Extract propositions using LLM."""
+        user_message = self.user_template.format(text=text)
+        
+        try:
+            response = self.llm_fn(self.system_prompt, user_message)
+            
+            # Parse numbered list
+            propositions = []
+            for line in response.strip().split("\n"):
+                line = line.strip()
+                if not line:
+                    continue
+                
+                # Remove numbering (1. 2. etc or 1) 2) etc)
+                import re
+                cleaned = re.sub(r'^[\d]+[\.\)]\s*', '', line)
+                if cleaned:
+                    propositions.append(cleaned)
+            
+            return propositions
+            
+        except Exception as e:
+            # Fallback to sentence splitting
+            return split_sentences(text)

autochunk/chunkers/python_ast.py

@@ -0,0 +1,248 @@
+
+from __future__ import annotations
+from typing import List, Optional
+import ast
+from .base import BaseChunker, Chunk
+from ..utils.text import count_tokens
+
+class PythonASTChunker(BaseChunker):
+    """
+    AST-Based Python Code Chunker.
+    
+    Uses Python's Abstract Syntax Tree to split code at natural boundaries
+    (classes, functions, imports) rather than arbitrary line counts.
+    
+    BEST-OF-BREED FEATURES:
+    1. Structural Awareness: Splits at class/function boundaries.
+    2. Docstring Preservation: Keeps docstrings with their functions.
+    3. Import Grouping: Groups imports together.
+    4. Nested Handling: Handles nested classes and functions.
+    5. Context Prepending: Optionally prepends module/class context.
+    """
+    name = "python_ast"
+
+    def __init__(self,
+                 include_imports_in_all: bool = True,
+                 split_classes: bool = True,
+                 split_functions: bool = True,
+                 max_tokens: int = 1000,
+                 prepend_context: bool = True):
+        """
+        Initialize the Python AST chunker.
+        
+        Args:
+            include_imports_in_all: If True, prepend imports to every chunk.
+            split_classes: If True, split classes into separate chunks.
+            split_functions: If True, split functions into separate chunks.
+            max_tokens: Maximum tokens per chunk (will further split if exceeded).
+            prepend_context: If True, prepend class name to method chunks.
+        """
+        self.include_imports_in_all = include_imports_in_all
+        self.split_classes = split_classes
+        self.split_functions = split_functions
+        self.max_tokens = max_tokens
+        self.prepend_context = prepend_context
+
+    def chunk(self,
+              doc_id: str,
+              text: str,
+              **params) -> List[Chunk]:
+        """
+        Parse Python code and split at structural boundaries.
+        
+        Args:
+            doc_id: Document identifier
+            text: Python source code
+        
+        Returns:
+            List of Chunk objects
+        """
+        if not text.strip():
+            return []
+        
+        try:
+            tree = ast.parse(text)
+        except SyntaxError:
+            # Fallback to line-based splitting for invalid Python
+            from .recursive_character import RecursiveCharacterChunker
+            return RecursiveCharacterChunker().chunk(doc_id, text, base_token_size=self.max_tokens)
+        
+        lines = text.split("\n")
+        chunks = []
+        
+        # Extract imports
+        imports = []
+        import_lines = set()
+        
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.Import, ast.ImportFrom)):
+                if hasattr(node, 'lineno'):
+                    start = node.lineno - 1
+                    end = getattr(node, 'end_lineno', node.lineno)
+                    import_text = "\n".join(lines[start:end])
+                    imports.append(import_text)
+                    for i in range(start, end):
+                        import_lines.add(i)
+        
+        import_block = "\n".join(imports)
+        
+        # Process top-level definitions
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, ast.ClassDef):
+                chunks.extend(self._process_class(doc_id, node, lines, import_block, len(chunks)))
+            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                chunks.extend(self._process_function(doc_id, node, lines, import_block, len(chunks), None))
+        
+        # If no structures found, treat as single chunk
+        if not chunks:
+            return [Chunk(
+                id=f"{doc_id}#py#0",
+                doc_id=doc_id,
+                text=text,
+                meta={"chunk_index": 0, "strategy": "python_ast", "type": "module"}
+            )]
+        
+        return chunks
+
+    def _process_class(self, doc_id: str, node: ast.ClassDef, lines: List[str], 
+                       import_block: str, start_idx: int) -> List[Chunk]:
+        """Process a class definition."""
+        chunks = []
+        
+        start = node.lineno - 1
+        end = node.end_lineno
+        class_text = "\n".join(lines[start:end])
+        class_name = node.name
+        
+        # Get class docstring
+        docstring = ast.get_docstring(node) or ""
+        
+        if not self.split_classes:
+            # Return whole class as one chunk
+            full_text = class_text
+            if self.include_imports_in_all and import_block:
+                full_text = import_block + "\n\n" + class_text
+            
+            chunks.append(Chunk(
+                id=f"{doc_id}#py#{start_idx}",
+                doc_id=doc_id,
+                text=full_text,
+                meta={
+                    "chunk_index": start_idx,
+                    "strategy": "python_ast",
+                    "type": "class",
+                    "name": class_name,
+                    "docstring": docstring[:200],
+                    "token_count": count_tokens(full_text)
+                }
+            ))
+            return chunks
+        
+        # Process methods within class
+        methods_processed = set()
+        
+        for item in node.body:
+            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                method_chunks = self._process_function(
+                    doc_id, item, lines, import_block, 
+                    start_idx + len(chunks), class_name
+                )
+                chunks.extend(method_chunks)
+                methods_processed.add(item.lineno)
+        
+        # If no methods or class has significant non-method content
+        if not methods_processed:
+            full_text = class_text
+            if self.include_imports_in_all and import_block:
+                full_text = import_block + "\n\n" + class_text
+            
+            chunks.append(Chunk(
+                id=f"{doc_id}#py#{start_idx}",
+                doc_id=doc_id,
+                text=full_text,
+                meta={
+                    "chunk_index": start_idx,
+                    "strategy": "python_ast",
+                    "type": "class",
+                    "name": class_name,
+                    "token_count": count_tokens(full_text)
+                }
+            ))
+        
+        return chunks
+
+    def _process_function(self, doc_id: str, node, lines: List[str],
+                          import_block: str, idx: int, class_name: Optional[str]) -> List[Chunk]:
+        """Process a function definition."""
+        start = node.lineno - 1
+        end = node.end_lineno
+        func_text = "\n".join(lines[start:end])
+        func_name = node.name
+        
+        # Get docstring
+        docstring = ast.get_docstring(node) or ""
+        
+        # Build context prefix
+        context_prefix = ""
+        if self.prepend_context and class_name:
+            context_prefix = f"# Method of class: {class_name}\n"
+        
+        # Build full text
+        full_text = func_text
+        if context_prefix:
+            full_text = context_prefix + full_text
+        if self.include_imports_in_all and import_block:
+            full_text = import_block + "\n\n" + full_text
+        
+        # Check if needs further splitting
+        if count_tokens(full_text) > self.max_tokens:
+            # Split large functions by logical blocks
+            from .recursive_character import RecursiveCharacterChunker
+            sub_chunker = RecursiveCharacterChunker()
+            sub_chunks = sub_chunker.chunk(
+                f"{doc_id}_func_{func_name}", 
+                func_text, 
+                base_token_size=self.max_tokens
+            )
+            
+            chunks = []
+            for i, sc in enumerate(sub_chunks):
+                chunk_text = sc.text
+                if self.include_imports_in_all and import_block and i == 0:
+                    chunk_text = import_block + "\n\n" + chunk_text
+                if context_prefix and i == 0:
+                    chunk_text = context_prefix + chunk_text
+                
+                chunks.append(Chunk(
+                    id=f"{doc_id}#py#{idx + i}",
+                    doc_id=doc_id,
+                    text=chunk_text,
+                    meta={
+                        "chunk_index": idx + i,
+                        "strategy": "python_ast",
+                        "type": "function_part",
+                        "name": func_name,
+                        "class_name": class_name,
+                        "part": i + 1,
+                        "token_count": count_tokens(chunk_text)
+                    }
+                ))
+            return chunks
+        
+        qualified_name = f"{class_name}.{func_name}" if class_name else func_name
+        
+        return [Chunk(
+            id=f"{doc_id}#py#{idx}",
+            doc_id=doc_id,
+            text=full_text,
+            meta={
+                "chunk_index": idx,
+                "strategy": "python_ast",
+                "type": "method" if class_name else "function",
+                "name": func_name,
+                "qualified_name": qualified_name,
+                "class_name": class_name,
+                "docstring": docstring[:200] if docstring else "",
+                "token_count": count_tokens(full_text)
+            }
+        )]