ingestion.py

"""
Document ingestion module for processing PDFs and URLs.
"""
import os
from typing import List, Dict
import requests
from bs4 import BeautifulSoup
from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
import faiss
import pickle


class DocumentIngestion:
    """Handles ingestion of PDFs and URLs into a searchable vector store."""
    
    def __init__(self, embedding_model: str = "all-mpnet-base-v2"):
        """
        Initialize the document ingestion system.
        
        Args:
            embedding_model: Hugging Face model name for embeddings
        """
        self.embedding_model = SentenceTransformer(embedding_model)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=600,
            chunk_overlap=150,
            length_function=len,
        )
        self.documents = []
        self.embeddings = None
        self.index = None
        
    def get_pdf_document_title(self, file_path: str) -> str:
        """
        Get a human-readable document title for a PDF (from metadata or filename).
        
        Args:
            file_path: Path to the PDF file
            
        Returns:
            Document title (e.g. standard name or filename without extension)
        """
        try:
            reader = PdfReader(file_path)
            if reader.metadata and getattr(reader.metadata, "title", None):
                title = reader.metadata.title
                if title and title.strip():
                    return title.strip()
        except Exception:
            pass
        # Fallback: filename without extension, cleaned for standards (e.g. CAN-CGSB-32.312 -> CAN/CGSB-32.312)
        base = os.path.splitext(os.path.basename(file_path))[0]
        if base:
            return base.replace("-", "/") if "CGSB" in base or "CAN" in base else base
        return file_path

    def read_pdf(self, file_path: str) -> str:
        """
        Extract text from a PDF file.
        
        Args:
            file_path: Path to the PDF file
            
        Returns:
            Extracted text content
        """
        try:
            reader = PdfReader(file_path)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text
        except Exception as e:
            raise Exception(f"Error reading PDF {file_path}: {str(e)}")
    
    def read_url(self, url: str):
        """
        Extract text and page title from a URL.

        Args:
            url: URL to fetch and extract text from

        Returns:
            Tuple of (text content, page title or None)
        """
        try:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            response = requests.get(url, headers=headers, timeout=10)
            response.raise_for_status()

            soup = BeautifulSoup(response.content, 'html.parser')

            # Extract page title before stripping elements
            page_title = None
            if soup.title and soup.title.string:
                page_title = soup.title.string.strip()

            # Remove script and style elements
            for script in soup(["script", "style"]):
                script.decompose()

            # Get text
            text = soup.get_text()

            # Clean up whitespace
            lines = (line.strip() for line in text.splitlines())
            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
            text = ' '.join(chunk for chunk in chunks if chunk)

            return text, page_title
        except Exception as e:
            raise Exception(f"Error reading URL {url}: {str(e)}")
    
    def process_documents(self, pdf_paths: List[str] = None, urls: List[str] = None, pdf_urls: Dict[str, str] = None) -> List[Dict]:
        """
        Process PDFs and URLs into chunks.

        Args:
            pdf_paths: List of PDF file paths
            urls: List of URLs to process
            pdf_urls: Optional dict mapping PDF filenames to their public URLs (for hyperlinking references)

        Returns:
            List of document chunks with metadata
        """
        all_texts = []
        all_metadata = []
        pdf_urls = pdf_urls or {}

        # Process PDFs
        if pdf_paths:
            for pdf_path in pdf_paths:
                if not os.path.exists(pdf_path):
                    print(f"Warning: PDF file not found: {pdf_path}")
                    continue
                document_title = self.get_pdf_document_title(pdf_path)
                text = self.read_pdf(pdf_path)
                chunks = self.text_splitter.split_text(text)
                filename = os.path.basename(pdf_path)
                public_url = pdf_urls.get(filename)
                for i, chunk in enumerate(chunks):
                    all_texts.append(chunk)
                    meta = {
                        'source': pdf_path,
                        'document_title': document_title,
                        'type': 'pdf',
                        'chunk_index': i
                    }
                    if public_url:
                        meta['url'] = public_url
                    all_metadata.append(meta)
        
        # Process URLs
        if urls:
            for url in urls:
                try:
                    text, page_title = self.read_url(url)
                    chunks = self.text_splitter.split_text(text)
                    # Use the page's <title> tag if available, otherwise fall back to domain + path
                    if page_title:
                        document_title = page_title
                    else:
                        try:
                            from urllib.parse import urlparse
                            parsed = urlparse(url)
                            document_title = parsed.netloc or url
                            if parsed.path and parsed.path != "/":
                                document_title += parsed.path.rstrip("/")
                        except Exception:
                            document_title = url
                    for i, chunk in enumerate(chunks):
                        all_texts.append(chunk)
                        all_metadata.append({
                            'source': url,
                            'document_title': document_title,
                            'type': 'url',
                            'chunk_index': i
                        })
                except Exception as e:
                    print(f"Warning: Failed to process URL {url}: {str(e)}")
                    continue
        
        # Create document objects
        documents = []
        for text, metadata in zip(all_texts, all_metadata):
            documents.append({
                'text': text,
                'metadata': metadata
            })
        
        self.documents = documents
        return documents
    
    def build_vector_store(self):
        """Build FAISS vector store from processed documents."""
        if not self.documents:
            raise ValueError("No documents processed. Call process_documents() first.")
        
        # Extract texts
        texts = [doc['text'] for doc in self.documents]
        
        # Generate embeddings
        print("Generating embeddings...")
        self.embeddings = self.embedding_model.encode(texts, show_progress_bar=True)
        
        # Build FAISS index
        dimension = self.embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(self.embeddings.astype('float32'))
        
        print(f"Vector store built with {len(self.documents)} documents")
    
    def search(self, query: str, k: int = 5) -> List[Dict]:
        """
        Search for similar documents.
        
        Args:
            query: Search query
            k: Number of results to return
            
        Returns:
            List of relevant document chunks with scores
        """
        if self.index is None:
            raise ValueError("Vector store not built. Call build_vector_store() first.")
        
        # Encode query
        query_embedding = self.embedding_model.encode([query])
        
        # Search
        distances, indices = self.index.search(query_embedding.astype('float32'), k)
        
        # Format results
        results = []
        for i, idx in enumerate(indices[0]):
            if idx < len(self.documents):
                results.append({
                    'text': self.documents[idx]['text'],
                    'metadata': self.documents[idx]['metadata'],
                    'score': float(distances[0][i])
                })
        
        return results
    
    def save(self, directory: str = "data/vector_store"):
        """Save the vector store to disk."""
        os.makedirs(directory, exist_ok=True)
        
        # Save index
        faiss.write_index(self.index, os.path.join(directory, "index.faiss"))
        
        # Save documents and embeddings
        with open(os.path.join(directory, "documents.pkl"), "wb") as f:
            pickle.dump(self.documents, f)
        
        with open(os.path.join(directory, "embeddings.pkl"), "wb") as f:
            pickle.dump(self.embeddings, f)
        
        print(f"Vector store saved to {directory}")
    
    def load(self, directory: str = "data/vector_store"):
        """Load the vector store from disk."""
        # Load index
        self.index = faiss.read_index(os.path.join(directory, "index.faiss"))
        
        # Load documents and embeddings
        with open(os.path.join(directory, "documents.pkl"), "rb") as f:
            self.documents = pickle.load(f)
        
        with open(os.path.join(directory, "embeddings.pkl"), "rb") as f:
            self.embeddings = pickle.load(f)
        
        print(f"Vector store loaded from {directory}")