""" Module 1 — Document Ingestion POST /upload-document Accepts PDF, .md, or .txt uploads. PDF handling: Docling first → PyMuPDF fallback if Docling unavailable. Text/Markdown: passed through as-is. """ import io import logging import tempfile import os from pathlib import Path from fastapi import APIRouter, File, HTTPException, UploadFile from models.schemas import IngestionResponse logger = logging.getLogger(__name__) router = APIRouter() # --------------------------------------------------------------------------- # Parser availability detection (done once at import time, not per-request) # --------------------------------------------------------------------------- _DOCLING_AVAILABLE = False _PYMUPDF_AVAILABLE = False try: from docling.document_converter import DocumentConverter # type: ignore _DOCLING_AVAILABLE = True logger.info("✅ Docling is available — will use for PDF parsing") except Exception as e: logger.warning(f"⚠️ Docling not available ({e}). Checking PyMuPDF fallback...") try: import fitz # PyMuPDF # type: ignore _PYMUPDF_AVAILABLE = True logger.info("✅ PyMuPDF (fitz) is available — will use as PDF fallback") except Exception as e: logger.error( f"❌ Neither Docling nor PyMuPDF is available ({e}). " "PDF uploads will fail. Install at least one parser." ) # --------------------------------------------------------------------------- # Helper: PDF → Markdown using Docling # --------------------------------------------------------------------------- def _parse_pdf_docling(file_bytes: bytes, filename: str) -> str: """ Write bytes to a temp file, run Docling's DocumentConverter on it, and return clean Markdown preserving section hierarchy. """ from docling.document_converter import DocumentConverter # type: ignore with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: tmp.write(file_bytes) tmp_path = tmp.name try: converter = DocumentConverter() result = converter.convert(tmp_path) # Docling's result object exposes .document.export_to_markdown() markdown = result.document.export_to_markdown() return markdown finally: try: os.unlink(tmp_path) except OSError: pass # --------------------------------------------------------------------------- # Helper: PDF → plain text using PyMuPDF # --------------------------------------------------------------------------- def _parse_pdf_pymupdf(file_bytes: bytes) -> str: """ Extract text from every page using PyMuPDF. Attempts to preserve heading-like structure by detecting font size jumps, but primarily provides a reliable text fallback. """ import fitz # type: ignore doc = fitz.open(stream=file_bytes, filetype="pdf") pages: list[str] = [] for page in doc: text = page.get_text("text") pages.append(text) doc.close() return "\n\n---\n\n".join(pages) # --------------------------------------------------------------------------- # Endpoint # --------------------------------------------------------------------------- @router.post("/upload-document", response_model=IngestionResponse) async def upload_document(file: UploadFile = File(...)) -> IngestionResponse: """ Accept a policy document upload and return clean Markdown. Supported formats: - .pdf → Docling (preferred) or PyMuPDF (fallback) - .md → passthrough - .txt → passthrough """ filename = file.filename or "unknown" extension = Path(filename).suffix.lower() content = await file.read() # ---- Plain text / Markdown passthrough ----------------------------------- if extension in (".md", ".txt"): try: text = content.decode("utf-8") except UnicodeDecodeError: text = content.decode("latin-1") return IngestionResponse( markdown=text, filename=filename, parser_used="passthrough", ) # ---- PDF ----------------------------------------------------------------- if extension == ".pdf": if _DOCLING_AVAILABLE: try: markdown = _parse_pdf_docling(content, filename) return IngestionResponse( markdown=markdown, filename=filename, parser_used="docling", ) except Exception as docling_err: logger.warning( f"Docling failed ({docling_err}), trying PyMuPDF fallback..." ) if _PYMUPDF_AVAILABLE: try: markdown = _parse_pdf_pymupdf(content) return IngestionResponse( markdown=markdown, filename=filename, parser_used="pymupdf", ) except Exception as pymupdf_err: raise HTTPException( status_code=500, detail=f"Both PDF parsers failed. PyMuPDF error: {pymupdf_err}", ) raise HTTPException( status_code=500, detail=( "No PDF parser is available. " "Install 'docling' or 'PyMuPDF' and restart the server." ), ) # ---- Unsupported format -------------------------------------------------- raise HTTPException( status_code=400, detail=f"Unsupported file type '{extension}'. Please upload PDF, .md, or .txt.", )