""" environments/trace_env/tools/doc_tool.py Document Extraction Pipeline for Trace — Docling-First, Local-Only. Architecture: ───────────────────────────────────────────────────────────────── Stage 1 │ PARSE │ Docling (IBM) — Universal Document Parser │ │ Converts PDF/DOCX/PPTX → Markdown │ │ Preserves layout, tables, reading order ▼ ▼ Stage 2 │ IMAGES │ Extract embedded images from document │ │ Forward each to image_tool (Ollama VLM) │ │ Moondream → qwen2-vl / llama3.2-vision ▼ ▼ Stage 3 │ UNIFIED │ Combine text + image descriptions │ OUTPUT │ Standard dict → WorldModel Fallback: If Docling is not installed, falls back to PyMuPDF (PDF only) or python-docx (DOCX only) for basic text extraction. Memory strategy: - Docling runs locally, no GPU required - Peak RAM: ~2-3GB for most documents (well within 16GB) - Images are analysed one-at-a-time through Ollama (sequential, not parallel) Install: pip install docling # Fallbacks (optional): pip install PyMuPDF python-docx """ from __future__ import annotations import re import io import logging import os import tempfile import uuid from pathlib import Path from typing import Optional logger = logging.getLogger(__name__) # ── Configuration (overridden via env_config.yaml → configure()) ────────────── _MAX_IMAGES_PER_DOC = 2 # limit VLM calls per document _ANALYSE_EMBEDDED_IMAGES = False # DISABLED: amounts come from text layer, not logos _MAX_TEXT_LENGTH = 50000 # truncate very long documents def configure(config: dict): """ Called once at startup with the document_extraction section of env_config.yaml. config = { "enabled": true, "analyse_embedded_images": true, "max_images_per_doc": 5, "max_text_length": 50000, } """ global _MAX_IMAGES_PER_DOC, _ANALYSE_EMBEDDED_IMAGES, _MAX_TEXT_LENGTH _ANALYSE_EMBEDDED_IMAGES = config.get("analyse_embedded_images", _ANALYSE_EMBEDDED_IMAGES) _MAX_IMAGES_PER_DOC = config.get("max_images_per_doc", _MAX_IMAGES_PER_DOC) _MAX_TEXT_LENGTH = config.get("max_text_length", _MAX_TEXT_LENGTH) logger.info( f"[DOC_TOOL] Configured — " f"analyse_images={_ANALYSE_EMBEDDED_IMAGES}, " f"max_images={_MAX_IMAGES_PER_DOC}" ) # ── Public API ──────────────────────────────────────────────────────────────── def extract_document( file_bytes: bytes, filename: str, mime_type: str = "", analyse_images: bool = True, ) -> dict: """ Extract text and images from a document file (PDF, DOCX, PPTX). Uses Docling as primary parser. Falls back to PyMuPDF (PDF) or python-docx (DOCX) if Docling is not installed. Args: file_bytes: Raw document bytes (e.g. from Gmail attachment). filename: Original filename (used for format detection + result ID). mime_type: MIME type hint (optional, auto-detected from extension). analyse_images: If True, embedded images are analysed via image_tool. Returns: Unified extraction dict: { "id": str, "filename": str, "mime_type": str, "parser_used": "docling" | "pymupdf" | "python-docx", "extracted_text": str, # full text or Markdown "summary": str, # first ~300 chars "page_count": int | None, "images_found": int, "image_analyses": list[dict], # from image_tool "entities": dict, # placeholder for downstream NER "error": str | None, } """ result = _empty_result(filename, mime_type) try: ext = Path(filename).suffix.lower() # ── Try PyMuPDF first for PDFs (ultra-fast text extraction) ───── pdf_ok = False if ext == ".pdf": pdf_ok = _try_pymupdf(file_bytes, result) # If PyMuPDF got text, we're done — skip heavy Docling # If PyMuPDF returned False (empty/image-based PDF), try RapidOCR first if not pdf_ok: pdf_ok = _try_rapidocr(file_bytes, result) if not pdf_ok: # ── Fallback to Docling or other parsers ────────────────── docling_ok = _try_docling(file_bytes, filename, ext, result) if not docling_ok and ext in (".docx", ".doc"): _try_python_docx(file_bytes, result) elif not docling_ok and ext != ".pdf": result["error"] = f"Unsupported format: {ext}" return result # Truncate extremely long text if len(result["extracted_text"]) > _MAX_TEXT_LENGTH: result["extracted_text"] = result["extracted_text"][:_MAX_TEXT_LENGTH] + "\n\n[... truncated]" # Generate summary from extracted text text = result["extracted_text"] if text: result["summary"] = text[:300].strip() # ── Analyse embedded images via image_tool ──────────────────── if analyse_images and _ANALYSE_EMBEDDED_IMAGES: _analyse_embedded_images(file_bytes, filename, ext, result) logger.info( f"[DOC_TOOL] Extraction complete for '{filename}' — " f"parser={result['parser_used']}, " f"text_len={len(result['extracted_text'])}, " f"images={result['images_found']}" ) except Exception as e: err = f"{type(e).__name__}: {e}" result["error"] = err logger.error(f"[DOC_TOOL] Extraction failed for '{filename}': {err}") return result # ── Docling (primary parser) ────────────────────────────────────────────────── def _try_docling(file_bytes: bytes, filename: str, ext: str, result: dict) -> bool: """ Try to parse document with Docling. Returns True if successful. """ try: from docling.document_converter import DocumentConverter except ImportError: logger.info("[DOC_TOOL] Docling not installed, trying fallback parsers.") return False try: # Docling needs a file path — write to temp file with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp: tmp.write(file_bytes) tmp_path = tmp.name try: converter = DocumentConverter() doc_result = converter.convert(tmp_path) # Export to Markdown — preserves layout, tables, reading order md_content = doc_result.document.export_to_markdown() result["extracted_text"] = md_content # Improved regex for Indian/Global amounts # Matches: ₹1,500.00, Rs. 1500, Rs 1500/-, INR 1,500 found_amounts = re.findall(r'(?:₹|Rs\.?|INR)\s*([\d,]+(?:\.\d{1,2})?)(?:\s*/-)?', md_content, re.IGNORECASE) if found_amounts: # Sort to find the highest value (likely the Total) sorted_amounts = sorted(found_amounts, key=lambda x: float(x.replace(',', '')), reverse=True) result["entities"]["amounts"] = [f"₹{a}" for a in sorted_amounts] result["parser_used"] = "docling" if result["entities"]["amounts"]: logger.info(f"[DOC_TOOL] Docling found {len(result['entities']['amounts'])} amounts. Top: {result['entities']['amounts'][0]}") # Count pages if available if hasattr(doc_result.document, "pages"): result["page_count"] = len(doc_result.document.pages) # Count and extract images if hasattr(doc_result.document, "images"): images = doc_result.document.images result["images_found"] = len(images) if images else 0 logger.info( f"[DOC_TOOL] Docling parsed '{filename}' — " f"{len(md_content)} chars Markdown, " f"{result['page_count'] or '?'} pages" ) return True finally: # Clean up temp file try: os.unlink(tmp_path) except OSError: pass except Exception as e: logger.warning(f"[DOC_TOOL] Docling failed for '{filename}': {e}") return False # ── Fallback: RapidOCR (image-based PDFs) ──────────────────────────────────── def _try_rapidocr(file_bytes: bytes, result: dict) -> bool: """ Use RapidOCR to OCR image-based PDFs page by page. Much lighter than Docling's layout predictor (~200MB vs 3GB). Already installed: pip install rapidocr-onnxruntime """ try: import fitz # PyMuPDF to render pages as images from rapidocr import RapidOCR except ImportError: return False try: ocr = RapidOCR() doc = fitz.open(stream=file_bytes, filetype="pdf") result["page_count"] = len(doc) result["parser_used"] = "rapidocr" pages_text = [] for page_num, page in enumerate(doc): # Render page to image at 150 DPI (good quality, manageable size) mat = fitz.Matrix(150 / 72, 150 / 72) pix = page.get_pixmap(matrix=mat) img_bytes = pix.tobytes("png") # Run OCR ocr_result, _ = ocr(img_bytes) if ocr_result: page_text = "\n".join([line[1] for line in ocr_result if line and len(line) > 1]) if page_text.strip(): pages_text.append(f"--- Page {page_num + 1} ---\n{page_text}") doc.close() result["extracted_text"] = "\n\n".join(pages_text) if not result["extracted_text"].strip(): logger.warning("[DOC_TOOL] RapidOCR: no text found in PDF images.") return False # Pre-extract amounts found_amounts = re.findall( r'(?:[\u20b9]|Rs\.?|INR)\s*([\d,]+(?:\.\d{1,2})?)(?:\s*/-)?', result["extracted_text"], re.IGNORECASE ) if found_amounts: sorted_amounts = sorted(found_amounts, key=lambda x: float(x.replace(',', '')), reverse=True) result["entities"]["amounts"] = [f"\u20b9{a}" for a in sorted_amounts] logger.info(f"[DOC_TOOL] RapidOCR found {len(sorted_amounts)} amounts. Top: \u20b9{sorted_amounts[0]}") logger.info(f"[DOC_TOOL] RapidOCR parsed — {len(result['extracted_text'])} chars, {result['page_count']} pages") return True except Exception as e: logger.warning(f"[DOC_TOOL] RapidOCR failed: {e}") return False # ── Fallback: PyMuPDF (PDF only) ───────────────────────────────────────────── def _try_pymupdf(file_bytes: bytes, result: dict) -> bool: """ Fallback PDF parser using PyMuPDF (fitz). Ultra-fast text extraction. """ try: import fitz # PyMuPDF except ImportError: result["error"] = ( "Neither Docling nor PyMuPDF installed. " "Install with: pip install docling OR pip install PyMuPDF" ) return False try: doc = fitz.open(stream=file_bytes, filetype="pdf") result["parser_used"] = "pymupdf" result["page_count"] = len(doc) # Extract text from all pages pages_text = [] image_count = 0 for page_num, page in enumerate(doc): text = page.get_text("text") if text.strip(): pages_text.append(f"--- Page {page_num + 1} ---\n{text}") image_count += len(page.get_images(full=True)) result["extracted_text"] = "\n\n".join(pages_text) result["images_found"] = image_count # If no text was extracted, PDF is likely image-based — signal failure # so the caller can try OCR (RapidOCR/Docling) instead. if not result["extracted_text"].strip(): logger.info(f"[DOC_TOOL] PyMuPDF: no text layer found ({image_count} images). Needs OCR.") doc.close() return False # Pre-extract amounts from text to aid the verifier text = result["extracted_text"] found_amounts = re.findall(r'(?:₹|Rs\.?|INR)\s*([\d,]+(?:\.\d{1,2})?)(?:\s*/-)?', text, re.IGNORECASE) if found_amounts: sorted_amounts = sorted(found_amounts, key=lambda x: float(x.replace(',', '')), reverse=True) result["entities"]["amounts"] = [f"₹{a}" for a in sorted_amounts] logger.info(f"[DOC_TOOL] PyMuPDF found {len(sorted_amounts)} amounts. Top: ₹{sorted_amounts[0]}") doc.close() logger.info( f"[DOC_TOOL] PyMuPDF parsed — " f"{len(result['extracted_text'])} chars, " f"{result['page_count']} pages, " f"{image_count} images" ) return True except Exception as e: result["error"] = f"PyMuPDF failed: {e}" logger.error(f"[DOC_TOOL] PyMuPDF error: {e}") return False # ── Fallback: python-docx (DOCX only) ──────────────────────────────────────── def _try_python_docx(file_bytes: bytes, result: dict) -> bool: """ Fallback DOCX parser using python-docx. """ try: from docx import Document except ImportError: result["error"] = ( "Neither Docling nor python-docx installed. " "Install with: pip install docling OR pip install python-docx" ) return False try: doc = Document(io.BytesIO(file_bytes)) result["parser_used"] = "python-docx" # Extract paragraphs paragraphs = [p.text for p in doc.paragraphs if p.text.strip()] result["extracted_text"] = "\n\n".join(paragraphs) # Count inline images image_count = 0 for rel in doc.part.rels.values(): if "image" in rel.reltype: image_count += 1 result["images_found"] = image_count logger.info( f"[DOC_TOOL] python-docx parsed — " f"{len(paragraphs)} paragraphs, " f"{image_count} images" ) return True except Exception as e: result["error"] = f"python-docx failed: {e}" logger.error(f"[DOC_TOOL] python-docx error: {e}") return False # ── Image analysis for embedded images ──────────────────────────────────────── def _analyse_embedded_images( file_bytes: bytes, filename: str, ext: str, result: dict, ): """ Extract embedded images from the document and analyse them with the existing image_tool (Ollama VLM pipeline). """ images_to_analyse = [] try: if ext == ".pdf": images_to_analyse = _extract_pdf_images(file_bytes) elif ext in (".docx", ".doc"): images_to_analyse = _extract_docx_images(file_bytes) except Exception as e: logger.warning(f"[DOC_TOOL] Could not extract images from '{filename}': {e}") return if not images_to_analyse: return # Limit number of images to analyse images_to_analyse = images_to_analyse[:_MAX_IMAGES_PER_DOC] logger.info( f"[DOC_TOOL] Analysing {len(images_to_analyse)} embedded images " f"from '{filename}'..." ) try: from .image_tool import analyse_image_from_bytes except ImportError: logger.warning("[DOC_TOOL] image_tool not available for embedded image analysis.") return for i, (img_bytes, img_name) in enumerate(images_to_analyse): try: analysis = analyse_image_from_bytes( image_bytes=img_bytes, mime_type="image/png", filename=f"{filename}:{img_name}", ) result["image_analyses"].append(analysis) logger.info( f"[DOC_TOOL] Analysed embedded image {i+1}/{len(images_to_analyse)}: " f"{img_name} — {len(analysis.get('extracted_text', ''))} chars" ) except Exception as e: logger.warning( f"[DOC_TOOL] Image analysis failed for embedded image " f"'{img_name}' in '{filename}': {e}" ) def _extract_pdf_images(file_bytes: bytes) -> list[tuple[bytes, str]]: """Extract embedded images from a PDF as (bytes, name) tuples.""" images = [] try: import fitz # PyMuPDF doc = fitz.open(stream=file_bytes, filetype="pdf") for page_num, page in enumerate(doc): for img_idx, img_info in enumerate(page.get_images(full=True)): xref = img_info[0] try: base_image = doc.extract_image(xref) if base_image and base_image.get("image"): img_name = f"page{page_num+1}_img{img_idx+1}.{base_image.get('ext', 'png')}" images.append((base_image["image"], img_name)) except Exception: continue doc.close() except ImportError: logger.info("[DOC_TOOL] PyMuPDF not available for PDF image extraction.") return images def _extract_docx_images(file_bytes: bytes) -> list[tuple[bytes, str]]: """Extract embedded images from a DOCX as (bytes, name) tuples.""" images = [] try: from docx import Document doc = Document(io.BytesIO(file_bytes)) for rel_id, rel in doc.part.rels.items(): if "image" in rel.reltype: try: image_part = rel.target_part img_bytes = image_part.blob img_name = os.path.basename(image_part.partname) images.append((img_bytes, img_name)) except Exception: continue except ImportError: logger.info("[DOC_TOOL] python-docx not available for DOCX image extraction.") return images # ── Helpers ─────────────────────────────────────────────────────────────────── def _empty_result(filename: str, mime_type: str) -> dict: """Return a blank result skeleton.""" return { "id": f"doc_{uuid.uuid4().hex[:12]}", "filename": filename, "mime_type": mime_type, "parser_used": "none", "extracted_text": "", "summary": "", "page_count": None, "images_found": 0, "image_analyses": [], "entities": { "amounts": [], "dates": [], "vendors": [], "items": [], "other": [], }, "error": None, }