Spaces:
Sleeping
Sleeping
| import json | |
| import re | |
| from html import escape | |
| from html.parser import HTMLParser | |
| import logging | |
| import shutil | |
| import subprocess | |
| import threading | |
| import time | |
| import uuid | |
| import zipfile | |
| from collections import defaultdict | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any | |
| from docling.document_converter import DocumentConverter | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_core.documents import Document | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from pypdf import PdfReader | |
| from reportlab.lib.enums import TA_LEFT | |
| from reportlab.lib.pagesizes import letter | |
| from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet | |
| from reportlab.lib.units import inch | |
| from reportlab.platypus import PageBreak, Paragraph, SimpleDocTemplate, Spacer | |
| from transformers import AutoTokenizer | |
| from config import ( | |
| EMBEDDING_MODEL_NAME, | |
| MAX_UPLOAD_PAGES, | |
| MAX_UPLOAD_SIZE_MB, | |
| UPLOAD_CHUNK_OVERLAP, | |
| UPLOAD_CHUNK_SIZE, | |
| UPLOAD_SESSION_ROOT, | |
| UPLOAD_SESSION_TTL_SECONDS, | |
| UPLOAD_SUPPORTED_EXTENSIONS, | |
| ) | |
| from services import initialize_embeddings, normalize_metadata | |
| logger = logging.getLogger(__name__) | |
| _sessions: dict[str, dict[str, Any]] = {} | |
| _sessions_lock = threading.RLock() | |
| _tokenizer = None | |
| _tokenizer_lock = threading.Lock() | |
| SESSION_METADATA_FILE = "session.json" | |
| NORMALIZED_PDF_NAME = "normalized_document.pdf" | |
| def _get_tokenizer(): | |
| global _tokenizer | |
| if _tokenizer is not None: | |
| return _tokenizer | |
| with _tokenizer_lock: | |
| if _tokenizer is None: | |
| logger.info( | |
| "Loading tokenizer for uploaded-document chunking: %s", | |
| EMBEDDING_MODEL_NAME, | |
| ) | |
| _tokenizer = AutoTokenizer.from_pretrained( | |
| EMBEDDING_MODEL_NAME | |
| ) | |
| return _tokenizer | |
| def _session_metadata_path(session_dir: Path) -> Path: | |
| return session_dir / SESSION_METADATA_FILE | |
| def _write_session_metadata(session: dict[str, Any]) -> None: | |
| session_dir = Path(session["session_dir"]) | |
| payload = { | |
| "session_id": session["session_id"], | |
| "filename": session["filename"], | |
| "original_extension": session["original_extension"], | |
| "processing_extension": session["processing_extension"], | |
| "conversion_method": session["conversion_method"], | |
| "page_reference_basis": session["page_reference_basis"], | |
| "normalized_pdf_path": str(session["normalized_pdf_path"]), | |
| "page_count": int(session["page_count"]), | |
| "chunk_count": int(session["chunk_count"]), | |
| "session_dir": str(session_dir), | |
| "vector_dir": str(session["vector_dir"]), | |
| "zip_path": str(session["zip_path"]), | |
| "created_at": float(session["created_at"]), | |
| "expires_at": float(session["expires_at"]), | |
| } | |
| _session_metadata_path(session_dir).write_text( | |
| json.dumps(payload, indent=2, ensure_ascii=False), | |
| encoding="utf-8", | |
| ) | |
| def _read_session_metadata( | |
| session_dir: Path, | |
| ) -> dict[str, Any] | None: | |
| metadata_path = _session_metadata_path(session_dir) | |
| if not metadata_path.is_file(): | |
| return None | |
| try: | |
| payload = json.loads( | |
| metadata_path.read_text(encoding="utf-8") | |
| ) | |
| except Exception as exc: | |
| logger.warning( | |
| "Could not read uploaded-document session metadata %s: %s", | |
| metadata_path, | |
| exc, | |
| ) | |
| return None | |
| required = { | |
| "session_id", | |
| "filename", | |
| "original_extension", | |
| "processing_extension", | |
| "conversion_method", | |
| "page_reference_basis", | |
| "normalized_pdf_path", | |
| "page_count", | |
| "chunk_count", | |
| "session_dir", | |
| "vector_dir", | |
| "zip_path", | |
| "created_at", | |
| "expires_at", | |
| } | |
| if not required.issubset(payload): | |
| logger.warning( | |
| "Uploaded-document session metadata is incomplete: %s", | |
| metadata_path, | |
| ) | |
| return None | |
| return payload | |
| def _load_vector_db(vector_dir: Path): | |
| faiss_path = vector_dir / "index.faiss" | |
| pickle_path = vector_dir / "index.pkl" | |
| if not faiss_path.is_file() or not pickle_path.is_file(): | |
| raise FileNotFoundError( | |
| "The temporary vector database is incomplete. " | |
| f"Expected '{faiss_path}' and '{pickle_path}'." | |
| ) | |
| return FAISS.load_local( | |
| str(vector_dir), | |
| embeddings=initialize_embeddings(), | |
| allow_dangerous_deserialization=True, | |
| ) | |
| def _restore_session_from_disk( | |
| session_id: str, | |
| ) -> dict[str, Any] | None: | |
| session_dir = UPLOAD_SESSION_ROOT / session_id | |
| payload = _read_session_metadata(session_dir) | |
| if payload is None: | |
| return None | |
| if float(payload["expires_at"]) <= time.time(): | |
| shutil.rmtree(session_dir, ignore_errors=True) | |
| return None | |
| normalized_pdf_path = Path(payload["normalized_pdf_path"]) | |
| if not normalized_pdf_path.is_file(): | |
| logger.warning( | |
| "Normalized PDF is missing for session %s", | |
| session_id, | |
| ) | |
| return None | |
| vector_dir = Path(payload["vector_dir"]) | |
| session = dict(payload) | |
| session["vector_db"] = _load_vector_db(vector_dir) | |
| with _sessions_lock: | |
| _sessions[session_id] = session | |
| logger.info( | |
| "Restored uploaded-document session from disk: %s", | |
| session_id, | |
| ) | |
| return session | |
| def _resolve_uploaded_path(uploaded_file: Any) -> Path: | |
| if uploaded_file is None: | |
| raise ValueError("Please select a document.") | |
| if isinstance(uploaded_file, (str, Path)): | |
| path = Path(uploaded_file) | |
| elif isinstance(uploaded_file, dict): | |
| path = Path( | |
| uploaded_file.get("path") | |
| or uploaded_file.get("name") | |
| or "" | |
| ) | |
| else: | |
| candidate = getattr(uploaded_file, "name", None) | |
| path = Path(candidate or "") | |
| if not path.is_file(): | |
| raise FileNotFoundError( | |
| "The uploaded document could not be read." | |
| ) | |
| return path | |
| def cleanup_expired_sessions() -> int: | |
| """Remove expired sessions from memory and temporary storage.""" | |
| now = time.time() | |
| expired_dirs: set[Path] = set() | |
| UPLOAD_SESSION_ROOT.mkdir(parents=True, exist_ok=True) | |
| with _sessions_lock: | |
| for session_id, session in list(_sessions.items()): | |
| if float(session["expires_at"]) <= now: | |
| expired_dirs.add(Path(session["session_dir"])) | |
| _sessions.pop(session_id, None) | |
| for session_dir in UPLOAD_SESSION_ROOT.iterdir(): | |
| if not session_dir.is_dir(): | |
| continue | |
| payload = _read_session_metadata(session_dir) | |
| if payload is None: | |
| try: | |
| age = now - session_dir.stat().st_mtime | |
| except OSError: | |
| age = UPLOAD_SESSION_TTL_SECONDS + 1 | |
| if age > UPLOAD_SESSION_TTL_SECONDS: | |
| expired_dirs.add(session_dir) | |
| continue | |
| if float(payload["expires_at"]) <= now: | |
| expired_dirs.add(session_dir) | |
| with _sessions_lock: | |
| _sessions.pop( | |
| str(payload["session_id"]), | |
| None, | |
| ) | |
| for session_dir in expired_dirs: | |
| shutil.rmtree(session_dir, ignore_errors=True) | |
| logger.info( | |
| "Removed expired uploaded-document session directory: %s", | |
| session_dir, | |
| ) | |
| return len(expired_dirs) | |
| def _pdf_page_count(pdf_path: Path) -> int: | |
| try: | |
| reader = PdfReader(str(pdf_path)) | |
| if reader.is_encrypted: | |
| try: | |
| reader.decrypt("") | |
| except Exception as exc: | |
| raise ValueError( | |
| "Password-protected PDF files are not supported." | |
| ) from exc | |
| return len(reader.pages) | |
| except ValueError: | |
| raise | |
| except Exception as exc: | |
| raise ValueError( | |
| f"The PDF page count could not be determined: {exc}" | |
| ) from exc | |
| def _find_libreoffice() -> str: | |
| executable = ( | |
| shutil.which("soffice") | |
| or shutil.which("libreoffice") | |
| ) | |
| if not executable: | |
| raise RuntimeError( | |
| "Document normalization requires LibreOffice, but no " | |
| "'soffice' or 'libreoffice' executable was found." | |
| ) | |
| return executable | |
| class _TextFocusedHtmlExtractor(HTMLParser): | |
| """ | |
| Extract readable policy text from HTML without loading web resources. | |
| Headings, paragraphs, lists, and table cells are retained. Scripts, | |
| styles, images, frames, media, forms, advertisements, and other web-only | |
| elements are discarded before LibreOffice receives the document. | |
| """ | |
| SKIP_TAGS = { | |
| "script", | |
| "style", | |
| "iframe", | |
| "object", | |
| "embed", | |
| "video", | |
| "audio", | |
| "canvas", | |
| "noscript", | |
| "svg", | |
| "form", | |
| "button", | |
| } | |
| BLOCK_TAGS = { | |
| "article", | |
| "aside", | |
| "blockquote", | |
| "body", | |
| "dd", | |
| "div", | |
| "dl", | |
| "dt", | |
| "fieldset", | |
| "figcaption", | |
| "figure", | |
| "footer", | |
| "h1", | |
| "h2", | |
| "h3", | |
| "h4", | |
| "h5", | |
| "h6", | |
| "header", | |
| "li", | |
| "main", | |
| "nav", | |
| "ol", | |
| "p", | |
| "pre", | |
| "section", | |
| "table", | |
| "tbody", | |
| "td", | |
| "tfoot", | |
| "th", | |
| "thead", | |
| "tr", | |
| "ul", | |
| } | |
| def __init__(self) -> None: | |
| super().__init__(convert_charrefs=True) | |
| self.skip_depth = 0 | |
| self.parts: list[str] = [] | |
| def _line_break(self) -> None: | |
| if self.parts and not self.parts[-1].endswith("\n"): | |
| self.parts.append("\n") | |
| def handle_starttag( | |
| self, | |
| tag: str, | |
| attrs: list[tuple[str, str | None]], | |
| ) -> None: | |
| tag = tag.lower() | |
| if tag in self.SKIP_TAGS: | |
| self.skip_depth += 1 | |
| return | |
| if self.skip_depth: | |
| return | |
| if tag in self.BLOCK_TAGS: | |
| self._line_break() | |
| if tag == "li": | |
| self.parts.append("• ") | |
| if tag in {"td", "th"}: | |
| self.parts.append(" | ") | |
| if tag == "br": | |
| self._line_break() | |
| def handle_endtag(self, tag: str) -> None: | |
| tag = tag.lower() | |
| if tag in self.SKIP_TAGS: | |
| if self.skip_depth: | |
| self.skip_depth -= 1 | |
| return | |
| if self.skip_depth: | |
| return | |
| if tag in self.BLOCK_TAGS: | |
| self._line_break() | |
| def handle_data(self, data: str) -> None: | |
| if self.skip_depth: | |
| return | |
| cleaned = re.sub( | |
| r"\s+", | |
| " ", | |
| str(data or ""), | |
| ).strip() | |
| if not cleaned: | |
| return | |
| if ( | |
| self.parts | |
| and not self.parts[-1].endswith( | |
| (" ", "\n", "| ") | |
| ) | |
| ): | |
| self.parts.append(" ") | |
| self.parts.append(cleaned) | |
| def text(self) -> str: | |
| raw = "".join(self.parts) | |
| raw = re.sub(r"[ \t]+\n", "\n", raw) | |
| raw = re.sub(r"\n[ \t]+", "\n", raw) | |
| raw = re.sub(r"\n{3,}", "\n\n", raw) | |
| return raw.strip() | |
| def _prepare_local_html( | |
| source_path: Path, | |
| session_dir: Path, | |
| ) -> Path: | |
| """ | |
| Build a small self-contained HTML document for predictable PDF conversion. | |
| This avoids LibreOffice waiting for external JavaScript, stylesheets, | |
| images, fonts, frames, advertisements, or tracking resources. | |
| """ | |
| raw = source_path.read_text( | |
| encoding="utf-8", | |
| errors="replace", | |
| ) | |
| extractor = _TextFocusedHtmlExtractor() | |
| extractor.feed(raw) | |
| extractor.close() | |
| extracted_text = extractor.text() | |
| if not extracted_text: | |
| raise ValueError( | |
| "The HTML file did not contain readable text after removing " | |
| "scripts, styles, and web-only resources." | |
| ) | |
| paragraphs = [] | |
| for block in re.split( | |
| r"\n\s*\n", | |
| extracted_text, | |
| ): | |
| block = block.strip() | |
| if not block: | |
| continue | |
| escaped_block = escape( | |
| block, | |
| quote=False, | |
| ).replace("\n", "<br>") | |
| paragraphs.append( | |
| f"<p>{escaped_block}</p>" | |
| ) | |
| normalized_html = ( | |
| "<!DOCTYPE html>" | |
| "<html><head>" | |
| '<meta charset="utf-8">' | |
| "<style>" | |
| "@page{size:Letter;margin:0.65in}" | |
| "body{font-family:Arial,sans-serif;" | |
| "font-size:10.5pt;line-height:1.45;color:#111}" | |
| "p{margin:0 0 9pt;page-break-inside:auto}" | |
| "</style>" | |
| "</head><body>" | |
| + "".join(paragraphs) | |
| + "</body></html>" | |
| ) | |
| normalized_path = ( | |
| session_dir | |
| / f"{source_path.stem}_text_normalized.html" | |
| ) | |
| normalized_path.write_text( | |
| normalized_html, | |
| encoding="utf-8", | |
| ) | |
| return normalized_path | |
| def _render_html_text_to_pdf( | |
| source_path: Path, | |
| session_dir: Path, | |
| ) -> Path: | |
| """ | |
| Render readable HTML text directly to a stable normalized PDF. | |
| This bypasses LibreOffice entirely for HTML files. The HTML parser removes | |
| scripts, styles, images, frames, media, forms, and external web resources, | |
| then ReportLab paginates the remaining policy text. | |
| """ | |
| raw = source_path.read_text( | |
| encoding="utf-8", | |
| errors="replace", | |
| ) | |
| extractor = _TextFocusedHtmlExtractor() | |
| extractor.feed(raw) | |
| extractor.close() | |
| extracted_text = extractor.text() | |
| if not extracted_text: | |
| raise ValueError( | |
| "The HTML file did not contain readable text after removing " | |
| "scripts, styles, and web-only resources." | |
| ) | |
| normalized_path = ( | |
| session_dir | |
| / NORMALIZED_PDF_NAME | |
| ) | |
| styles = getSampleStyleSheet() | |
| body_style = ParagraphStyle( | |
| "WNPBody", | |
| parent=styles["BodyText"], | |
| fontName="Helvetica", | |
| fontSize=9.5, | |
| leading=13.5, | |
| alignment=TA_LEFT, | |
| spaceAfter=6, | |
| allowWidows=1, | |
| allowOrphans=1, | |
| ) | |
| heading_style = ParagraphStyle( | |
| "WNPHeading", | |
| parent=styles["Heading2"], | |
| fontName="Helvetica-Bold", | |
| fontSize=12, | |
| leading=15, | |
| spaceBefore=8, | |
| spaceAfter=6, | |
| keepWithNext=True, | |
| ) | |
| document = SimpleDocTemplate( | |
| str(normalized_path), | |
| pagesize=letter, | |
| rightMargin=0.62 * inch, | |
| leftMargin=0.62 * inch, | |
| topMargin=0.58 * inch, | |
| bottomMargin=0.58 * inch, | |
| title=source_path.stem, | |
| author="WNP-ARAG", | |
| ) | |
| story = [] | |
| blocks = [ | |
| block.strip() | |
| for block in re.split( | |
| r"\n\s*\n", | |
| extracted_text, | |
| ) | |
| if block.strip() | |
| ] | |
| for block in blocks: | |
| # Preserve explicit lines and escape XML-sensitive characters. | |
| safe_block = escape( | |
| block, | |
| quote=False, | |
| ).replace("\n", "<br/>") | |
| # Treat short title-like blocks as headings. | |
| word_count = len(block.split()) | |
| is_heading = ( | |
| word_count <= 18 | |
| and len(block) <= 140 | |
| and not block.endswith((".", ";", ",")) | |
| ) | |
| story.append( | |
| Paragraph( | |
| safe_block, | |
| heading_style | |
| if is_heading | |
| else body_style, | |
| ) | |
| ) | |
| if not is_heading: | |
| story.append( | |
| Spacer(1, 2) | |
| ) | |
| if not story: | |
| raise ValueError( | |
| "No printable text remained after HTML preprocessing." | |
| ) | |
| document.build(story) | |
| if not normalized_path.is_file(): | |
| raise RuntimeError( | |
| "The normalized PDF could not be created from the HTML text." | |
| ) | |
| return normalized_path | |
| def _convert_to_pdf( | |
| source_path: Path, | |
| session_dir: Path, | |
| ) -> Path: | |
| """Convert an editable source document to a normalized PDF.""" | |
| executable = _find_libreoffice() | |
| conversion_source = source_path | |
| if source_path.suffix.lower() in {".html", ".htm"}: | |
| conversion_source = _prepare_local_html( | |
| source_path, | |
| session_dir, | |
| ) | |
| conversion_dir = session_dir / "converted" | |
| conversion_dir.mkdir( | |
| parents=True, | |
| exist_ok=True, | |
| ) | |
| profile_dir = session_dir / "libreoffice_profile" | |
| profile_dir.mkdir( | |
| parents=True, | |
| exist_ok=True, | |
| ) | |
| profile_uri = profile_dir.resolve().as_uri() | |
| command = [ | |
| executable, | |
| f"-env:UserInstallation={profile_uri}", | |
| "--headless", | |
| "--convert-to", | |
| "pdf", | |
| "--outdir", | |
| str(conversion_dir), | |
| str(conversion_source), | |
| ] | |
| conversion_timeout = ( | |
| 180 | |
| if source_path.suffix.lower() in {".html", ".htm"} | |
| else 420 | |
| ) | |
| try: | |
| result = subprocess.run( | |
| command, | |
| check=False, | |
| capture_output=True, | |
| text=True, | |
| timeout=conversion_timeout, | |
| ) | |
| except subprocess.TimeoutExpired as exc: | |
| raise RuntimeError( | |
| "Document-to-PDF conversion exceeded the " | |
| f"{conversion_timeout}-second limit. " | |
| "The document could not be converted to a stable normalized PDF." | |
| ) from exc | |
| expected_path = ( | |
| conversion_dir | |
| / f"{conversion_source.stem}.pdf" | |
| ) | |
| candidates = sorted( | |
| conversion_dir.glob("*.pdf"), | |
| key=lambda path: path.stat().st_mtime, | |
| reverse=True, | |
| ) | |
| converted_path = ( | |
| expected_path | |
| if expected_path.is_file() | |
| else candidates[0] | |
| if candidates | |
| else None | |
| ) | |
| if result.returncode != 0 or converted_path is None: | |
| details = ( | |
| result.stderr | |
| or result.stdout | |
| or "" | |
| ).strip() | |
| raise RuntimeError( | |
| "LibreOffice could not create a normalized PDF." | |
| + ( | |
| f" Details: {details}" | |
| if details | |
| else "" | |
| ) | |
| ) | |
| normalized_path = ( | |
| session_dir | |
| / NORMALIZED_PDF_NAME | |
| ) | |
| shutil.copy2( | |
| converted_path, | |
| normalized_path, | |
| ) | |
| return normalized_path | |
| def _prepare_normalized_pdf( | |
| copied_source: Path, | |
| original_extension: str, | |
| session_dir: Path, | |
| ) -> tuple[Path, str, str]: | |
| """ | |
| Return: | |
| - normalized PDF path | |
| - conversion method | |
| - page-reference basis | |
| """ | |
| normalized_path = ( | |
| session_dir | |
| / NORMALIZED_PDF_NAME | |
| ) | |
| if original_extension == ".pdf": | |
| shutil.copy2( | |
| copied_source, | |
| normalized_path, | |
| ) | |
| return ( | |
| normalized_path, | |
| "None", | |
| "original_pdf", | |
| ) | |
| if original_extension in { | |
| ".html", | |
| ".htm", | |
| }: | |
| normalized_path = ( | |
| _render_html_text_to_pdf( | |
| copied_source, | |
| session_dir, | |
| ) | |
| ) | |
| return ( | |
| normalized_path, | |
| "Text-focused HTML rendering with ReportLab", | |
| "normalized_pdf", | |
| ) | |
| normalized_path = _convert_to_pdf( | |
| copied_source, | |
| session_dir, | |
| ) | |
| return ( | |
| normalized_path, | |
| "LibreOffice headless PDF conversion", | |
| "normalized_pdf", | |
| ) | |
| def _iterate_docling_items(document: Any): | |
| """Iterate native Docling items across compatible API versions.""" | |
| try: | |
| yield from document.iterate_items(with_groups=False) | |
| except TypeError: | |
| yield from document.iterate_items() | |
| def _item_page_numbers(item: Any) -> list[int]: | |
| """Return all unique provenance pages attached to one Docling item.""" | |
| pages = [] | |
| for provenance in list( | |
| getattr(item, "prov", None) | |
| or [] | |
| ): | |
| value = getattr(provenance, "page_no", None) | |
| if value is None: | |
| continue | |
| try: | |
| page = int(value) | |
| except (TypeError, ValueError): | |
| continue | |
| if page > 0 and page not in pages: | |
| pages.append(page) | |
| return pages | |
| def _item_to_markdown( | |
| item: Any, | |
| document: Any, | |
| ) -> str: | |
| """ | |
| Convert one native Docling item to readable text. | |
| Text items use their text directly. Tables and other structured | |
| items are exported to Markdown when supported. | |
| """ | |
| text = str( | |
| getattr(item, "text", "") | |
| or "" | |
| ).strip() | |
| if text: | |
| return text | |
| export_method = getattr( | |
| item, | |
| "export_to_markdown", | |
| None, | |
| ) | |
| if callable(export_method): | |
| for args in ( | |
| (document,), | |
| (), | |
| ): | |
| try: | |
| value = str(export_method(*args)).strip() | |
| if value: | |
| return value | |
| except Exception: | |
| continue | |
| caption_method = getattr( | |
| item, | |
| "caption_text", | |
| None, | |
| ) | |
| if callable(caption_method): | |
| try: | |
| value = str( | |
| caption_method(doc=document) | |
| ).strip() | |
| if value: | |
| return value | |
| except Exception: | |
| pass | |
| return "" | |
| def _extract_page_documents( | |
| normalized_pdf_path: Path, | |
| original_name: str, | |
| clean_title: str, | |
| original_extension: str, | |
| page_reference_basis: str, | |
| ) -> tuple[list[Document], dict[str, int]]: | |
| """ | |
| Extract the normalized PDF with native Docling and create one | |
| LangChain Document per page. | |
| Page provenance is attached before chunking, so every resulting | |
| chunk retains a reliable page number. | |
| """ | |
| logger.info( | |
| "Converting normalized PDF with native Docling: %s", | |
| normalized_pdf_path, | |
| ) | |
| converter = DocumentConverter() | |
| result = converter.convert(str(normalized_pdf_path)) | |
| docling_document = result.document | |
| page_parts: dict[int, list[str]] = defaultdict(list) | |
| items_with_page = 0 | |
| items_without_page = 0 | |
| for item, _level in _iterate_docling_items( | |
| docling_document | |
| ): | |
| text = _item_to_markdown( | |
| item, | |
| docling_document, | |
| ) | |
| if not text: | |
| continue | |
| pages = _item_page_numbers(item) | |
| if not pages: | |
| items_without_page += 1 | |
| continue | |
| items_with_page += 1 | |
| # An item may have provenance on more than one page. | |
| # Attach it to each referenced page so the evidence location | |
| # remains recoverable. | |
| for page in pages: | |
| page_parts[page].append(text) | |
| page_documents = [] | |
| for page in sorted(page_parts): | |
| parts = [ | |
| part.strip() | |
| for part in page_parts[page] | |
| if part and part.strip() | |
| ] | |
| page_text = "\n\n".join(parts).strip() | |
| if not page_text: | |
| continue | |
| page_label = ( | |
| f"Page {page}" | |
| if page_reference_basis == "original_pdf" | |
| else f"Page {page} (normalized PDF)" | |
| ) | |
| page_documents.append( | |
| Document( | |
| page_content=page_text, | |
| metadata={ | |
| "title": clean_title, | |
| "source": original_name, | |
| "page": page, | |
| "page_label": page_label, | |
| "document_type": "uploaded_document", | |
| "original_extension": original_extension, | |
| "processing_extension": ".pdf", | |
| "page_reference_basis": page_reference_basis, | |
| }, | |
| ) | |
| ) | |
| if not page_documents: | |
| raise ValueError( | |
| "Native Docling extracted no page-aware text from the " | |
| "normalized PDF." | |
| ) | |
| diagnostics = { | |
| "page_documents": len(page_documents), | |
| "items_with_page": items_with_page, | |
| "items_without_page": items_without_page, | |
| } | |
| logger.info( | |
| "Native Docling page extraction created %s page documents " | |
| "from %s page-aware items; %s text-bearing items lacked " | |
| "page provenance.", | |
| diagnostics["page_documents"], | |
| diagnostics["items_with_page"], | |
| diagnostics["items_without_page"], | |
| ) | |
| return page_documents, diagnostics | |
| def _build_download_package( | |
| session_dir: Path, | |
| vector_dir: Path, | |
| normalized_pdf_path: Path, | |
| manifest: dict, | |
| ) -> Path: | |
| manifest_path = session_dir / "manifest.json" | |
| manifest_path.write_text( | |
| json.dumps( | |
| manifest, | |
| indent=2, | |
| ensure_ascii=False, | |
| ), | |
| encoding="utf-8", | |
| ) | |
| readme_path = session_dir / "README.txt" | |
| readme_path.write_text( | |
| "\n".join( | |
| [ | |
| "WNP-ARAG uploaded-document vector database", | |
| "", | |
| "Contents:", | |
| "- index.faiss", | |
| "- index.pkl", | |
| "- normalized_document.pdf", | |
| "- manifest.json", | |
| "- README.txt", | |
| "", | |
| ( | |
| "Source-card page references correspond to " | |
| "normalized_document.pdf." | |
| ), | |
| ( | |
| "For an originally uploaded PDF, " | |
| "normalized_document.pdf is an exact copy of that PDF." | |
| ), | |
| ( | |
| "For DOC, DOCX, PPTX, HTML, and TXT, " | |
| "normalized_document.pdf was created with LibreOffice." | |
| ), | |
| ( | |
| "Text and page provenance were extracted with " | |
| "Docling's native DocumentConverter." | |
| ), | |
| ( | |
| "This database was created with embedding model " | |
| f"{manifest['embedding_model']}." | |
| ), | |
| ( | |
| "Use the same embedding model and a compatible " | |
| "LangChain FAISS loader when reusing the database." | |
| ), | |
| ( | |
| "Security notice: index.pkl is a Python pickle file. " | |
| "Load pickle files only from sources you trust." | |
| ), | |
| ] | |
| ), | |
| encoding="utf-8", | |
| ) | |
| zip_path = ( | |
| session_dir | |
| / "uploaded_document_vector_db.zip" | |
| ) | |
| with zipfile.ZipFile( | |
| zip_path, | |
| "w", | |
| compression=zipfile.ZIP_DEFLATED, | |
| ) as archive: | |
| archive.write( | |
| vector_dir / "index.faiss", | |
| "index.faiss", | |
| ) | |
| archive.write( | |
| vector_dir / "index.pkl", | |
| "index.pkl", | |
| ) | |
| archive.write( | |
| normalized_pdf_path, | |
| NORMALIZED_PDF_NAME, | |
| ) | |
| archive.write( | |
| manifest_path, | |
| "manifest.json", | |
| ) | |
| archive.write( | |
| readme_path, | |
| "README.txt", | |
| ) | |
| return zip_path | |
| def process_uploaded_document( | |
| uploaded_file: Any, | |
| requested_session_id: str = "", | |
| ): | |
| """Process one document and stream progress to Gradio clients.""" | |
| started_at = time.time() | |
| active_session_id = "" | |
| current_stage = "Preparing document" | |
| def progress(stage: str, message: str, detail: str = ""): | |
| nonlocal current_stage | |
| current_stage = stage | |
| elapsed_seconds = max(0, int(time.time() - started_at)) | |
| minutes, seconds = divmod(elapsed_seconds, 60) | |
| lines = [f"### {stage}", message] | |
| if detail: | |
| lines.extend(["", detail]) | |
| return ( | |
| active_session_id, | |
| "\n".join(lines), | |
| None, | |
| False, | |
| ) | |
| yield progress( | |
| "Preparing document", | |
| "Validating the uploaded file and creating a temporary session.", | |
| ) | |
| cleanup_expired_sessions() | |
| source_path = _resolve_uploaded_path(uploaded_file) | |
| original_name = source_path.name | |
| original_extension = source_path.suffix.lower() | |
| if original_extension not in UPLOAD_SUPPORTED_EXTENSIONS: | |
| supported = ", ".join(sorted(UPLOAD_SUPPORTED_EXTENSIONS)) | |
| raise ValueError( | |
| f"Unsupported file type '{original_extension}'. " | |
| f"Supported formats: {supported}." | |
| ) | |
| size_bytes = source_path.stat().st_size | |
| max_size_bytes = MAX_UPLOAD_SIZE_MB * 1024 * 1024 | |
| if size_bytes > max_size_bytes: | |
| size_mb = size_bytes / (1024 * 1024) | |
| raise ValueError( | |
| "File exceeds the upload limit. " | |
| f"The selected document is {size_mb:.1f} MB. " | |
| f"The maximum supported file size is {MAX_UPLOAD_SIZE_MB} MB. " | |
| "Please choose a smaller document or divide it into multiple files." | |
| ) | |
| UPLOAD_SESSION_ROOT.mkdir(parents=True, exist_ok=True) | |
| requested_session_id = str( | |
| requested_session_id or "" | |
| ).strip() | |
| if requested_session_id: | |
| try: | |
| session_id = uuid.UUID( | |
| requested_session_id | |
| ).hex | |
| except (ValueError, AttributeError, TypeError): | |
| session_id = uuid.uuid4().hex | |
| else: | |
| session_id = uuid.uuid4().hex | |
| active_session_id = session_id | |
| session_dir = UPLOAD_SESSION_ROOT / session_id | |
| vector_dir = session_dir / "vector_db" | |
| session_dir.mkdir(parents=True, exist_ok=False) | |
| copied_source = session_dir / original_name | |
| shutil.copy2(source_path, copied_source) | |
| try: | |
| if original_extension == ".pdf": | |
| yield progress( | |
| "Preparing PDF", | |
| "Copying the uploaded PDF into the temporary workspace.", | |
| f"**File:** {original_name}", | |
| ) | |
| elif original_extension in {".html", ".htm"}: | |
| yield progress( | |
| "Rendering HTML document", | |
| "Extracting readable policy text and creating a normalized PDF.", | |
| ( | |
| f"**File:** {original_name}\n\n" | |
| "Scripts, styles, images, and external web resources " | |
| "are excluded from the temporary document." | |
| ), | |
| ) | |
| else: | |
| yield progress( | |
| "Converting document", | |
| "Creating a normalized PDF with LibreOffice.", | |
| ( | |
| f"**File:** {original_name}\n\n" | |
| "Large office documents can take several minutes." | |
| ), | |
| ) | |
| normalized_pdf_path, conversion_method, page_reference_basis = _prepare_normalized_pdf( | |
| copied_source, original_extension, session_dir | |
| ) | |
| yield progress( | |
| "Checking document length", | |
| "Counting pages in the normalized PDF before text extraction.", | |
| f"**Maximum supported length:** {MAX_UPLOAD_PAGES} pages", | |
| ) | |
| page_count = _pdf_page_count( | |
| normalized_pdf_path | |
| ) | |
| if page_count > MAX_UPLOAD_PAGES: | |
| raise ValueError( | |
| "Document exceeds the page limit. " | |
| f"The normalized PDF contains {page_count} pages. " | |
| f"The maximum supported length is {MAX_UPLOAD_PAGES} pages. " | |
| "Please divide the document into smaller files and process " | |
| "them separately." | |
| ) | |
| yield progress( | |
| "Extracting page-aware content", | |
| "Reading document structure, text, tables, and page provenance with native Docling.", | |
| f"**Normalized PDF pages:** {page_count}", | |
| ) | |
| clean_title = copied_source.stem.replace("_", " ").strip() | |
| page_documents, extraction_diagnostics = _extract_page_documents( | |
| normalized_pdf_path=normalized_pdf_path, | |
| original_name=original_name, | |
| clean_title=clean_title, | |
| original_extension=original_extension, | |
| page_reference_basis=page_reference_basis, | |
| ) | |
| yield progress( | |
| "Creating searchable passages", | |
| "Splitting each page independently and removing duplicate passages.", | |
| f"**Pages with extracted text:** {extraction_diagnostics['page_documents']}", | |
| ) | |
| splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer( | |
| _get_tokenizer(), | |
| chunk_size=UPLOAD_CHUNK_SIZE, | |
| chunk_overlap=UPLOAD_CHUNK_OVERLAP, | |
| add_start_index=True, | |
| strip_whitespace=True, | |
| separators=["\n\n", "\n", ".", " ", ""], | |
| ) | |
| chunks = [] | |
| seen_texts = set() | |
| for page_document in page_documents: | |
| for chunk in splitter.split_documents([page_document]): | |
| text = str(chunk.page_content or "").strip() | |
| if not text or text in seen_texts: | |
| continue | |
| seen_texts.add(text) | |
| chunks.append(chunk) | |
| if not chunks: | |
| raise ValueError( | |
| "The document was processed, but no usable text chunks were created." | |
| ) | |
| yield progress( | |
| "Generating embeddings", | |
| "Encoding the page-aware passages with the WNP-ARAG embedding model.", | |
| f"**Passages to embed:** {len(chunks)}", | |
| ) | |
| embeddings = initialize_embeddings() | |
| vector_db = FAISS.from_documents(chunks, embeddings) | |
| vector_db.save_local(str(vector_dir)) | |
| yield progress( | |
| "Packaging temporary database", | |
| "Saving the FAISS index and preparing the downloadable package.", | |
| f"**Indexed passages:** {len(chunks)}\n\n**Included PDF:** {NORMALIZED_PDF_NAME}", | |
| ) | |
| now = time.time() | |
| expires_at = now + UPLOAD_SESSION_TTL_SECONDS | |
| manifest = { | |
| "original_filename": original_name, | |
| "original_extension": original_extension, | |
| "processing_extension": ".pdf", | |
| "conversion_method": conversion_method, | |
| "page_reference_basis": "Original uploaded PDF" if page_reference_basis == "original_pdf" else NORMALIZED_PDF_NAME, | |
| "normalized_pdf_included": True, | |
| "normalized_pdf_filename": NORMALIZED_PDF_NAME, | |
| "file_size_bytes": size_bytes, | |
| "page_count": page_count, | |
| "page_documents": extraction_diagnostics["page_documents"], | |
| "docling_items_with_page": extraction_diagnostics["items_with_page"], | |
| "docling_items_without_page": extraction_diagnostics["items_without_page"], | |
| "chunk_count": len(chunks), | |
| "embedding_model": EMBEDDING_MODEL_NAME, | |
| "chunk_size": UPLOAD_CHUNK_SIZE, | |
| "chunk_overlap": UPLOAD_CHUNK_OVERLAP, | |
| "text_extraction": "Native Docling DocumentConverter with item-level page provenance", | |
| "chunking_strategy": "Page-aware chunking; chunks never cross page boundaries", | |
| "vector_store": "LangChain FAISS", | |
| "created_at": datetime.now(timezone.utc).isoformat(), | |
| "session_lifetime_minutes": UPLOAD_SESSION_TTL_SECONDS // 60, | |
| } | |
| zip_path = _build_download_package( | |
| session_dir=session_dir, | |
| vector_dir=vector_dir, | |
| normalized_pdf_path=normalized_pdf_path, | |
| manifest=manifest, | |
| ) | |
| session = { | |
| "session_id": session_id, | |
| "vector_db": vector_db, | |
| "filename": original_name, | |
| "original_extension": original_extension, | |
| "processing_extension": ".pdf", | |
| "conversion_method": conversion_method, | |
| "page_reference_basis": page_reference_basis, | |
| "normalized_pdf_path": str(normalized_pdf_path), | |
| "page_count": page_count, | |
| "chunk_count": len(chunks), | |
| "session_dir": str(session_dir), | |
| "vector_dir": str(vector_dir), | |
| "zip_path": str(zip_path), | |
| "created_at": now, | |
| "expires_at": expires_at, | |
| } | |
| _write_session_metadata(session) | |
| with _sessions_lock: | |
| _sessions[session_id] = session | |
| total_seconds = max(0, int(time.time() - started_at)) | |
| minutes, seconds = divmod(total_seconds, 60) | |
| status_lines = [ | |
| "### Document ready", | |
| f"**File:** {original_name}", | |
| f"**Format:** {original_extension.lstrip('.').upper()}", | |
| f"**Pages:** {page_count}", | |
| ] | |
| if page_reference_basis == "normalized_pdf": | |
| status_lines.append("**Page references:** normalized PDF") | |
| if conversion_method != "None": | |
| status_lines.append(f"**Conversion:** {conversion_method}") | |
| status_lines.extend([ | |
| f"**Page-aware extraction:** {extraction_diagnostics['page_documents']} pages with text", | |
| f"**Indexed passages:** {len(chunks)}", | |
| f"**Download package:** includes {NORMALIZED_PDF_NAME}", | |
| f"**Session lifetime:** {UPLOAD_SESSION_TTL_SECONDS // 60} minutes from the latest activity", | |
| f"**Text embedding time elapsed:** {minutes}:{seconds:02d}", | |
| "", | |
| "Uploaded documents are processed temporarily for this session and are not added to the public regulatory databases.", | |
| ]) | |
| yield ( | |
| session_id, | |
| "\n".join(status_lines), | |
| str(zip_path), | |
| True, | |
| ) | |
| except Exception as exc: | |
| error_text = str(exc).strip() or exc.__class__.__name__ | |
| shutil.rmtree( | |
| session_dir, | |
| ignore_errors=True, | |
| ) | |
| failure_status = "\n".join( | |
| [ | |
| "### Document processing failed", | |
| f"**File:** {original_name}", | |
| f"**Stage:** {current_stage}", | |
| f"**Error:** {error_text}", | |
| "", | |
| "No user vector database was created.", | |
| ] | |
| ) | |
| yield ( | |
| active_session_id, | |
| failure_status, | |
| None, | |
| False, | |
| ) | |
| return | |
| # Backward-compatible alias for older clients. | |
| process_uploaded_pdf = process_uploaded_document | |
| def get_document_session( | |
| session_id: str, | |
| ) -> dict[str, Any]: | |
| cleanup_expired_sessions() | |
| key = str(session_id or "").strip() | |
| if not key: | |
| raise ValueError( | |
| "No uploaded-document session is active. " | |
| "Process a document first." | |
| ) | |
| with _sessions_lock: | |
| session = _sessions.get(key) | |
| if session is None: | |
| session = _restore_session_from_disk(key) | |
| if session is None: | |
| raise ValueError( | |
| "The uploaded-document session has expired or is " | |
| "unavailable. Please process the document again." | |
| ) | |
| session["expires_at"] = ( | |
| time.time() | |
| + UPLOAD_SESSION_TTL_SECONDS | |
| ) | |
| _write_session_metadata(session) | |
| with _sessions_lock: | |
| _sessions[key] = session | |
| return session | |
| def retrieve_uploaded_documents( | |
| query: str, | |
| top_k: int, | |
| session_id: str, | |
| ) -> list[dict]: | |
| session = get_document_session( | |
| session_id | |
| ) | |
| results = ( | |
| session["vector_db"] | |
| .similarity_search_with_score( | |
| query, | |
| k=top_k, | |
| ) | |
| ) | |
| chunks = [] | |
| for document, score in results: | |
| text = str( | |
| document.page_content | |
| or "" | |
| ).strip() | |
| if not text: | |
| continue | |
| chunks.append( | |
| { | |
| "text": text, | |
| "metadata": normalize_metadata( | |
| document.metadata | |
| ), | |
| "score": float(score), | |
| "jurisdiction": "uploaded", | |
| "collection_label": ( | |
| session["filename"] | |
| ), | |
| } | |
| ) | |
| return chunks | |
| def document_session_summary( | |
| session_id: str, | |
| ) -> dict[str, Any]: | |
| session = get_document_session( | |
| session_id | |
| ) | |
| return { | |
| "filename": session["filename"], | |
| "original_extension": ( | |
| session["original_extension"] | |
| ), | |
| "processing_extension": ( | |
| session["processing_extension"] | |
| ), | |
| "conversion_method": ( | |
| session["conversion_method"] | |
| ), | |
| "page_reference_basis": ( | |
| session["page_reference_basis"] | |
| ), | |
| "page_count": session["page_count"], | |
| "chunk_count": session["chunk_count"], | |
| "expires_at": session["expires_at"], | |
| "zip_path": session["zip_path"], | |
| "normalized_pdf_path": ( | |
| session["normalized_pdf_path"] | |
| ), | |
| } | |
| def validate_document_session( | |
| session_id: str, | |
| ) -> tuple[bool, str]: | |
| """Confirm that a temporary uploaded-document FAISS session still exists.""" | |
| try: | |
| summary = document_session_summary(session_id) | |
| return ( | |
| True, | |
| ( | |
| "User vector database available. " | |
| f"File: {summary['filename']}; " | |
| f"pages: {summary['page_count']}; " | |
| f"passages: {summary['chunk_count']}." | |
| ), | |
| ) | |
| except Exception as exc: | |
| technical_reason = ( | |
| str(exc).strip() | |
| or "The temporary document session is unavailable." | |
| ) | |
| return ( | |
| False, | |
| ( | |
| "The temporary user vector database is no longer available. " | |
| "This can happen when the Hugging Face Space sleeps, restarts, " | |
| "or when the 30-minute inactive session expires. " | |
| "Please rebuild the vector database from the selected document. " | |
| f"Technical detail: {technical_reason}" | |
| ), | |
| ) | |
| def clear_document_session( | |
| session_id: str, | |
| ) -> str: | |
| key = str(session_id or "").strip() | |
| if not key: | |
| return ( | |
| "No uploaded-document session " | |
| "was active." | |
| ) | |
| with _sessions_lock: | |
| session = _sessions.pop( | |
| key, | |
| None, | |
| ) | |
| session_dir = ( | |
| Path(session["session_dir"]) | |
| if session is not None | |
| else UPLOAD_SESSION_ROOT / key | |
| ) | |
| if not session_dir.exists(): | |
| return ( | |
| "The uploaded-document session " | |
| "was already unavailable." | |
| ) | |
| shutil.rmtree( | |
| session_dir, | |
| ignore_errors=True, | |
| ) | |
| return ( | |
| "Uploaded document and temporary " | |
| "vector database removed." | |
| ) | |
| def uploaded_session_count() -> int: | |
| cleanup_expired_sessions() | |
| UPLOAD_SESSION_ROOT.mkdir( | |
| parents=True, | |
| exist_ok=True, | |
| ) | |
| valid_ids = set() | |
| for session_dir in ( | |
| UPLOAD_SESSION_ROOT.iterdir() | |
| ): | |
| if not session_dir.is_dir(): | |
| continue | |
| payload = _read_session_metadata( | |
| session_dir | |
| ) | |
| if ( | |
| payload is not None | |
| and float(payload["expires_at"]) | |
| > time.time() | |
| ): | |
| valid_ids.add( | |
| str(payload["session_id"]) | |
| ) | |
| with _sessions_lock: | |
| valid_ids.update( | |
| _sessions.keys() | |
| ) | |
| return len(valid_ids) | |