pipeline.py

"""Main encoding pipeline orchestrator."""
import time
from pathlib import Path
from typing import List, Optional
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from rich.table import Table

from config import PipelineConfig, DEFAULT_CONFIG
from document_loader import load_book, DocumentSection
from semantic_chunker import SemanticChunker, ChonkieSemanticChunker, Chunk
from embeddings import EmbeddingEngine, ContextualEmbedder, EmbeddedChunk
from vector_store import VectorStore

console = Console()


class BookEncodingPipeline:
    """
    End-to-end pipeline for encoding a cybersecurity book into a vector database.
    
    Pipeline stages:
    1. Load document and extract structure
    2. Semantic chunking with parent-document strategy
    3. Generate high-dimensional embeddings
    4. Store in vector database with metadata
    """
    
    def __init__(self, config: PipelineConfig = DEFAULT_CONFIG):
        self.config = config
        
        # Initialize components lazily
        self._chunker = None
        self._embedding_engine = None
        self._contextual_embedder = None
        self._vector_store = None
        
        # Pipeline state
        self.sections: List[DocumentSection] = []
        self.child_chunks: List[Chunk] = []
        self.parent_chunks: List[Chunk] = []
        self.embedded_chunks: List[EmbeddedChunk] = []
    
    @property
    def chunker(self) -> SemanticChunker:
        if self._chunker is None:
            self._chunker = ChonkieSemanticChunker(
                child_chunk_size=self.config.chunking.chunk_size,
                parent_chunk_size=self.config.chunking.parent_chunk_size,
                overlap_ratio=self.config.chunking.overlap,
                min_chunk_size=self.config.chunking.min_chunk_size
            )
        return self._chunker
    
    @property
    def embedding_engine(self) -> EmbeddingEngine:
        if self._embedding_engine is None:
            self._embedding_engine = EmbeddingEngine(
                model_name=self.config.embedding.model_name,
                device=self.config.embedding.device,
                normalize=self.config.embedding.normalize,
                use_hyperbolic=self.config.use_hyperbolic_embeddings
            )
        return self._embedding_engine
    
    @property
    def contextual_embedder(self) -> ContextualEmbedder:
        if self._contextual_embedder is None:
            self._contextual_embedder = ContextualEmbedder(self.embedding_engine)
        return self._contextual_embedder
    
    @property
    def vector_store(self) -> VectorStore:
        if self._vector_store is None:
            self._vector_store = VectorStore(
                collection_name=self.config.vector_db.collection_name,
                persist_directory=self.config.vector_db.persist_directory,
                distance_metric=self.config.vector_db.distance_metric
            )
        return self._vector_store
    
    def run(self, book_path: str) -> None:
        """
        Run the complete encoding pipeline.
        
        Args:
            book_path: Path to the book (PDF or text file)
        """
        console.print("\n[bold blue]═══════════════════════════════════════════════════════════════[/bold blue]")
        console.print("[bold blue]       CYBERSECURITY BOOK ENCODING PIPELINE[/bold blue]")
        console.print("[bold blue]═══════════════════════════════════════════════════════════════[/bold blue]\n")
        
        start_time = time.time()
        
        # Stage 1: Load document
        self._stage_load(book_path)
        
        # Stage 2: Chunk document
        self._stage_chunk()
        
        # Stage 3: Generate embeddings
        self._stage_embed()
        
        # Stage 4: Store in vector database
        self._stage_store()
        
        # Print summary
        elapsed = time.time() - start_time
        self._print_summary(elapsed)
    
    def _stage_load(self, book_path: str) -> None:
        """Stage 1: Load and parse the document."""
        console.print("[bold cyan]Stage 1:[/bold cyan] Loading Document\n")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console
        ) as progress:
            task = progress.add_task("Parsing document structure...", total=None)
            
            self.sections = load_book(book_path)
            
            progress.update(task, completed=True)
        
        console.print(f"  ✓ Loaded [green]{len(self.sections)}[/green] sections")
        
        # Show chapter breakdown
        chapters = set(s.chapter for s in self.sections if s.chapter)
        if chapters:
            console.print(f"  ✓ Detected [green]{len(chapters)}[/green] chapters")
        console.print()
    
    def _stage_chunk(self) -> None:
        """Stage 2: Semantic chunking with parent-document strategy."""
        console.print("[bold cyan]Stage 2:[/bold cyan] Semantic Chunking\n")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=console
        ) as progress:
            task = progress.add_task("Creating semantic chunks...", total=len(self.sections))
            
            self.child_chunks, self.parent_chunks = self.chunker.chunk_sections(self.sections)
            
            progress.update(task, completed=len(self.sections))
        
        console.print(f"  ✓ Created [green]{len(self.child_chunks)}[/green] child chunks (for retrieval)")
        console.print(f"  ✓ Created [green]{len(self.parent_chunks)}[/green] parent chunks (for context)")
        
        # Show average chunk sizes
        avg_child = sum(len(c.content) for c in self.child_chunks) / max(len(self.child_chunks), 1)
        avg_parent = sum(len(c.content) for c in self.parent_chunks) / max(len(self.parent_chunks), 1)
        console.print(f"  ✓ Average child chunk: [green]{avg_child:.0f}[/green] chars")
        console.print(f"  ✓ Average parent chunk: [green]{avg_parent:.0f}[/green] chars")
        console.print()
    
    def _stage_embed(self) -> None:
        """Stage 3: Generate embeddings with contextual enrichment."""
        console.print("[bold cyan]Stage 3:[/bold cyan] Generating Embeddings\n")
        
        console.print(f"  Model: [yellow]{self.config.embedding.model_name}[/yellow]")
        console.print(f"  Dimension: [yellow]{self.embedding_engine.dimension}[/yellow]")
        
        if self.config.use_hyperbolic_embeddings:
            console.print(f"  Hyperbolic projection: [yellow]Enabled[/yellow]")
        
        console.print()
        
        # Embed child chunks (these are what we search)
        console.print("  Embedding child chunks...")
        child_embedded = self.contextual_embedder.embed_chunks(
            self.child_chunks,
            batch_size=self.config.embedding.batch_size
        )
        
        # Embed parent chunks (for context storage)
        console.print("  Embedding parent chunks...")
        parent_embedded = self.contextual_embedder.embed_chunks(
            self.parent_chunks,
            batch_size=self.config.embedding.batch_size
        )
        
        self.embedded_chunks = child_embedded + parent_embedded
        
        console.print(f"\n  ✓ Generated [green]{len(self.embedded_chunks)}[/green] embeddings")
        console.print()
    
    def _stage_store(self) -> None:
        """Stage 4: Store embeddings in vector database."""
        console.print("[bold cyan]Stage 4:[/bold cyan] Storing in Vector Database\n")
        
        console.print(f"  Database: [yellow]ChromaDB[/yellow]")
        console.print(f"  Collection: [yellow]{self.config.vector_db.collection_name}[/yellow]")
        console.print(f"  Path: [yellow]{self.config.vector_db.persist_directory}[/yellow]")
        console.print()
        
        import numpy as np
        
        # Prepare data for storage
        chunk_ids = [ec.chunk_id for ec in self.embedded_chunks]
        embeddings = np.array([ec.embedding for ec in self.embedded_chunks])
        contents = [ec.content for ec in self.embedded_chunks]
        metadatas = [ec.metadata for ec in self.embedded_chunks]
        
        # Store in vector database
        self.vector_store.add_chunks(
            chunk_ids=chunk_ids,
            embeddings=embeddings,
            contents=contents,
            metadatas=metadatas
        )
        
        console.print(f"  ✓ Stored [green]{self.vector_store.get_chunk_count()}[/green] searchable chunks")
        console.print()
    
    def _print_summary(self, elapsed: float) -> None:
        """Print pipeline summary."""
        console.print("[bold blue]═══════════════════════════════════════════════════════════════[/bold blue]")
        console.print("[bold green]                    ENCODING COMPLETE![/bold green]")
        console.print("[bold blue]═══════════════════════════════════════════════════════════════[/bold blue]\n")
        
        table = Table(title="Pipeline Summary")
        table.add_column("Metric", style="cyan")
        table.add_column("Value", style="green")
        
        table.add_row("Total Sections", str(len(self.sections)))
        table.add_row("Child Chunks", str(len(self.child_chunks)))
        table.add_row("Parent Chunks", str(len(self.parent_chunks)))
        table.add_row("Total Embeddings", str(len(self.embedded_chunks)))
        table.add_row("Vector Dimension", str(self.embedding_engine.dimension))
        table.add_row("Time Elapsed", f"{elapsed:.2f}s")
        
        console.print(table)
        console.print()


def encode_book(book_path: str, config: Optional[PipelineConfig] = None) -> VectorStore:
    """
    Convenience function to encode a book.
    
    Args:
        book_path: Path to the book file
        config: Optional pipeline configuration
        
    Returns:
        The populated VectorStore instance
    """
    if config is None:
        config = DEFAULT_CONFIG
    
    pipeline = BookEncodingPipeline(config)
    pipeline.run(book_path)
    
    return pipeline.vector_store


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        console.print("[red]Usage: python pipeline.py <path_to_book>[/red]")
        sys.exit(1)
    
    book_path = sys.argv[1]
    encode_book(book_path)