Kelvin-programmer/rag-chatbot

1"""PDF text extraction and chunking utilities."""2 3import logging4from pathlib import Path5 6from pypdf import PdfReader7 8logger = logging.getLogger(__name__)9 10 11class PDFProcessor:12    """Extracts and chunks text from PDF documents."""13 14    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):15        self.chunk_size = chunk_size16        self.chunk_overlap = chunk_overlap17 18    @staticmethod19    def clean_text(text: str) -> str:20        """Collapse whitespace and strip surrounding blanks."""21        return " ".join(text.split())22 23    def extract_chunks(24        self,25        pdf_path: str,26        start_page: int = 0,27        end_page: int | None = None,28    ) -> list[dict]:29        """Extract text chunks from a PDF with page-level metadata.30 31        Returns a list of dicts, each with ``text`` and ``metadata`` keys.32        """33        reader = PdfReader(pdf_path)34        filename = Path(pdf_path).name35        pages = reader.pages[start_page:end_page]36 37        results: list[dict] = []38        for page_num, page in enumerate(pages, start_page):39            text = self.clean_text(page.extract_text() or "")40            if len(text) < 50:41                continue42 43            chunks = self._split_text(text)44            for i, chunk in enumerate(chunks):45                results.append(46                    {47                        "text": chunk,48                        "metadata": {49                            "source": filename,50                            "page": page_num + 1,51                            "chunk_index": i,52                        },53                    }54                )55 56        logger.info(57            "Extracted %d chunks from %s (%d pages)",58            len(results),59            filename,60            len(pages),61        )62        return results63 64    def _split_text(self, text: str) -> list[str]:65        """Split *text* into overlapping chunks."""66        step = max(self.chunk_size - self.chunk_overlap, 1)67        chunks = []68        for i in range(0, len(text), step):69            chunk = text[i : i + self.chunk_size].strip()70            if chunk:71                chunks.append(chunk)72        return chunks73