Spaces:
Build error
Build error
| """ | |
| ingestion.py — PDF ingestion and clean text extraction. | |
| """ | |
| import io | |
| import re | |
| import logging | |
| from typing import Union | |
| import pypdf | |
| logger = logging.getLogger("enterprise-rag.ingestion") | |
| def extract_text_from_pdf(pdf_source: Union[str, bytes, io.BytesIO]) -> dict: | |
| """ | |
| Extract text from a PDF file. | |
| Args: | |
| pdf_source: File path string, raw bytes, or BytesIO object. | |
| Returns dict: | |
| text — full extracted text | |
| page_count — number of pages | |
| pages — list of per-page text | |
| metadata — PDF metadata dict | |
| success — bool | |
| error — error message or None | |
| """ | |
| result = { | |
| "text": "", | |
| "page_count": 0, | |
| "pages": [], | |
| "metadata": {}, | |
| "success": False, | |
| "error": None, | |
| } | |
| try: | |
| if isinstance(pdf_source, str): | |
| pdf_file = open(pdf_source, "rb") | |
| elif isinstance(pdf_source, bytes): | |
| pdf_file = io.BytesIO(pdf_source) | |
| else: | |
| pdf_file = pdf_source | |
| reader = pypdf.PdfReader(pdf_file) | |
| result["page_count"] = len(reader.pages) | |
| result["metadata"] = _safe_metadata(reader) | |
| pages_text = [] | |
| full_parts = [] | |
| for page_num, page in enumerate(reader.pages): | |
| try: | |
| raw = page.extract_text() or "" | |
| cleaned = _clean_text(raw) | |
| pages_text.append(cleaned) | |
| if cleaned.strip(): | |
| full_parts.append(cleaned) | |
| except Exception as e: | |
| logger.warning(f"Page {page_num + 1} extraction failed: {e}") | |
| pages_text.append("") | |
| result["pages"] = pages_text | |
| result["text"] = "\n\n".join(full_parts) | |
| result["success"] = bool(result["text"].strip()) | |
| if not result["success"]: | |
| result["error"] = ( | |
| "No readable text found. This PDF may be scanned (image-only). " | |
| "OCR support is not included in this version." | |
| ) | |
| else: | |
| logger.info( | |
| f"Ingested PDF: {result['page_count']} pages, " | |
| f"{len(result['text'])} characters" | |
| ) | |
| except pypdf.errors.PdfReadError as e: | |
| result["error"] = f"PDF is corrupted or password-protected: {e}" | |
| logger.error(result["error"]) | |
| except Exception as e: | |
| result["error"] = f"Ingestion error: {e}" | |
| logger.error(result["error"]) | |
| return result | |
| def _clean_text(text: str) -> str: | |
| """ | |
| Normalize raw PDF text: | |
| - Collapse extra whitespace | |
| - Fix hyphenated line breaks | |
| - Remove lone page-number lines | |
| """ | |
| text = re.sub(r"[ \t]+", " ", text) | |
| text = re.sub(r"-\n(\w)", r"\1", text) | |
| text = re.sub(r"\n{3,}", "\n\n", text) | |
| lines = [ | |
| ln for ln in text.split("\n") | |
| if not (ln.strip().isdigit() and len(ln.strip()) <= 4) | |
| ] | |
| return "\n".join(lines).strip() | |
| def _safe_metadata(reader: pypdf.PdfReader) -> dict: | |
| """Extract PDF metadata fields silently ignoring missing ones.""" | |
| meta = {} | |
| try: | |
| if reader.metadata: | |
| for key in ["/Title", "/Author", "/Subject", "/CreationDate"]: | |
| val = reader.metadata.get(key) | |
| if val: | |
| meta[key.lstrip("/")] = str(val) | |
| except Exception: | |
| pass | |
| return meta | |
| def validate_pdf(file_bytes: bytes, max_mb: int = 30) -> tuple: | |
| """Check file size before processing to avoid memory issues.""" | |
| size_mb = len(file_bytes) / (1024 * 1024) | |
| if size_mb > max_mb: | |
| return False, f"File is {size_mb:.1f}MB — limit is {max_mb}MB." | |
| return True, f"File OK ({size_mb:.1f}MB)" |