"""Utilities for ingesting and normalizing documents.""" from __future__ import annotations import io import os import re import tempfile from urllib.parse import urlparse from typing import Iterable import requests from docx import Document as DocxDocument from lxml import html from readability import Document as ReadabilityDocument from pypdf import PdfReader from config import MAX_UPLOAD_SIZE_MB, TIMEOUT_SECONDS SUPPORTED_EXTENSIONS = { ".pdf": "application/pdf", ".txt": "text/plain", ".md": "text/markdown", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", } _CONTROL_CHARS_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]") class IngestionError(RuntimeError): pass def _ensure_size_limit(file_size: int) -> None: max_bytes = MAX_UPLOAD_SIZE_MB * 1024 * 1024 if file_size > max_bytes: raise IngestionError( f"File exceeds maximum size of {MAX_UPLOAD_SIZE_MB} MB." ) def extract_text_from_pdf(data: bytes) -> str: reader = PdfReader(io.BytesIO(data)) text_chunks: list[str] = [] for page in reader.pages: chunk = page.extract_text() or "" text_chunks.append(chunk.strip()) return _normalize_text(_repair_known_collapsed_terms("\n".join(text_chunks))) def extract_text_from_docx(data: bytes) -> str: with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp: tmp.write(data) tmp_path = tmp.name try: document = DocxDocument(tmp_path) paragraphs = [p.text.strip() for p in document.paragraphs if p.text.strip()] return _normalize_text("\n".join(paragraphs)) finally: os.unlink(tmp_path) def extract_text_from_text(data: bytes) -> str: return _normalize_text(data.decode("utf-8", errors="ignore")) def ingest_file(filename: str, file_bytes: bytes) -> str: _ensure_size_limit(len(file_bytes)) ext = os.path.splitext(filename)[1].lower() if ext not in SUPPORTED_EXTENSIONS: raise IngestionError(f"Unsupported file type: {ext}") if ext == ".pdf": return extract_text_from_pdf(file_bytes) if ext == ".docx": return extract_text_from_docx(file_bytes) return extract_text_from_text(file_bytes) def fetch_url_content(url: str, timeout_seconds: int = TIMEOUT_SECONDS) -> str: response = requests.get(url, timeout=timeout_seconds) response.raise_for_status() content_type = (response.headers.get("content-type") or "").lower() path = urlparse(url).path.lower() if "application/pdf" in content_type or path.endswith(".pdf"): return extract_text_from_pdf(response.content) if ( "application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type or path.endswith(".docx") ): return extract_text_from_docx(response.content) if any(text_type in content_type for text_type in ("text/plain", "text/markdown")) or any( path.endswith(ext) for ext in (".txt", ".md") ): return extract_text_from_text(response.content) readable = ReadabilityDocument(_strip_xml_control_chars(response.text)) summary_html = _strip_xml_control_chars(readable.summary(html_partial=True)) parsed = html.fromstring(summary_html) title = readable.short_title() or readable.title() or "" text_content = parsed.text_content() return _normalize_text("\n".join([title, text_content])) def combine_sources(sources: Iterable[str]) -> str: return _normalize_text("\n\n".join(filter(None, sources))) def _normalize_text(value: str) -> str: value = _strip_xml_control_chars(value) clean_lines = [line.strip() for line in value.splitlines()] return "\n".join(line for line in clean_lines if line) def _strip_xml_control_chars(value: str) -> str: return _CONTROL_CHARS_RE.sub("", value) def _repair_known_collapsed_terms(value: str) -> str: replacements = { "NATIONALINFORMATIONTECHNOLOGYAUTHORITY": "National Information Technology Authority", "NationalInformationTechnologyAuthority": "National Information Technology Authority", } for collapsed, expanded in replacements.items(): value = value.replace(collapsed, expanded) return value