latest

routers/database/files/split_map.py

@@ -0,0 +1,588 @@
+# api/routers/database/files/split_map.py
+"""
+Automatic split_map.json generator for uploaded documents.
+
+This module creates chapter/section boundaries for documents using existing artefacts
+(Tika JSON, Docling frontmatter OCR) and optional PDF outline extraction.
+
+Strategy (waterfall, stop at confidence ≥ 0.7):
+1. PDF Outline/Bookmarks (best): confidence ≈ 0.95
+2. Headings from Docling JSON: confidence ≈ 0.8
+3. TOC from Tika text: confidence ≈ 0.7-0.8
+4. Fixed windows: confidence ≈ 0.2
+
+Hard constraints:
+- For any fallback Docling "no-OCR" call: limit page_range to [1, min(30, page_count)]
+- Never process more than 30 pages in one Docling request
+- Use existing artefacts whenever possible
+"""
+
+import re
+import json
+import uuid
+import datetime
+import os
+import requests
+from typing import List, Dict, Any, Optional, Tuple
+from modules.database.supabase.utils.client import SupabaseServiceRoleClient
+from modules.database.supabase.utils.storage import StorageAdmin
+from modules.logger_tool import initialise_logger
+
+logger = initialise_logger(__name__, os.getenv("LOG_LEVEL"), os.getenv("LOG_PATH"), 'default', True)
+
+# ---------- Utilities
+
+def _now_iso():
+    """Return current UTC timestamp in ISO format."""
+    return datetime.datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
+
+def _load_artefact_json(storage: StorageAdmin, bucket: str, rel_path: str) -> Optional[Dict[str, Any]]:
+    """Load JSON artefact from storage."""
+    try:
+        raw = storage.download_file(bucket, rel_path)
+        return json.loads(raw.decode("utf-8"))
+    except Exception as e:
+        logger.debug(f"Failed to load artefact {rel_path}: {e}")
+        return None
+
+def _page_count_from_tika(tika_json: Dict[str, Any]) -> Optional[int]:
+    """Extract page count from Tika JSON metadata."""
+    for k in ("xmpTPg:NPages", "Page-Count", "pdf:PageCount", "pdf:pagecount"):
+        v = tika_json.get(k) or tika_json.get(k.lower())
+        try:
+            if v is not None:
+                return int(v)
+        except Exception:
+            pass
+    return None
+
+# ---------- A) Outline via PyMuPDF (optional but recommended)
+
+def _try_outline(pdf_bytes: bytes) -> Optional[List[Tuple[str, int]]]:
+    """
+    Extract PDF outline/bookmarks using PyMuPDF.
+    Returns [(title, start_page)] for level-1 bookmarks only.
+    """
+    try:
+        import fitz  # PyMuPDF
+        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+        toc = doc.get_toc(simple=True)  # list of [level, title, page]
+        doc.close()
+        
+        # Keep level-1 only, ensure valid pages
+        out = []
+        for level, title, page in toc:
+            if level == 1 and page >= 1:
+                clean_title = title.strip()
+                if clean_title and len(clean_title) > 1:
+                    out.append((clean_title, page))
+        
+        return out if len(out) >= 2 else None  # Need at least 2 chapters
+    except ImportError:
+        logger.debug("PyMuPDF not available, skipping outline extraction")
+        return None
+    except Exception as e:
+        logger.debug(f"Outline extraction failed: {e}")
+        return None
+
+# ---------- B) Headings from Docling JSON
+
+def _try_headings(docling_json: Dict[str, Any]) -> Optional[List[Tuple[str, int, int]]]:
+    """
+    Extract headings from Docling JSON.
+    Returns [(title, start_page, level)] — we only return starts; end pages are computed later.
+    """
+    if not docling_json:
+        return None
+    
+    # Handle different Docling JSON structures
+    blocks = (docling_json.get("blocks") or 
+              docling_json.get("elements") or 
+              docling_json.get("body", {}).get("blocks") or [])
+    
+    candidates: List[Tuple[str, int, int]] = []
+    
+    for b in blocks:
+        # Check if this is a heading block
+        role = (b.get("role") or b.get("type") or "").lower()
+        if not ("heading" in role or role in ("h1", "h2", "title", "section-header")):
+            continue
+            
+        # Extract text content
+        text = (b.get("text") or b.get("content") or "").strip()
+        if not text or len(text) < 3:
+            continue
+            
+        # Extract page number with robust handling of 0-based pageIndex
+        p = None
+        if b.get("pageIndex") is not None:
+            try:
+                p = int(b.get("pageIndex")) + 1
+            except Exception:
+                p = None
+        if p is None:
+            for key in ("page", "page_no", "page_number"):
+                if b.get(key) is not None:
+                    try:
+                        p = int(b.get(key))
+                    except Exception:
+                        p = None
+                    break
+        if p is None or p < 1:
+            continue
+        
+        # Determine heading level
+        level = 1  # default
+        if "1" in role or "h1" in role:
+            level = 1
+        elif "2" in role or "h2" in role:
+            level = 2
+        
+        # Chapter regex boosts to level 1
+        if re.match(r"^\s*(chapter|ch\.?|section|part)\s+\d+", text, re.I):
+            level = 1
+        
+        candidates.append((text, p, level))
+    
+    if not candidates:
+        return None
+    
+    # Prefer level 1; if none, promote level 2 to level 1
+    l1 = [(t, p, l) for (t, p, l) in candidates if l == 1]
+    if not l1:
+        l1 = [(t, p, 1) for (t, p, _) in candidates]
+    
+    # Sort by page and keep strictly increasing pages only
+    l1_sorted = []
+    seen = set()
+    for (t, p, l) in sorted(l1, key=lambda x: x[1]):
+        if p not in seen and p >= 1:
+            l1_sorted.append((t, p, l))
+            seen.add(p)
+    
+    return l1_sorted if len(l1_sorted) >= 2 else None
+
+def _try_headings_fallback(file_id: str, cabinet_id: str, bucket: str, 
+                          processing_bytes: bytes, processing_mime: str, 
+                          page_count: int) -> Optional[List[Tuple[str, int, int]]]:
+    """
+    Make a limited Docling no-OCR call (max 30 pages) to extract headings.
+    This is used only when existing artefacts don't have sufficient heading data.
+    """
+    try:
+        docling_url = os.getenv('DOCLING_URL') or os.getenv('NEOFS_DOCLING_URL')
+        if not docling_url:
+            logger.debug("No Docling URL configured for headings fallback")
+            return None
+        
+        # Strictly limit to first 30 pages
+        max_pages = min(30, page_count)
+        logger.info(f"Headings fallback: limited Docling call for file_id={file_id}, pages=1-{max_pages}")
+        
+        # Build Docling request (no-OCR, limited pages)
+        docling_api_key = os.getenv('DOCLING_API_KEY')
+        headers = {'Accept': 'application/json'}
+        if docling_api_key:
+            headers['X-Api-Key'] = docling_api_key
+        
+        form_data = [
+            ('target_type', 'inbody'),
+            ('to_formats', 'json'),
+            ('do_ocr', 'false'),
+            ('force_ocr', 'false'),
+            ('image_export_mode', 'embedded'),
+            ('pdf_backend', 'dlparse_v4'),
+            ('table_mode', 'fast'),
+            ('page_range', '1'),
+            ('page_range', str(max_pages))
+        ]
+        
+        files = [('files', ('file', processing_bytes, processing_mime))]
+        
+        # Make the request with timeout
+        timeout = int(os.getenv('DOCLING_HEADINGS_TIMEOUT', '1800'))  # 30 minutes default
+        resp = requests.post(
+            f"{docling_url.rstrip('/')}/v1/convert/file",
+            files=files,
+            data=form_data,
+            headers=headers,
+            timeout=timeout
+        )
+        resp.raise_for_status()
+        
+        docling_json = resp.json()
+        logger.debug(f"Headings fallback: received Docling response for file_id={file_id}")
+        
+        return _try_headings(docling_json)
+        
+    except Exception as e:
+        logger.error(f"Headings fallback failed for file_id={file_id}: {e}")
+        return None
+
+# ---------- C) TOC from Tika text (dot leaders & page num)
+
+TOC_LINE = re.compile(r"^\s*(.+?)\s?(\.{2,}|\s{3,})\s*(\d{1,4})\s*$")
+
+def _try_toc_text(tika_text: str) -> Optional[List[Tuple[str, int]]]:
+    """
+    Parse TOC from Tika text using dot leaders and page numbers.
+    Returns [(title, start_page)] if successful.
+    """
+    if not tika_text:
+        return None
+    
+    # Heuristic: only scan first ~1500 lines (roughly first 15 pages)
+    head = "\n".join(tika_text.splitlines()[:1500])
+    pairs = []
+    
+    for line in head.splitlines():
+        m = TOC_LINE.match(line)
+        if not m:
+            continue
+        
+        title = m.group(1).strip()
+        try:
+            page = int(m.group(3))
+        except Exception:
+            continue
+        
+        # Reject obvious junk
+        if len(title) < 3 or page < 1 or page > 9999:
+            continue
+        
+        # Skip common false positives
+        if any(skip in title.lower() for skip in ['copyright', 'isbn', 'published', 'printed']):
+            continue
+            
+        pairs.append((title, page))
+    
+    # Require at least 5 entries and monotonic pages
+    if len(pairs) >= 5:
+        pages = [p for _, p in pairs]
+        if pages == sorted(pages):
+            logger.debug(f"TOC extraction found {len(pairs)} entries")
+            return pairs
+    
+    return None
+
+# ---------- Build entries with ends, apply smoothing
+
+def _entries_from_starts(starts: List[Tuple[str, int, int]], page_count: int, source: str = "headings") -> List[Dict[str, Any]]:
+    """
+    Build entries from start points with computed end pages.
+    starts: [(title, page, level)]
+    """
+    entries = []
+    base_confidence = 0.8 if source == "headings" else 0.75
+    
+    for i, (title, start, level) in enumerate(starts):
+        end = (starts[i + 1][1] - 1) if i + 1 < len(starts) else page_count
+        entries.append({
+            "id": f"sec{i + 1:02d}",
+            "title": title,
+            "level": level,
+            "start_page": int(start),
+            "end_page": int(end),
+            "source": source,
+            "confidence": base_confidence
+        })
+    
+    # Merge tiny sections (< 3 pages) into previous
+    merged = []
+    for e in entries:
+        section_size = e["end_page"] - e["start_page"] + 1
+        if merged and section_size < 3:
+            # Merge into previous section
+            merged[-1]["end_page"] = e["end_page"]
+            merged[-1]["title"] += " / " + e["title"]
+            merged[-1]["confidence"] *= 0.95  # Slight confidence penalty for merging
+        else:
+            merged.append(e)
+    
+    return merged
+
+def _entries_from_pairs(pairs: List[Tuple[str, int]], page_count: int, source: str = "outline") -> List[Dict[str, Any]]:
+    """
+    Build entries from (title, start_page) pairs.
+    """
+    entries = []
+    base_confidence = 0.95 if source == "outline" else (0.8 if source == "toc" else 0.75)
+    
+    for i, (title, start) in enumerate(pairs):
+        end = (pairs[i + 1][1] - 1) if i + 1 < len(pairs) else page_count
+        entries.append({
+            "id": f"sec{i + 1:02d}",
+            "title": title,
+            "level": 1,
+            "start_page": int(start),
+            "end_page": int(end),
+            "source": source,
+            "confidence": base_confidence
+        })
+    
+    # Apply same merging logic for tiny sections
+    merged = []
+    for e in entries:
+        section_size = e["end_page"] - e["start_page"] + 1
+        if merged and section_size < 3:
+            merged[-1]["end_page"] = e["end_page"]
+            merged[-1]["title"] += " / " + e["title"]
+            merged[-1]["confidence"] *= 0.95
+        else:
+            merged.append(e)
+    
+    return merged
+
+# ---------- Post-processing normalization
+
+def _normalize_entries(entries: List[Dict[str, Any]], page_count: int) -> List[Dict[str, Any]]:
+    """Normalize entries to ensure:
+    - coverage from page 1
+    - 1 <= start_page <= end_page <= page_count
+    - strictly increasing, non-overlapping ranges
+    - fill initial gap with a synthetic front matter section if needed
+    """
+    if not entries:
+        return entries
+
+    # Sanitize and sort by start_page
+    safe: List[Dict[str, Any]] = []
+    for e in entries:
+        try:
+            s = int(e.get("start_page", 1))
+            t = int(e.get("end_page", s))
+        except Exception:
+            continue
+        s = max(1, min(s, page_count))
+        t = max(1, min(t, page_count))
+        if t < s:
+            t = s
+        ne = dict(e)
+        ne["start_page"], ne["end_page"] = s, t
+        safe.append(ne)
+    safe.sort(key=lambda x: (x["start_page"], x.get("level", 1)))
+
+    # De-overlap by adjusting starts; ensure monotonic ranges
+    normalized: List[Dict[str, Any]] = []
+    for e in safe:
+        if not normalized:
+            normalized.append(e)
+            continue
+        prev = normalized[-1]
+        if e["start_page"] <= prev["end_page"]:
+            e["start_page"] = prev["end_page"] + 1
+            if e["start_page"] > page_count:
+                continue
+            if e["end_page"] < e["start_page"]:
+                e["end_page"] = e["start_page"]
+        e["end_page"] = min(e["end_page"], page_count)
+        normalized.append(e)
+
+    # Insert synthetic front matter if first start > 1
+    if normalized and normalized[0]["start_page"] > 1:
+        front = {
+            "id": "sec00",
+            "title": "Front matter",
+            "level": 1,
+            "start_page": 1,
+            "end_page": normalized[0]["start_page"] - 1,
+            "source": "synthetic",
+            "confidence": 0.6,
+        }
+        normalized.insert(0, front)
+
+    # Ensure last section ends at page_count
+    if normalized and normalized[-1]["end_page"] < page_count:
+        normalized[-1]["end_page"] = page_count
+
+    # Renumber ids sequentially
+    out: List[Dict[str, Any]] = []
+    for idx, e in enumerate(normalized, start=1):
+        ne = dict(e)
+        ne["id"] = f"sec{idx:02d}"
+        out.append(ne)
+    return out
+
+# ---------- Main entry point
+
+def create_split_map_for_file(file_id: str) -> Dict[str, Any]:
+    """
+    Create split_map.json for a file using waterfall strategy:
+    1. PDF outline (best)
+    2. Docling headings (from existing or limited fallback)
+    3. Tika TOC parsing
+    4. Fixed windows (fallback)
+    """
+    logger.info(f"Creating split_map for file_id={file_id}")
+    
+    client = SupabaseServiceRoleClient()
+    storage = StorageAdmin()
+
+    # 1) Lookup file row & bucket
+    fr = client.supabase.table('files').select('id,bucket,cabinet_id,name,path,mime_type').eq('id', file_id).single().execute()
+    file_row = fr.data or {}
+    bucket = file_row.get('bucket')
+    cabinet_id = file_row.get('cabinet_id')
+
+    # 2) Find artefacts
+    arts = client.supabase.table('document_artefacts') \
+        .select('*').eq('file_id', file_id).order('created_at', desc=True).execute().data or []
+
+    def find_art(t): 
+        for a in arts:
+            if a.get('type') == t: 
+                return a
+        return None
+
+    a_pdf = find_art('document_pdf')  # if converted to PDF
+    a_tika = find_art('tika_json')
+    a_noocr = find_art('docling_noocr_json')
+    a_fm = find_art('docling_frontmatter_json')
+
+    # 3) Load JSON/text
+    tika_json = _load_artefact_json(storage, bucket, a_tika['rel_path']) if a_tika else None
+    docling_noocr = _load_artefact_json(storage, bucket, a_noocr['rel_path']) if a_noocr else None
+    docling_fm = _load_artefact_json(storage, bucket, a_fm['rel_path']) if a_fm else None
+    
+    # Get page count
+    page_count = _page_count_from_tika(tika_json or {}) or 100  # reasonable default
+
+    # Get PDF bytes for outline extraction
+    pdf_bytes = None
+    processing_bytes = None
+    processing_mime = None
+    
+    if a_pdf:
+        # Use converted PDF
+        pdf_bytes = storage.download_file(bucket, a_pdf['rel_path'])
+        processing_bytes = pdf_bytes
+        processing_mime = 'application/pdf'
+    else:
+        # Check if original file is PDF
+        if file_row.get('mime_type') == 'application/pdf':
+            pdf_bytes = storage.download_file(bucket, file_row['path'])
+            processing_bytes = pdf_bytes
+            processing_mime = 'application/pdf'
+
+    # 4) Try methods in waterfall order
+    method = "fixed"
+    confidence = 0.2
+    entries: List[Dict[str, Any]] = []
+
+    # A) PDF Outline/Bookmarks (best)
+    if pdf_bytes and not entries:
+        logger.debug(f"Trying outline extraction for file_id={file_id}")
+        pairs = _try_outline(pdf_bytes)
+        if pairs:
+            entries = _entries_from_pairs(pairs, page_count, source="outline")
+            method, confidence = "outline", 0.95
+            logger.info(f"Split map: outline method found {len(entries)} sections")
+
+    # B) Headings from existing Docling JSON
+    if not entries:
+        logger.debug(f"Trying headings from existing Docling JSON for file_id={file_id}")
+        # Try no-OCR first, then frontmatter
+        for docling_json, source_name in [(docling_noocr, "noocr"), (docling_fm, "frontmatter")]:
+            if docling_json:
+                starts = _try_headings(docling_json)
+                if starts:
+                    entries = _entries_from_starts(starts, page_count, source="headings")
+                    method, confidence = "headings", 0.8
+                    logger.info(f"Split map: headings method ({source_name}) found {len(entries)} sections")
+                    break
+
+    # B2) Headings fallback with limited Docling call (if we have processing bytes)
+    if not entries and processing_bytes and processing_mime:
+        logger.debug(f"Trying headings fallback with limited Docling call for file_id={file_id}")
+        starts = _try_headings_fallback(file_id, cabinet_id, bucket, processing_bytes, processing_mime, page_count)
+        if starts:
+            entries = _entries_from_starts(starts, page_count, source="headings")
+            method, confidence = "headings", 0.75  # Slightly lower confidence for fallback
+            logger.info(f"Split map: headings fallback found {len(entries)} sections")
+
+    # C) TOC from Tika text
+    if not entries and tika_json:
+        logger.debug(f"Trying TOC extraction from Tika text for file_id={file_id}")
+        # Try common Tika text keys
+        text = tika_json.get("X-TIKA:content") or tika_json.get("content") or ""
+        pairs = _try_toc_text(text)
+        if pairs:
+            entries = _entries_from_pairs(pairs, page_count, source="toc")
+            method, confidence = "toc", 0.75
+            logger.info(f"Split map: TOC method found {len(entries)} sections")
+
+    # D) Fixed windows (fallback)
+    if not entries:
+        logger.info(f"Using fixed window fallback for file_id={file_id}")
+        step = max(10, min(20, page_count // 10))  # Adaptive step size
+        pairs = []
+        for i in range(1, page_count + 1, step):
+            end_page = min(i + step - 1, page_count)
+            title = f"Pages {i}-{end_page}" if i != end_page else f"Page {i}"
+            pairs.append((title, i))
+        
+        entries = _entries_from_pairs(pairs, page_count, source="fixed")
+        method, confidence = "fixed", 0.2
+        logger.info(f"Split map: fixed method created {len(entries)} sections")
+
+    # 5) Normalize entries and build split_map.json
+    entries = _normalize_entries(entries, page_count)
+    split_map = {
+        "version": 1,
+        "file_id": file_id,
+        "source_pdf_artefact_id": a_pdf['id'] if a_pdf else None,
+        "sources": {
+            "docling_noocr_json": a_noocr['id'] if a_noocr else None,
+            "docling_frontmatter_json": a_fm['id'] if a_fm else None,
+            "tika_json": a_tika['id'] if a_tika else None
+        },
+        "method": method,
+        "confidence": confidence,
+        "page_count": page_count,
+        "entries": entries,
+        "created_at": _now_iso(),
+        "notes": f"auto-generated using {method} method; user can edit in Split Marker UI"
+    }
+
+    # 6) Store as artefact
+    artefact_id = str(uuid.uuid4())
+    rel_path = f"{cabinet_id}/{file_id}/{artefact_id}/split_map.json"
+    
+    storage.upload_file(
+        bucket, 
+        rel_path, 
+        json.dumps(split_map, ensure_ascii=False, indent=2).encode("utf-8"), 
+        "application/json", 
+        upsert=True
+    )
+    
+    # Enhanced metadata for UI display
+    enhanced_extra = {
+        "method": method, 
+        "confidence": confidence, 
+        "entries_count": len(entries),
+        "display_name": "Document Structure Map",
+        "bundle_label": "Split Map",
+        "section_title": "Document Structure Map", 
+        "page_count": page_count,
+        "bundle_type": "split_map_json",
+        "processing_mode": "document_analysis",
+        "pipeline": "structure_analysis",
+        "is_structure_map": True,
+        "ui_category": "document_analysis",
+        "ui_order": 2,
+        "description": f"Document section boundaries identified using {method} method with {confidence:.1%} confidence ({len(entries)} sections)",
+        "viewer_type": "json"
+    }
+    
+    client.supabase.table('document_artefacts').insert({
+        "id": artefact_id,
+        "file_id": file_id,
+        "type": "split_map_json",
+        "rel_path": rel_path,
+        "extra": enhanced_extra,
+        "status": "completed"
+    }).execute()
+
+    logger.info(f"Split map stored: file_id={file_id}, method={method}, confidence={confidence:.2f}, entries={len(entries)}")
+    return split_map