Add main orchestrator

doc_enricher/enricher.py

@@ -0,0 +1,239 @@
+"""
+Document Enricher — Main orchestrator.
+
+Coordinates the handler, chunker, and LLM client to produce
+a re-enriched copy of a document with proper heading formatting.
+
+Usage:
+    from doc_enricher import DocumentEnricher, DocxHandler
+    
+    enricher = DocumentEnricher(
+        handler=DocxHandler(),
+        model="llama3",
+        ollama_url="http://localhost:11434",
+    )
+    
+    output_path = enricher.enrich("input.docx", "output.docx")
+"""
+
+import os
+import logging
+import time
+from typing import Optional
+
+from .base_handler import BaseHandler, ParagraphInfo
+from .chunker import build_chunks, Chunk, estimate_tokens
+from .llm_client import OllamaClassifier
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentEnricher:
+    """
+    Orchestrates the document re-enrichment pipeline:
+    
+    1. Extract paragraphs from the original document (via handler)
+    2. Chunk paragraphs into LLM-digestible batches
+    3. Classify each batch using the local LLM
+    4. Apply classifications to a copy of the document (via handler)
+    """
+
+    def __init__(
+        self,
+        handler: BaseHandler,
+        model: str = "llama3",
+        ollama_url: str = "http://localhost:11434",
+        max_tokens_per_chunk: int = 3000,
+        overlap: int = 3,
+        temperature: float = 0.0,
+        num_ctx: int = 8192,
+        timeout_per_request: int = 180,
+        include_formatting_hints: bool = True,
+    ):
+        """
+        Args:
+            handler: Document format handler (e.g. DocxHandler)
+            model: Ollama model name (e.g. "llama3", "llama3:8b")
+            ollama_url: Ollama API base URL
+            max_tokens_per_chunk: Max tokens of paragraph text per LLM call
+            overlap: Paragraphs of overlap between chunks
+            temperature: LLM temperature (0.0 = deterministic)
+            num_ctx: LLM context window size
+            timeout_per_request: HTTP timeout per LLM call in seconds
+            include_formatting_hints: Send existing formatting metadata to LLM
+        """
+        self.handler = handler
+        self.max_tokens_per_chunk = max_tokens_per_chunk
+        self.overlap = overlap
+        self.include_formatting_hints = include_formatting_hints
+
+        self.classifier = OllamaClassifier(
+            model=model,
+            ollama_url=ollama_url,
+            temperature=temperature,
+            num_ctx=num_ctx,
+            timeout=timeout_per_request,
+        )
+
+    def enrich(
+        self,
+        src_path: str,
+        dst_path: Optional[str] = None,
+    ) -> str:
+        """
+        Produce a re-enriched copy of the document.
+        
+        Args:
+            src_path: Path to the original document
+            dst_path: Path for the re-enriched copy.
+                      If None, uses "{name}_enriched.{ext}"
+                      
+        Returns:
+            Path to the re-enriched document
+        """
+        if not os.path.exists(src_path):
+            raise FileNotFoundError(f"Source document not found: {src_path}")
+
+        if dst_path is None:
+            base, ext = os.path.splitext(src_path)
+            dst_path = f"{base}_enriched{ext}"
+
+        logger.info(f"=== Starting re-enrichment: {src_path} → {dst_path} ===")
+        start_time = time.time()
+
+        # Step 1: Extract paragraphs
+        logger.info("Step 1/4: Extracting paragraphs...")
+        paragraphs = self.handler.extract_paragraphs(src_path)
+        logger.info(f"  Found {len(paragraphs)} non-empty paragraphs")
+
+        if not paragraphs:
+            logger.warning("No paragraphs found. Producing unmodified copy.")
+            import shutil
+            shutil.copy2(src_path, dst_path)
+            return dst_path
+
+        # Step 2: Chunk paragraphs
+        logger.info("Step 2/4: Chunking paragraphs...")
+        chunks = build_chunks(
+            paragraphs,
+            max_tokens_per_chunk=self.max_tokens_per_chunk,
+            overlap=self.overlap,
+        )
+        logger.info(f"  Created {len(chunks)} chunk(s)")
+
+        # Build lookup for paragraph info by index
+        para_lookup = {p.index: p for p in paragraphs}
+
+        # Step 3: Classify each chunk
+        logger.info("Step 3/4: Classifying paragraphs with LLM...")
+        all_classifications: dict[int, str] = {}
+
+        for chunk_num, chunk in enumerate(chunks, 1):
+            logger.info(f"  Chunk {chunk_num}/{len(chunks)}: "
+                       f"{len(chunk.all_indices)} paragraphs "
+                       f"({len(chunk.classify_indices)} to classify)")
+
+            # Build the batch for this chunk
+            batch = []
+            for idx in chunk.all_indices:
+                p = para_lookup[idx]
+                entry = {
+                    "index": idx,
+                    "text": p.text,
+                }
+                if self.include_formatting_hints:
+                    entry["style_name"] = p.style_name
+                    entry["is_bold"] = p.is_bold
+                    entry["avg_font_size_pt"] = p.avg_font_size_pt
+                batch.append(entry)
+
+            # Call the LLM
+            try:
+                result = self.classifier.classify_batch(
+                    batch,
+                    formatting_hints=self.include_formatting_hints,
+                )
+            except Exception as e:
+                logger.error(f"  LLM classification failed for chunk {chunk_num}: {e}")
+                logger.info("  Defaulting all paragraphs in this chunk to BODY")
+                for idx in chunk.classify_indices:
+                    if idx not in all_classifications:
+                        all_classifications[idx] = "BODY"
+                continue
+
+            # Store results, but ONLY for primary (non-overlap) indices
+            classify_set = set(chunk.classify_indices)
+            for item in result["classifications"]:
+                idx = item["index"]
+                label = item["label"]
+                if idx in classify_set and idx not in all_classifications:
+                    all_classifications[idx] = label
+
+        # Fill in any missing classifications with BODY
+        for p in paragraphs:
+            if p.index not in all_classifications:
+                logger.warning(f"  Paragraph {p.index} not classified, defaulting to BODY")
+                all_classifications[p.index] = "BODY"
+
+        # Log classification summary
+        label_counts = {}
+        for label in all_classifications.values():
+            label_counts[label] = label_counts.get(label, 0) + 1
+        logger.info(f"  Classification summary: {label_counts}")
+
+        # Step 4: Apply formatting to copy
+        logger.info("Step 4/4: Applying formatting to document copy...")
+        output = self.handler.apply_classifications(
+            src_path, dst_path, all_classifications
+        )
+
+        elapsed = time.time() - start_time
+        logger.info(f"=== Re-enrichment complete in {elapsed:.1f}s: {output} ===")
+
+        return output
+
+    def enrich_batch(
+        self,
+        src_dir: str,
+        dst_dir: str,
+        extension: str = ".docx",
+    ) -> list[str]:
+        """
+        Re-enrich all documents in a directory.
+        
+        Args:
+            src_dir: Directory containing original documents
+            dst_dir: Directory where re-enriched copies will be saved
+            extension: File extension to filter by
+            
+        Returns:
+            List of paths to re-enriched documents
+        """
+        os.makedirs(dst_dir, exist_ok=True)
+
+        src_files = sorted(
+            f for f in os.listdir(src_dir)
+            if f.lower().endswith(extension)
+        )
+
+        if not src_files:
+            logger.warning(f"No {extension} files found in {src_dir}")
+            return []
+
+        logger.info(f"Batch processing {len(src_files)} files from {src_dir}")
+        outputs = []
+
+        for i, filename in enumerate(src_files, 1):
+            src_path = os.path.join(src_dir, filename)
+            dst_path = os.path.join(dst_dir, filename)
+
+            logger.info(f"[{i}/{len(src_files)}] Processing: {filename}")
+            try:
+                output = self.enrich(src_path, dst_path)
+                outputs.append(output)
+            except Exception as e:
+                logger.error(f"  Failed to process {filename}: {e}")
+                continue
+
+        logger.info(f"Batch complete: {len(outputs)}/{len(src_files)} files processed")
+        return outputs