rnsr 0.1.0__py3-none-any.whl

rnsr/ingestion/pipeline.py

@@ -0,0 +1,936 @@
+"""
+Ingestion Pipeline - Master Function with Enhanced Latent Hierarchy Generation
+
+This module provides the main `ingest_document()` function that implements
+the full Latent Hierarchy Generator from the research paper (Section 4-6).
+
+TIER 1: Visual-Geometric Analysis (Primary)
+    1a. PyMuPDF Font Histogram (Section 6.1)
+        - If headers detected via font variance → Build hierarchical tree
+    1b. Recursive XY-Cut (Section 4.1.1) - Optional for complex layouts
+        - For multi-column documents, L-shaped text wraps
+        
+TIER 2: Semantic Boundary Detection (Fallback 1 - Flat Text)
+    2a. LlamaIndex SemanticSplitterNodeParser (Section 4.2.1)
+        - Embedding-based splitting at topic shifts
+    2b. Hierarchical Clustering (Section 4.2.2) - Enhanced option
+        - Multi-resolution: micro-clusters → macro-clusters
+    2c. Synthetic Header Generation (Section 6.3)
+        - LLM-generated titles for each section
+
+TIER 3: OCR + Re-analyze (Fallback 2 - Scanned PDFs)
+    - Apply Tesseract or Doctr OCR
+    - Generate text layer from images
+    - Build tree from OCR output
+
+ALWAYS call `ingest_document()` - never call individual tiers directly.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import structlog
+
+from rnsr.exceptions import IngestionError
+from rnsr.ingestion.document_boundary import (
+    DocumentBoundaryDetector,
+    segment_by_documents,
+)
+from rnsr.ingestion.font_histogram import FontHistogramAnalyzer
+from rnsr.ingestion.header_classifier import classify_headers
+from rnsr.ingestion.layout_detector import detect_layout_complexity
+from rnsr.ingestion.ocr_fallback import has_extractable_text, try_ocr_ingestion
+from rnsr.ingestion.semantic_fallback import try_semantic_splitter_ingestion
+from rnsr.ingestion.tree_builder import build_document_tree, build_multi_document_tree
+from rnsr.models import DocumentTree, IngestionResult
+
+logger = structlog.get_logger(__name__)
+
+
+def ingest_document(
+    pdf_path: Path | str,
+    use_visual_analysis: bool = True,
+    complexity_threshold: float = 0.3,
+) -> IngestionResult:
+    """
+    Master ingestion function implementing 3-tier graceful degradation.
+    
+    ALWAYS call this function - never call individual tiers directly.
+    
+    Ingestion Flow:
+    0. Pre-analysis: Detect layout complexity (multi-column, empty pages)
+    1. Tier 1a: Font Histogram (simple layouts)
+    1. Tier 1b: LayoutLM + XY-Cut (complex layouts, if use_visual_analysis=True)
+    2. Tier 2: Semantic Splitter (flat text, no structure)
+    3. Tier 3: OCR (scanned/image-only PDFs)
+    
+    Args:
+        pdf_path: Path to the PDF file to ingest.
+        use_visual_analysis: Enable LayoutLM for complex layouts (default: True).
+        complexity_threshold: Threshold for triggering visual analysis (0.0-1.0).
+        
+    Returns:
+        IngestionResult containing the DocumentTree and metadata.
+        
+    Raises:
+        IngestionError: If all tiers fail.
+        
+    Example:
+        # Auto-detect layout complexity
+        result = ingest_document("contract.pdf")
+        
+        # Force visual analysis
+        result = ingest_document("report.pdf", use_visual_analysis=True)
+        
+        # Disable visual analysis
+        result = ingest_document("simple.pdf", use_visual_analysis=False)
+    """
+    pdf_path = Path(pdf_path)
+    
+    if not pdf_path.exists():
+        raise IngestionError(f"PDF file not found: {pdf_path}")
+    
+    logger.info("ingestion_started", path=str(pdf_path))
+    
+    warnings: list[str] = []
+    stats: dict = {"path": str(pdf_path)}
+    
+    # Check if document has extractable text
+    if not has_extractable_text(pdf_path):
+        # No text - go directly to Tier 3 (OCR)
+        logger.info("no_extractable_text", path=str(pdf_path))
+        return _try_tier_3(pdf_path, warnings, stats)
+    
+    # PRE-ANALYSIS: Detect layout complexity
+    if use_visual_analysis:
+        try:
+            complexity = detect_layout_complexity(pdf_path, threshold=complexity_threshold)
+            
+            stats["layout_complexity"] = complexity.complexity_score
+            stats["needs_visual"] = complexity.needs_visual_analysis
+            stats["complexity_reason"] = complexity.reason
+            
+            if complexity.needs_visual_analysis:
+                logger.info(
+                    "complex_layout_detected",
+                    path=str(pdf_path),
+                    score=complexity.complexity_score,
+                    reason=complexity.reason,
+                )
+                
+                # Try visual analysis first
+                result = _try_tier_1b_visual(pdf_path, warnings, stats)
+                if result is not None:
+                    return result
+                
+                # Fall through to standard font histogram if visual fails
+                warnings.append(f"Visual analysis failed, using font histogram fallback")
+        except Exception as e:
+            logger.warning("layout_detection_failed", error=str(e))
+            warnings.append(f"Layout detection failed: {e}")
+    
+    # TIER 1: Try PyMuPDF Font Histogram
+    result = _try_tier_1(pdf_path, warnings, stats)
+    if result is not None:
+        return result
+    
+    # TIER 2: Try Semantic Splitter
+    result = _try_tier_2(pdf_path, warnings, stats)
+    if result is not None:
+        return result
+    
+    # This shouldn't happen, but just in case
+    raise IngestionError("All ingestion tiers failed")
+
+
+def _try_tier_1b_visual(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+) -> IngestionResult | None:
+    """
+    TIER 1b: Try LayoutLM + XY-Cut for complex layouts.
+    
+    Uses visual analysis to detect document structure when
+    layout is too complex for simple font histogram.
+    """
+    logger.debug("trying_tier_1b_visual", path=str(pdf_path))
+    
+    try:
+        from rnsr.ingestion.layout_model import check_layout_model_available
+        
+        if not check_layout_model_available():
+            logger.warning("layout_model_unavailable")
+            warnings.append("LayoutLM not available - falling back to font histogram")
+            return None
+        
+        from rnsr.ingestion.xy_cut import analyze_document_with_xycut
+        
+        # Use XY-Cut + LayoutLM for visual analysis
+        tree = analyze_document_with_xycut(pdf_path)
+        tree.ingestion_tier = 1
+        tree.ingestion_method = "layoutlm_xycut"
+        
+        logger.info(
+            "tier_1b_visual_success",
+            path=str(pdf_path),
+            nodes=tree.total_nodes,
+        )
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=1,
+            method="layoutlm_xycut",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.warning("tier_1b_visual_failed", path=str(pdf_path), error=str(e))
+        warnings.append(f"LayoutLM visual analysis failed: {e}")
+        return None
+
+
+def _try_tier_1(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+    detect_multi_document: bool = True,
+    boundary_confidence: float = 0.5,
+) -> IngestionResult | None:
+    """
+    TIER 1: Try Font Histogram ingestion.
+    
+    Now includes multi-document detection for combined PDFs.
+    
+    Returns None if should fall back to Tier 2.
+    """
+    logger.debug("trying_tier_1", path=str(pdf_path))
+    
+    try:
+        analyzer = FontHistogramAnalyzer()
+        analysis, spans = analyzer.analyze(pdf_path)
+        
+        stats["span_count"] = len(spans)
+        stats["unique_sizes"] = analysis.unique_sizes
+        stats["body_size"] = analysis.body_size
+        
+        # Check if we have font variance
+        if not analyzer.has_font_variance(analysis):
+            logger.info("no_font_variance", path=str(pdf_path))
+            warnings.append("No font variance detected - using semantic splitter")
+            return None  # Trigger Tier 2
+        
+        # Check if we can detect headers
+        if not analyzer.has_detectable_headers(analysis, spans):
+            logger.info("no_headers_detected", path=str(pdf_path))
+            warnings.append("No headers detected - using semantic splitter")
+            return None  # Trigger Tier 2
+        
+        # NEW: Detect document boundaries for multi-document PDFs
+        if detect_multi_document:
+            segments = segment_by_documents(
+                spans, 
+                min_confidence=boundary_confidence,
+            )
+            
+            stats["documents_detected"] = len(segments)
+            
+            if len(segments) > 1:
+                logger.info(
+                    "multi_document_detected",
+                    path=str(pdf_path),
+                    document_count=len(segments),
+                    titles=[s.title[:30] for s in segments],
+                )
+                
+                # Build multi-document tree
+                tree = build_multi_document_tree(
+                    segments,
+                    container_title=pdf_path.stem,
+                )
+                tree.ingestion_tier = 1
+                tree.ingestion_method = "font_histogram"
+                
+                logger.info(
+                    "tier_1_success",
+                    path=str(pdf_path),
+                    nodes=tree.total_nodes,
+                    documents=len(segments),
+                )
+                
+                return IngestionResult(
+                    tree=tree,
+                    tier_used=1,
+                    method="font_histogram",
+                    warnings=warnings,
+                    stats=stats,
+                )
+        
+        # Single document: standard processing
+        # Classify spans
+        classified = classify_headers(spans, analysis)
+        
+        header_count = sum(1 for s in classified if s.role == "header")
+        stats["header_count"] = header_count
+        
+        # Build tree
+        tree = build_document_tree(classified, title=pdf_path.stem)
+        tree.ingestion_tier = 1
+        tree.ingestion_method = "font_histogram"
+        
+        logger.info(
+            "tier_1_success",
+            path=str(pdf_path),
+            nodes=tree.total_nodes,
+        )
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=1,
+            method="font_histogram",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.warning("tier_1_failed", path=str(pdf_path), error=str(e))
+        warnings.append(f"Font histogram failed: {e}")
+        return None
+
+
+def _try_tier_2(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+    use_hierarchical_clustering: bool = False,
+) -> IngestionResult | None:
+    """
+    TIER 2: Try Semantic Splitter or Hierarchical Clustering ingestion.
+    
+    Implements Section 4.2 of the research paper:
+    - 4.2.1: SemanticSplitterNodeParser for breakpoint detection
+    - 4.2.2: Hierarchical Clustering for multi-resolution topics
+    - 6.3: Synthetic Header Generation via LLM
+    """
+    logger.debug("trying_tier_2", path=str(pdf_path))
+    
+    # Option: Use hierarchical clustering for richer structure
+    if use_hierarchical_clustering:
+        try:
+            from rnsr.ingestion.hierarchical_cluster import cluster_document_hierarchically
+            
+            tree = cluster_document_hierarchically(pdf_path)
+            
+            logger.info(
+                "tier_2_hierarchical_success",
+                path=str(pdf_path),
+                nodes=tree.total_nodes,
+            )
+            
+            return IngestionResult(
+                tree=tree,
+                tier_used=2,
+                method="hierarchical_clustering",
+                warnings=warnings,
+                stats=stats,
+            )
+        except Exception as e:
+            logger.warning("hierarchical_clustering_failed", error=str(e))
+            warnings.append(f"Hierarchical clustering failed: {e}")
+            # Fall through to semantic splitter
+    
+    # Default: Semantic Splitter (with LLM-generated headers)
+    try:
+        tree = try_semantic_splitter_ingestion(pdf_path)
+        
+        logger.info(
+            "tier_2_success",
+            path=str(pdf_path),
+            nodes=tree.total_nodes,
+        )
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=2,
+            method="semantic_splitter",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.warning("tier_2_failed", path=str(pdf_path), error=str(e))
+        warnings.append(f"Semantic splitter failed: {e}")
+        # Continue to Tier 3
+        return _try_tier_3(pdf_path, warnings, stats)
+
+
+def _try_tier_3(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+) -> IngestionResult:
+    """
+    TIER 3: Try OCR ingestion (last resort).
+    """
+    logger.debug("trying_tier_3", path=str(pdf_path))
+    
+    try:
+        tree = try_ocr_ingestion(pdf_path)
+        
+        logger.info(
+            "tier_3_success",
+            path=str(pdf_path),
+            nodes=tree.total_nodes,
+        )
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=3,
+            method="ocr",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.error("tier_3_failed", path=str(pdf_path), error=str(e))
+        raise IngestionError(f"All ingestion tiers failed. Last error: {e}") from e
+
+
+def ingest_document_enhanced(
+    pdf_path: Path | str,
+    use_xy_cut: bool = False,
+    use_hierarchical_clustering: bool = False,
+) -> IngestionResult:
+    """
+    Enhanced ingestion with all research paper features.
+    
+    This exposes the full Latent Hierarchy Generator from the paper:
+    - XY-Cut for complex multi-column layouts (Section 4.1.1)
+    - Hierarchical Clustering for multi-resolution topics (Section 4.2.2)
+    - Synthetic Header Generation via LLM (Section 6.3)
+    
+    Args:
+        pdf_path: Path to the PDF file to ingest.
+        use_xy_cut: Enable Recursive XY-Cut for complex layouts.
+        use_hierarchical_clustering: Use clustering instead of simple splits.
+        
+    Returns:
+        IngestionResult containing the DocumentTree and metadata.
+        
+    Example:
+        # For a complex multi-column PDF:
+        result = ingest_document_enhanced("report.pdf", use_xy_cut=True)
+        
+        # For flat text that needs hierarchical structure:
+        result = ingest_document_enhanced(
+            "transcript.pdf",
+            use_hierarchical_clustering=True
+        )
+    """
+    pdf_path = Path(pdf_path)
+    
+    if not pdf_path.exists():
+        raise IngestionError(f"PDF file not found: {pdf_path}")
+    
+    logger.info(
+        "enhanced_ingestion_started",
+        path=str(pdf_path),
+        xy_cut=use_xy_cut,
+        hierarchical=use_hierarchical_clustering,
+    )
+    
+    warnings: list[str] = []
+    stats: dict = {"path": str(pdf_path)}
+    
+    # Check if document has extractable text
+    if not has_extractable_text(pdf_path):
+        return _try_tier_3(pdf_path, warnings, stats)
+    
+    # Try XY-Cut first if enabled (for complex layouts)
+    if use_xy_cut:
+        result = _try_xy_cut_ingestion(pdf_path, warnings, stats)
+        if result is not None:
+            return result
+    
+    # TIER 1: Try Font Histogram
+    result = _try_tier_1(pdf_path, warnings, stats)
+    if result is not None:
+        return result
+    
+    # TIER 2: Semantic analysis with optional hierarchical clustering
+    result = _try_tier_2(pdf_path, warnings, stats, use_hierarchical_clustering)
+    if result is not None:
+        return result
+    
+    raise IngestionError("All ingestion tiers failed")
+
+
+def _try_xy_cut_ingestion(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+) -> IngestionResult | None:
+    """
+    Optional: Use Recursive XY-Cut + LayoutLM for complex layouts.
+    
+    Implements Section 4.1.1:
+    "A top-down page segmentation technique that is particularly 
+    effective for discovering document structure."
+    """
+    logger.debug("trying_xy_cut_with_layoutlm", path=str(pdf_path))
+    
+    try:
+        # Check if LayoutLM is available
+        from rnsr.ingestion.layout_model import check_layout_model_available
+        
+        if not check_layout_model_available():
+            logger.warning("layout_model_unavailable")
+            warnings.append("LayoutLM not available for XY-Cut enhancement")
+            return None
+        
+        from rnsr.ingestion.xy_cut import analyze_document_with_xycut
+        
+        # Use XY-Cut + LayoutLM for visual analysis
+        tree = analyze_document_with_xycut(pdf_path)
+        tree.ingestion_tier = 1
+        tree.ingestion_method = "xy_cut_layoutlm"
+        
+        logger.info(
+            "xy_cut_layoutlm_success",
+            path=str(pdf_path),
+            nodes=tree.total_nodes,
+        )
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=1,
+            method="xy_cut_layoutlm",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.warning("xy_cut_layoutlm_failed", path=str(pdf_path), error=str(e))
+        warnings.append(f"XY-Cut + LayoutLM failed: {e}")
+        return None
+
+
+def _try_xy_cut_ingestion_legacy(
+    pdf_path: Path,
+    warnings: list[str],
+    stats: dict,
+) -> IngestionResult | None:
+    """
+    Legacy XY-Cut implementation without LayoutLM.
+    
+    Implements Section 4.1.1:
+    "A top-down page segmentation technique that is particularly 
+    effective for discovering document structure."
+    """
+    logger.debug("trying_xy_cut", path=str(pdf_path))
+    
+    try:
+        from rnsr.ingestion.xy_cut import RecursiveXYCutter
+        import fitz
+        
+        cutter = RecursiveXYCutter()
+        page_trees = cutter.segment_pdf(pdf_path)
+        
+        # Extract text for each leaf region
+        doc = fitz.open(pdf_path)
+        for page_num, tree in enumerate(page_trees):
+            cutter.extract_text_in_regions(doc[page_num], tree)
+        doc.close()
+        
+        # Convert XY-Cut tree to DocumentTree
+        from rnsr.models import DocumentNode, DocumentTree
+        
+        root = DocumentNode(
+            id="root",
+            level=0,
+            header=pdf_path.stem,
+        )
+        
+        section_num = 0
+        for page_tree in page_trees:
+            for leaf in _get_xy_cut_leaves(page_tree):
+                if leaf.text.strip():
+                    section_num += 1
+                    # Generate synthetic header
+                    from rnsr.ingestion.semantic_fallback import _generate_synthetic_header
+                    
+                    section = DocumentNode(
+                        id=f"xycut_{section_num:03d}",
+                        level=1,
+                        header=_generate_synthetic_header(leaf.text, section_num),
+                        content=leaf.text,
+                    )
+                    root.children.append(section)
+        
+        if section_num == 0:
+            warnings.append("XY-Cut found no text regions")
+            return None
+        
+        tree = DocumentTree(
+            title=pdf_path.stem,
+            root=root,
+            total_nodes=section_num + 1,
+            ingestion_tier=1,
+            ingestion_method="xy_cut",
+        )
+        
+        logger.info("xy_cut_success", path=str(pdf_path), nodes=tree.total_nodes)
+        
+        return IngestionResult(
+            tree=tree,
+            tier_used=1,
+            method="xy_cut",
+            warnings=warnings,
+            stats=stats,
+        )
+        
+    except Exception as e:
+        logger.warning("xy_cut_failed", path=str(pdf_path), error=str(e))
+        warnings.append(f"XY-Cut failed: {e}")
+        return None
+
+
+def _get_xy_cut_leaves(node) -> list:
+    """Get all leaf nodes from an XY-Cut segment tree."""
+    if node.is_leaf:
+        return [node]
+    leaves = []
+    for child in node.children:
+        leaves.extend(_get_xy_cut_leaves(child))
+    return leaves
+
+
+# =============================================================================
+# Entity Extraction Integration
+# =============================================================================
+
+
+def extract_entities_from_tree(
+    tree: DocumentTree,
+    doc_id: str | None = None,
+    extract_relationships: bool = True,
+    max_nodes: int = 100,
+    sample_strategy: str = "important",
+) -> dict:
+    """
+    Extract entities and relationships from an ingested document tree.
+    
+    Uses the RLM Unified Extractor - the most accurate approach:
+    1. LLM writes extraction code based on document
+    2. Code executes on DOC_VAR (grounded in actual text)
+    3. ToT validates with probabilities
+    4. Cross-validation between entities and relationships
+    
+    Args:
+        tree: The ingested DocumentTree.
+        doc_id: Document ID (defaults to tree.id).
+        extract_relationships: Whether to also extract relationships.
+        max_nodes: Maximum nodes to process (for large documents).
+        sample_strategy: How to select nodes - "important" (headers first), 
+                        "uniform" (evenly spaced), or "all" (process all).
+        
+    Returns:
+        Dictionary containing:
+        - entities: List of extracted Entity objects
+        - relationships: List of extracted Relationship objects
+        - stats: Extraction statistics
+        
+    Example:
+        result = ingest_document("contract.pdf")
+        extraction = extract_entities_from_tree(result.tree)
+        
+        # Store in knowledge graph
+        for entity in extraction["entities"]:
+            kg.add_entity(entity)
+    """
+    from rnsr.extraction import (
+        RLMUnifiedExtractor,
+        merge_entities,
+    )
+    
+    doc_id = doc_id or tree.id
+    
+    # Use RLM Unified Extractor (LLM writes code + ToT validation)
+    extractor = RLMUnifiedExtractor(
+        enable_type_learning=True,
+        enable_tot_validation=True,
+        enable_cross_validation=True,
+    )
+    
+    all_entities = []
+    all_relationships = []
+    
+    # Collect all nodes for processing
+    all_nodes = _collect_nodes(tree.root, doc_id)
+    
+    # Sample nodes if too many
+    if sample_strategy != "all" and len(all_nodes) > max_nodes:
+        nodes_to_process = _sample_nodes(all_nodes, max_nodes, sample_strategy)
+        logger.info(
+            "entity_extraction_sampling",
+            total_nodes=len(all_nodes),
+            sampled_nodes=len(nodes_to_process),
+            strategy=sample_strategy,
+        )
+    else:
+        nodes_to_process = all_nodes
+    
+    logger.info(
+        "entity_extraction_started",
+        doc_id=doc_id,
+        node_count=len(nodes_to_process),
+        total_nodes=len(all_nodes),
+        extractor="RLMUnifiedExtractor",
+    )
+    
+    # Process nodes in batches for efficiency
+    batch_size = 10
+    processed = 0
+    
+    for i in range(0, len(nodes_to_process), batch_size):
+        batch = nodes_to_process[i:i + batch_size]
+        
+        for node_data in batch:
+            try:
+                result = extractor.extract(
+                    node_id=node_data["node_id"],
+                    doc_id=doc_id,
+                    header=node_data["header"],
+                    content=node_data["content"],
+                    page_num=node_data.get("page_num"),
+                )
+                
+                if result.entities:
+                    all_entities.extend(result.entities)
+                
+                if extract_relationships and result.relationships:
+                    all_relationships.extend(result.relationships)
+                    
+            except Exception as e:
+                logger.warning(
+                    "node_extraction_failed",
+                    node_id=node_data.get("node_id"),
+                    error=str(e)[:100],
+                )
+        
+        processed += len(batch)
+        if processed % 20 == 0:
+            logger.info(
+                "entity_extraction_progress",
+                processed=processed,
+                total=len(nodes_to_process),
+                entities_so_far=len(all_entities),
+            )
+    
+    # Merge duplicate entities
+    merged_entities = merge_entities(all_entities)
+    
+    stats = {
+        "nodes_processed": len(nodes_to_process),
+        "entities_extracted": len(all_entities),
+        "entities_after_merge": len(merged_entities),
+        "relationships_extracted": len(all_relationships),
+        "entity_types": _count_entity_types(merged_entities),
+        "extraction_method": "rlm_unified",
+    }
+    
+    logger.info(
+        "entity_extraction_complete",
+        doc_id=doc_id,
+        **stats,
+    )
+    
+    return {
+        "entities": merged_entities,
+        "relationships": all_relationships,
+        "stats": stats,
+    }
+
+
+def _collect_nodes(node, doc_id: str, collected: list | None = None) -> list[dict]:
+    """
+    Recursively collect all nodes from a DocumentNode tree.
+    
+    Args:
+        node: Root DocumentNode.
+        doc_id: Document ID.
+        collected: List to collect into.
+        
+    Returns:
+        List of node data dictionaries.
+    """
+    if collected is None:
+        collected = []
+    
+    # Add this node if it has content
+    if node.content or node.header:
+        collected.append({
+            "node_id": node.id,
+            "header": node.header,
+            "content": node.content,
+            "page_num": node.page_num,
+            "level": node.level,
+        })
+    
+    # Process children
+    for child in node.children:
+        _collect_nodes(child, doc_id, collected)
+    
+    return collected
+
+
+def _sample_nodes(nodes: list[dict], max_nodes: int, strategy: str) -> list[dict]:
+    """
+    Sample nodes from a large document for efficient processing.
+    
+    Args:
+        nodes: All collected nodes.
+        max_nodes: Maximum number of nodes to return.
+        strategy: Sampling strategy - "important", "uniform", or "first".
+        
+    Returns:
+        Sampled list of nodes.
+    """
+    if len(nodes) <= max_nodes:
+        return nodes
+    
+    if strategy == "important":
+        # Prioritize nodes with headers and higher-level sections
+        scored_nodes = []
+        for node in nodes:
+            score = 0
+            
+            # Prefer nodes with headers
+            if node.get("header"):
+                score += 10
+            
+            # Prefer higher-level (lower number) sections
+            level = node.get("level", 3)
+            score += max(0, 5 - level)
+            
+            # Prefer nodes with substantial content
+            content_len = len(node.get("content", ""))
+            if content_len > 500:
+                score += 3
+            elif content_len > 200:
+                score += 2
+            elif content_len > 50:
+                score += 1
+            
+            scored_nodes.append((score, node))
+        
+        # Sort by score (descending) and take top nodes
+        scored_nodes.sort(key=lambda x: x[0], reverse=True)
+        return [node for _, node in scored_nodes[:max_nodes]]
+    
+    elif strategy == "uniform":
+        # Evenly sample across the document
+        step = len(nodes) // max_nodes
+        return [nodes[i] for i in range(0, len(nodes), step)][:max_nodes]
+    
+    else:  # "first"
+        return nodes[:max_nodes]
+
+
+def _count_entity_types(entities: list) -> dict[str, int]:
+    """Count entities by type."""
+    counts: dict[str, int] = {}
+    for entity in entities:
+        type_name = entity.type.value
+        counts[type_name] = counts.get(type_name, 0) + 1
+    return counts
+
+
+def ingest_with_entities(
+    pdf_path: Path | str,
+    knowledge_graph=None,
+    extract_relationships: bool = True,
+    link_entities: bool = True,
+    **ingest_kwargs,
+) -> dict:
+    """
+    Ingest a document and extract entities in a single operation.
+    
+    Combines document ingestion with entity extraction, optionally
+    storing results in a knowledge graph.
+    
+    Args:
+        pdf_path: Path to the PDF file.
+        knowledge_graph: Optional KnowledgeGraph to store entities.
+        extract_relationships: Whether to extract relationships.
+        link_entities: Whether to link entities across documents.
+        **ingest_kwargs: Additional arguments for ingest_document.
+        
+    Returns:
+        Dictionary containing:
+        - ingestion_result: The IngestionResult
+        - extraction: Entity extraction results
+        - links: Entity links (if link_entities=True)
+        
+    Example:
+        from rnsr.indexing.knowledge_graph import KnowledgeGraph
+        
+        kg = KnowledgeGraph("./data/kg.db")
+        result = ingest_with_entities("contract.pdf", knowledge_graph=kg)
+        
+        print(f"Found {len(result['extraction']['entities'])} entities")
+    """
+    # Ingest document
+    ingestion_result = ingest_document(pdf_path, **ingest_kwargs)
+    
+    # Extract entities
+    extraction = extract_entities_from_tree(
+        tree=ingestion_result.tree,
+        extract_relationships=extract_relationships,
+    )
+    
+    result = {
+        "ingestion_result": ingestion_result,
+        "extraction": extraction,
+        "links": [],
+    }
+    
+    # Store in knowledge graph if provided
+    if knowledge_graph is not None:
+        from rnsr.extraction import EntityLinker
+        
+        # Store entities
+        for entity in extraction["entities"]:
+            knowledge_graph.add_entity(entity)
+        
+        # Store relationships
+        for relationship in extraction["relationships"]:
+            knowledge_graph.add_relationship(relationship)
+        
+        # Link entities if enabled
+        if link_entities:
+            linker = EntityLinker(knowledge_graph)
+            links = linker.link_all_entities_in_document(ingestion_result.tree.id)
+            result["links"] = links
+        
+        logger.info(
+            "stored_in_knowledge_graph",
+            doc_id=ingestion_result.tree.id,
+            entities=len(extraction["entities"]),
+            relationships=len(extraction["relationships"]),
+            links=len(result["links"]),
+        )
+    
+    return result
+
+
+# Convenience exports
+__all__ = [
+    "ingest_document",
+    "ingest_document_enhanced",
+    "ingest_with_entities",
+    "extract_entities_from_tree",
+    "IngestionResult",
+]