""" 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 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}")