""" Document extraction service – extracts text (as markdown) and images from uploaded documents (PDF, DOCX, PPTX). - PDF: PyMuPDF + PyMuPDF4LLM (markdown with structure preserved) - DOCX: python-docx (paragraphs + embedded images) - PPTX: python-pptx (slide text + embedded images) """ import os import hashlib import tempfile import fitz # PyMuPDF import pymupdf4llm from docx import Document as DocxDocument from pptx import Presentation from fastapi import UploadFile from sqlalchemy.orm import Session from app.config import settings from app.models.project import Project, ProjectStatus from app.models.asset import Asset, AssetType from app.services import r2_storage # Minimum image bytes to keep (skip tiny icons / decorations) _MIN_IMAGE_BYTES = 5_000 # 5 KB # Recognised file extensions -> handler key _EXT_MAP = { ".pdf": "pdf", ".docx": "docx", ".pptx": "pptx", } def extract_from_documents( project: Project, files: list[UploadFile], db: Session, ) -> Project: """ Extract text and images from uploaded documents. - Concatenates all extracted text into ``project.blog_content`` - Saves extracted images locally and uploads to R2 - Sets ``project.status = SCRAPED`` """ all_markdown: list[str] = [] image_dir = os.path.join(settings.MEDIA_DIR, f"projects/{project.id}/images") os.makedirs(image_dir, exist_ok=True) image_count = 0 for upload_file in files: filename = upload_file.filename or "document" ext = os.path.splitext(filename)[1].lower() handler = _EXT_MAP.get(ext, "pdf") # default to PDF # Save upload to a temp file with tempfile.NamedTemporaryFile(delete=False, suffix=ext or ".pdf") as tmp: content = upload_file.file.read() tmp.write(content) tmp_path = tmp.name try: if handler == "pdf": md, imgs = _extract_pdf(tmp_path, image_dir) elif handler == "docx": md, imgs = _extract_docx(tmp_path, image_dir) elif handler == "pptx": md, imgs = _extract_pptx(tmp_path, image_dir) else: md, imgs = "", [] if md and md.strip(): all_markdown.append(md) # Create Asset records for extracted images for img_path, img_filename in imgs: r2_key = None r2_url = None if r2_storage.is_r2_configured(): try: r2_url = r2_storage.upload_project_image( project.user_id, project.id, img_path, img_filename ) r2_key = r2_storage.image_key( project.user_id, project.id, img_filename ) except Exception as e: print(f"[DOC_EXTRACTOR] R2 upload failed for {img_filename}: {e}") asset = Asset( project_id=project.id, asset_type=AssetType.IMAGE, original_url=None, local_path=img_path, filename=img_filename, r2_key=r2_key, r2_url=r2_url, ) db.add(asset) image_count += 1 finally: try: os.unlink(tmp_path) except OSError: pass # ── Persist results ─────────────────────────────────────── project.blog_content = "\n\n---\n\n".join(all_markdown) if all_markdown else "" project.status = ProjectStatus.SCRAPED db.commit() db.refresh(project) print( f"[DOC_EXTRACTOR] Project {project.id}: extracted " f"{len(all_markdown)} document(s), {image_count} images, " f"{len(project.blog_content)} chars of markdown" ) return project # ─── PDF extraction ────────────────────────────────────────── def _extract_pdf( file_path: str, image_dir: str ) -> tuple[str, list[tuple[str, str]]]: """Return (markdown_text, [(local_path, filename), ...]).""" images: list[tuple[str, str]] = [] # Markdown text via pymupdf4llm md_text = pymupdf4llm.to_markdown(file_path) # Images via PyMuPDF doc = fitz.open(file_path) for page_num in range(len(doc)): page = doc[page_num] image_list = page.get_images(full=True) for img_info in image_list: xref = img_info[0] try: base_image = doc.extract_image(xref) except Exception: continue if not base_image or not base_image.get("image"): continue image_bytes = base_image["image"] if len(image_bytes) < _MIN_IMAGE_BYTES: continue ext = base_image.get("ext", "png") if ext not in ("png", "jpg", "jpeg", "webp"): ext = "png" img_hash = hashlib.md5(image_bytes).hexdigest()[:10] filename = f"pdf_p{page_num + 1}_{img_hash}.{ext}" local_path = os.path.join(image_dir, filename) if os.path.exists(local_path): continue with open(local_path, "wb") as f: f.write(image_bytes) images.append((local_path, filename)) doc.close() return md_text or "", images # ─── DOCX extraction ───────────────────────────────────────── def _extract_docx( file_path: str, image_dir: str ) -> tuple[str, list[tuple[str, str]]]: """Return (markdown_text, [(local_path, filename), ...]).""" images: list[tuple[str, str]] = [] lines: list[str] = [] doc = DocxDocument(file_path) # Extract text — convert paragraphs to simple markdown for para in doc.paragraphs: text = para.text.strip() if not text: lines.append("") continue style_name = (para.style.name or "").lower() if "heading 1" in style_name: lines.append(f"# {text}") elif "heading 2" in style_name: lines.append(f"## {text}") elif "heading 3" in style_name: lines.append(f"### {text}") elif "heading" in style_name: lines.append(f"#### {text}") elif "list" in style_name: lines.append(f"- {text}") else: lines.append(text) md_text = "\n\n".join(lines) # Extract embedded images from the docx relationships for rel in doc.part.rels.values(): if "image" in rel.reltype: try: image_bytes = rel.target_part.blob if len(image_bytes) < _MIN_IMAGE_BYTES: continue content_type = rel.target_part.content_type or "" ext = _mime_to_ext(content_type) img_hash = hashlib.md5(image_bytes).hexdigest()[:10] filename = f"docx_{img_hash}.{ext}" local_path = os.path.join(image_dir, filename) if os.path.exists(local_path): continue with open(local_path, "wb") as f: f.write(image_bytes) images.append((local_path, filename)) except Exception as e: print(f"[DOC_EXTRACTOR] DOCX image extraction error: {e}") continue return md_text, images # ─── PPTX extraction ───────────────────────────────────────── def _extract_pptx( file_path: str, image_dir: str ) -> tuple[str, list[tuple[str, str]]]: """Return (markdown_text, [(local_path, filename), ...]).""" images: list[tuple[str, str]] = [] slides_text: list[str] = [] prs = Presentation(file_path) for slide_num, slide in enumerate(prs.slides, 1): slide_lines: list[str] = [] for shape in slide.shapes: # Extract text from text frames if shape.has_text_frame: for para in shape.text_frame.paragraphs: text = para.text.strip() if text: slide_lines.append(text) # Extract text from tables if shape.has_table: for row in shape.table.rows: row_text = " | ".join( cell.text.strip() for cell in row.cells ) if row_text.strip(" |"): slide_lines.append(row_text) # Extract images if shape.shape_type == 13: # MSO_SHAPE_TYPE.PICTURE try: image = shape.image image_bytes = image.blob if len(image_bytes) < _MIN_IMAGE_BYTES: continue ext = _mime_to_ext(image.content_type or "") img_hash = hashlib.md5(image_bytes).hexdigest()[:10] filename = f"pptx_s{slide_num}_{img_hash}.{ext}" local_path = os.path.join(image_dir, filename) if os.path.exists(local_path): continue with open(local_path, "wb") as f: f.write(image_bytes) images.append((local_path, filename)) except Exception as e: print(f"[DOC_EXTRACTOR] PPTX image extraction error: {e}") continue if slide_lines: header = f"## Slide {slide_num}" slides_text.append(header + "\n\n" + "\n\n".join(slide_lines)) md_text = "\n\n---\n\n".join(slides_text) return md_text, images # ─── Helpers ────────────────────────────────────────────────── def _mime_to_ext(content_type: str) -> str: """Map MIME type to file extension.""" ct = content_type.lower() mapping = { "image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg", "image/gif": "gif", "image/webp": "webp", "image/tiff": "tiff", "image/bmp": "bmp", "image/x-emf": "emf", "image/x-wmf": "wmf", } return mapping.get(ct, "png")