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", "
")
paragraphs.append(
f"{escaped_block}
"
)
normalized_html = (
""
""
''
""
""
+ "".join(paragraphs)
+ ""
)
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", "
")
# 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)