@piecemaker-legal/piecemaker 1.0.6 → 1.0.7

package/server/piecemaker/vendor/websocket-server/scripts/convert_and_scan_pipeline.py

@@ -3,7 +3,7 @@
 Convert & Scan Pipeline — orchestrates document conversion and PII scanning.
 
 This script chains smart_converter.py (conversion to Markdown) and
-presidio-gliner.py (PII scanning) for batch processing of multiple documents.
+scanner_worker.py (PII scanning) for batch processing of multiple documents.
 
 Usage:
     python3 convert_and_scan_pipeline.py <file1> [file2 ...] -o <output_dir> [options]
@@ -48,6 +48,10 @@ from pathlib import Path
 from typing import List, Optional, Tuple, Dict, Set
 
 
+GLINER_CHUNK_SIZE = 384
+GLINER_CHUNK_OVERLAP = 64
+
+
 def print_progress(phase: str, current: int, total: int) -> None:
     """Print standardized progress message.
 
@@ -93,7 +97,11 @@ def start_scanner_worker() -> Optional[subprocess.Popen]:
         return None
 
 
-def wait_for_worker_ready(proc: Optional[subprocess.Popen], timeout: int = 500) -> bool:
+def wait_for_worker_ready(
+    proc: Optional[subprocess.Popen],
+    progress_context: Dict,
+    timeout: int = 500,
+) -> bool:
     """Wait for the scanner worker to print READY on stdout.
 
     While waiting, streams stderr output (model loading progress) in real-time.
@@ -108,9 +116,9 @@ def wait_for_worker_ready(proc: Optional[subprocess.Popen], timeout: int = 500)
     if proc is None:
         return False
 
-    import re
     import select
     import threading
+    progress_context["event"] = threading.Event()
 
     _chunk_re = re.compile(r"^PROGRESS:CHUNKS:\d+:\d+:\d+")
 
@@ -126,9 +134,24 @@ def wait_for_worker_ready(proc: Optional[subprocess.Popen], timeout: int = 500)
             for line in proc.stderr:
                 stripped = line.strip()
                 if _chunk_re.match(stripped):
-                    print(stripped, flush=True)
+                    _, _, _, current, raw_total = stripped.split(":")
+                    global_total = progress_context.get("total", 0)
+                    if global_total:
+                        global_current = min(global_total, progress_context.get("offset", 0) + int(current))
+                        global_percent = int(global_current * 100 / global_total)
+                        print(
+                            f"PROGRESS:CHUNKS:{global_percent}:{global_current}:{global_total}",
+                            flush=True,
+                        )
+                    else:
+                        print(stripped, flush=True)
+                    if int(current) >= int(raw_total):
+                        progress_context["event"].set()
                 else:
-                    print(f"  [worker] {line}", end="", flush=True)
+                    # Keep diagnostics on stderr. The Node parent deliberately
+                    # ignores non-PROGRESS stdout because it may contain document
+                    # text, while it retains a bounded stderr tail on failure.
+                    print(f"  [worker] {line}", end="", file=sys.stderr, flush=True)
         except (ValueError, OSError):
             pass  # Pipe closed
 
@@ -141,11 +164,12 @@ def wait_for_worker_ready(proc: Optional[subprocess.Popen], timeout: int = 500)
             print(f"⚠️  Scanner worker exited prematurely (exit {proc.returncode})", file=sys.stderr)
             return False
 
-        # Non-blocking read from stdout
+        # TextIOWrapper does not reliably distinguish EAGAIN from EOF when its
+        # descriptor is toggled to non-blocking mode. select() keeps the stream
+        # blocking and reads only after a complete line is available.
         try:
-            os.set_blocking(proc.stdout.fileno(), False)
-            line = proc.stdout.readline()
-            os.set_blocking(proc.stdout.fileno(), True)
+            readable, _, _ = select.select([proc.stdout], [], [], 0.1)
+            line = proc.stdout.readline() if readable else ""
         except (IOError, OSError):
             line = ""
 
@@ -153,16 +177,23 @@ def wait_for_worker_ready(proc: Optional[subprocess.Popen], timeout: int = 500)
             print("✅ Scanner worker ready (models loaded)")
             return True
 
-        if not line:
-            time.sleep(0.1)
-
     # Timeout
     print("⚠️  Scanner worker timed out waiting for READY", file=sys.stderr)
     proc.kill()
+    try:
+        proc.wait(timeout=5)
+    except subprocess.TimeoutExpired:
+        pass
     return False
 
 
-def scan_file_via_worker(proc: subprocess.Popen, md_file: str, output_dir: str) -> bool:
+def scan_file_via_worker(
+    proc: subprocess.Popen,
+    md_file: str,
+    output_dir: str,
+    progress_context: Optional[Dict] = None,
+    chunk_offset: int = 0,
+) -> bool:
     """Send a scan command to the worker and wait for the result.
 
     Args:
@@ -173,6 +204,9 @@ def scan_file_via_worker(proc: subprocess.Popen, md_file: str, output_dir: str)
     Returns:
         True if scan succeeded, False otherwise.
     """
+    if progress_context is not None:
+        progress_context["offset"] = chunk_offset
+        progress_context["event"].clear()
     cmd = json.dumps({"cmd": "scan", "md_file": md_file, "output_dir": output_dir})
     try:
         proc.stdin.write(cmd + "\n")
@@ -198,6 +232,9 @@ def scan_file_via_worker(proc: subprocess.Popen, md_file: str, output_dir: str)
         print(f"⚠️  Invalid JSON from worker: {response_line.strip()}", file=sys.stderr)
         return False
 
+    if progress_context is not None and progress_context.get("active_chunks", 0):
+        progress_context["event"].wait(timeout=5)
+
     if response.get("status") == "ok":
         return True
     else:
@@ -207,15 +244,43 @@ def scan_file_via_worker(proc: subprocess.Popen, md_file: str, output_dir: str)
 
 def stop_scanner_worker(proc: Optional[subprocess.Popen]) -> None:
     """Gracefully stop the scanner worker."""
-    if proc is None or proc.poll() is not None:
+    if proc is None:
         return
 
     try:
-        proc.stdin.write(json.dumps({"cmd": "quit"}) + "\n")
-        proc.stdin.flush()
-        proc.wait(timeout=10)
+        if proc.poll() is None:
+            proc.stdin.write(json.dumps({"cmd": "quit"}) + "\n")
+            proc.stdin.flush()
+            proc.wait(timeout=10)
     except Exception:
-        proc.kill()
+        try:
+            proc.kill()
+            proc.wait(timeout=5)
+        except Exception:
+            pass
+    finally:
+        for stream in (proc.stdin, proc.stdout, proc.stderr):
+            try:
+                stream.close()
+            except (AttributeError, OSError):
+                pass
+
+
+class PipelineResources:
+    """Owns every heavyweight or sensitive transient pipeline resource."""
+
+    def __init__(self) -> None:
+        self.scanner_worker: Optional[subprocess.Popen] = None
+        self.scan_workspace: Optional[tempfile.TemporaryDirectory] = None
+
+    def close(self) -> None:
+        # The temporary maps contain raw PII and must be removed even when
+        # worker shutdown itself encounters a broken pipe or a dead process.
+        try:
+            stop_scanner_worker(self.scanner_worker)
+        finally:
+            if self.scan_workspace is not None:
+                self.scan_workspace.cleanup()
 
 
 def convert_file(
@@ -328,100 +393,6 @@ def convert_file(
         return False, None
 
 
-def scan_file(md_file: str, output_dir: str) -> bool:
-    """Run presidio-gliner.py subprocess to scan Markdown for PII with real-time output streaming.
-
-    Args:
-        md_file: Path to markdown file
-        output_dir: Directory for output JSON
-
-    Returns:
-        success: True if scan succeeded
-    """
-    script_dir = Path(__file__).parent
-    scanner_script = script_dir / "presidio-gliner" / "presidio-gliner.py"
-
-    if not scanner_script.exists():
-        print(f"❌ presidio-gliner.py not found at {scanner_script}", file=sys.stderr)
-        return False
-
-    # Build command
-    cmd = [sys.executable, str(scanner_script), md_file, "-o", output_dir]
-
-    try:
-        # Run scanner with real-time output streaming
-        process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-            bufsize=1,  # Line buffered
-        )
-
-        # Stream output in real-time
-        os.set_blocking(process.stdout.fileno(), False)
-        os.set_blocking(process.stderr.fileno(), False)
-
-        return_code = None
-        while return_code is None:
-            # Check if process has finished
-            return_code = process.poll()
-
-            # Read available output without blocking
-            try:
-                # Read stdout
-                while True:
-                    line = process.stdout.readline()
-                    if not line:
-                        break
-                    print(line, end="", flush=True)
-
-                # Read stderr
-                while True:
-                    line = process.stderr.readline()
-                    if not line:
-                        break
-                    print(line, end="", file=sys.stderr, flush=True)
-
-            except (IOError, OSError):
-                pass
-
-            # Small delay to prevent CPU spinning
-            if return_code is None:
-                time.sleep(0.01)
-
-        # Ensure we read any remaining output
-        try:
-            remaining_stdout, remaining_stderr = process.communicate(timeout=1)
-            if remaining_stdout:
-                print(remaining_stdout, end="", flush=True)
-            if remaining_stderr:
-                print(remaining_stderr, end="", file=sys.stderr, flush=True)
-        except subprocess.TimeoutExpired:
-            pass
-
-        if return_code != 0:
-            print(f"⚠️  Scan failed with exit code {return_code}", file=sys.stderr)
-            return False
-
-        # Verify JSON output was created
-        md_stem = Path(md_file).stem
-        json_path = Path(output_dir) / f"{md_stem}_sensitive_map.json"
-
-        if not json_path.exists():
-            print(f"⚠️  Expected JSON output not found: {json_path}", file=sys.stderr)
-            return False
-
-        return True
-
-    except subprocess.TimeoutExpired:
-        print(f"❌ Scan timeout (>10 minutes): {md_file}", file=sys.stderr)
-        return False
-    except Exception as e:
-        print(f"❌ Scan error: {e}", file=sys.stderr)
-        return False
-
-
 def load_existing_mapping(mapping_path: Path) -> Optional[Dict]:
     """Load the case/document mapping file if it exists.
 
@@ -568,6 +539,25 @@ def are_names_similar(name1: str, norm1: str, name2: str, norm2: str) -> bool:
     return False
 
 
+def count_gliner_chunks(md_file: str) -> int:
+    try:
+        word_count = len(re.findall(r"\S+", Path(md_file).read_text(encoding="utf-8")))
+    except (OSError, UnicodeError):
+        return 1
+    if not word_count:
+        return 1
+    step = max(1, GLINER_CHUNK_SIZE - GLINER_CHUNK_OVERLAP)
+    return (max(0, word_count - GLINER_CHUNK_SIZE) + step - 1) // step + 1
+
+
+def emit_chunk_progress(current: int, total: int) -> None:
+    if total <= 0:
+        return
+    bounded_current = min(total, max(0, current))
+    percent = int(bounded_current * 100 / total)
+    print(f"PROGRESS:CHUNKS:{percent}:{bounded_current}:{total}", flush=True)
+
+
 def consolidate_duplicate_entities(entities: List[Dict]) -> List[Dict]:
     """Consolidate duplicate entities into single entries with variants.
 
@@ -643,6 +633,84 @@ def consolidate_duplicate_entities(entities: List[Dict]) -> List[Dict]:
     return consolidated
 
 
+def _person_entry_texts(code: str, entry: Dict, reverse_mapping: Dict,
+                        mapping_texts: Optional[List[str]] = None) -> List[str]:
+    values = [entry.get("original", ""), *(entry.get("variants") or [])]
+    values.extend(reverse_mapping.get(code, []) or [])
+    values.extend(mapping_texts or [])
+    return list(dict.fromkeys(
+        value for value in values
+        if isinstance(value, str) and value.strip()
+    ))
+
+
+def _merge_person_entry(merged_mapping: Dict, merged_reverse: Dict,
+                        merged_extracted: Dict, code: str, texts: List[str],
+                        source_entry: Optional[Dict] = None) -> None:
+    entries = merged_extracted.setdefault("personnes_physiques", {})
+    entry = entries.setdefault(code, {
+        "original": texts[0] if texts else "",
+        "code": code,
+        "variants": [],
+    })
+    current_texts = _person_entry_texts(code, entry, merged_reverse)
+    combined = list(dict.fromkeys([*current_texts, *texts]))
+    if combined:
+        entry["original"] = max(combined, key=len)
+        entry["variants"] = combined
+        merged_reverse[code] = list(dict.fromkeys([
+            *(merged_reverse.get(code, []) or []), *combined
+        ]))
+        for text in combined:
+            merged_mapping[text] = code
+    if source_entry:
+        entry["score"] = max(entry.get("score", 0.0), source_entry.get("score", 0.0))
+        entry.setdefault("recognizer", source_entry.get("recognizer", "unknown"))
+
+
+def _prune_existing_person_duplicates(merged_mapping: Dict, merged_reverse: Dict,
+                                       merged_extracted: Dict) -> None:
+    entries = merged_extracted.setdefault("personnes_physiques", {})
+    mapping_texts_by_code: Dict[str, List[str]] = {}
+    for text, code in merged_mapping.items():
+        mapping_texts_by_code.setdefault(code, []).append(text)
+
+    canonical_by_name: Dict[str, str] = {}
+    for code in list(entries):
+        entry = entries.get(code, {})
+        texts = _person_entry_texts(
+            code,
+            entry,
+            merged_reverse,
+            mapping_texts_by_code.get(code),
+        )
+        names = [normalize_name(text) for text in texts]
+        names = list(dict.fromkeys(name for name in names if name))
+        survivor = next((canonical_by_name[name] for name in names if name in canonical_by_name), None)
+        if survivor is None:
+            for name in names:
+                canonical_by_name[name] = code
+            continue
+        if survivor == code:
+            continue
+
+        _merge_person_entry(
+            merged_mapping,
+            merged_reverse,
+            merged_extracted,
+            survivor,
+            texts,
+            entry,
+        )
+        for name in names:
+            canonical_by_name[name] = survivor
+        for text, mapped_code in list(merged_mapping.items()):
+            if mapped_code == code:
+                merged_mapping[text] = survivor
+        merged_reverse.pop(code, None)
+        entries.pop(code, None)
+
+
 def is_known_city_or_country(text: str) -> bool:
     """Check if text matches known city or country names.
 
@@ -1188,8 +1256,21 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
         "autres": {**existing_mapping.get('extracted_data', {}).get('autres', {})}
     }
 
-    # Track seen entities to avoid duplicates
+    _prune_existing_person_duplicates(merged_mapping, merged_reverse, merged_extracted)
+
     seen_entities_lower = {k.lower(): v for k, v in merged_mapping.items()}
+    existing_person_by_name: Dict[str, str] = {}
+    for code, entry in merged_extracted["personnes_physiques"].items():
+        for entity_text in _person_entry_texts(code, entry, merged_reverse):
+            normalized = normalize_name(entity_text)
+            if normalized:
+                existing_person_by_name.setdefault(normalized, code)
+    for entity_text, code in merged_mapping.items():
+        if "PERSONNE_PHYSIQUE_" not in str(code) and not str(code).startswith(("DIRIGEANT_", "AVOCAT_")):
+            continue
+        normalized = normalize_name(entity_text)
+        if normalized:
+            existing_person_by_name.setdefault(normalized, code)
 
     # Find highest existing code numbers per category
     code_counters = {
@@ -1204,7 +1285,7 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
     # URGOT SA), là où un test par sous-chaîne ne reconnaîtrait pas SA_3.
     societe_counters: Dict[str, int] = {}
 
-    for code in merged_reverse.keys():
+    for code in set(merged_reverse) | set(merged_mapping.values()):
         if "PERSONNE_PHYSIQUE_" in code or code.startswith("DIRIGEANT_"):
             num = int(code.split("_")[-1])
             code_counters["personnes_physiques"] = max(code_counters["personnes_physiques"], num + 1)
@@ -1226,15 +1307,6 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
     for text, code in new_mapping.get('mapping', {}).items():
         text_lower = text.lower()
 
-        # Skip if already exists (use existing code)
-        if text_lower in seen_entities_lower:
-            continue
-
-        # Skip entities the lawyer discarded from the mapping
-        if text_lower in ignored_lower:
-            continue
-
-        # Determine category
         category = None
         for cat_key in merged_extracted.keys():
             if code in new_mapping.get('extracted_data', {}).get(cat_key, {}):
@@ -1244,6 +1316,42 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
         if not category:
             category = "autres"
 
+        old_entry = new_mapping.get('extracted_data', {}).get(category, {}).get(code, {})
+        variants = old_entry.get("variants", [text])
+        candidate_texts = list(dict.fromkeys([text, *variants]))
+
+        if category == "personnes_physiques":
+            existing_code = next(
+                (
+                    existing_person_by_name[normalized]
+                    for candidate in candidate_texts
+                    for normalized in [normalize_name(candidate)]
+                    if normalized in existing_person_by_name
+                ),
+                None,
+            )
+            if existing_code:
+                accepted_texts = [candidate for candidate in candidate_texts if candidate.lower() not in ignored_lower]
+                _merge_person_entry(
+                    merged_mapping,
+                    merged_reverse,
+                    merged_extracted,
+                    existing_code,
+                    accepted_texts,
+                    old_entry,
+                )
+                for candidate in accepted_texts:
+                    seen_entities_lower[candidate.lower()] = existing_code
+                continue
+
+        # Skip if already exists (use existing code)
+        if text_lower in seen_entities_lower:
+            continue
+
+        # Skip entities the lawyer discarded from the mapping
+        if text_lower in ignored_lower:
+            continue
+
         # Generate new code with incremented counter
         if category == "personnes_physiques":
             new_code = f"PERSONNE_PHYSIQUE_{code_counters[category]:02d}"
@@ -1267,13 +1375,14 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
             new_code = f"{entity_type}_{code_counters[category]:02d}"
             code_counters[category] += 1
 
-        # Get variants from existing entry
-        old_entry = new_mapping.get('extracted_data', {}).get(category, {}).get(code, {})
-        variants = old_entry.get("variants", [text])
-
         # Add all variants to merged structures (no lowercase duplicates)
-        for variant in variants:
+        accepted_variants = [variant for variant in variants if variant.lower() not in ignored_lower]
+        for variant in accepted_variants:
             merged_mapping[variant] = new_code
+            seen_entities_lower[variant.lower()] = new_code
+            normalized = normalize_name(variant) if category == "personnes_physiques" else ""
+            if normalized:
+                existing_person_by_name.setdefault(normalized, new_code)
 
         seen_entities_lower[text_lower] = new_code
 
@@ -1285,7 +1394,7 @@ def merge_with_existing_mapping(new_mapping: Dict, existing_mapping: Optional[Di
         merged_extracted[category][new_code] = {
             "original": original,
             "code": new_code,
-            "variants": variants,
+            "variants": accepted_variants,
             "score": old_entry.get("score", 1.0),
             "recognizer": old_entry.get("recognizer", "unknown")
         }
@@ -1633,7 +1742,7 @@ def write_document_index(
     return written
 
 
-def main():
+def run_pipeline(resources: PipelineResources):
     """Main pipeline orchestration."""
     parser = argparse.ArgumentParser(
         description="Convert documents to Markdown and scan for sensitive data",
@@ -1697,6 +1806,7 @@ def main():
         help="Reuse existing Markdown and scans whose source fingerprint is unchanged",
     )
     args = parser.parse_args()
+    progress_context = {"offset": 0, "total": 0, "active_chunks": 0}
 
     # Validate input files
     input_files = []
@@ -1747,7 +1857,8 @@ def main():
 
     # Start scanner worker immediately so model loading overlaps with Phase 1.
     if scan_needed:
-        scanner_worker = start_scanner_worker()
+        resources.scanner_worker = start_scanner_worker()
+        scanner_worker = resources.scanner_worker
     else:
         scanner_worker = None
         print("✅ Every PII scan is already up to date — GLiNER not started")
@@ -1825,19 +1936,27 @@ def main():
     if skipped_scans:
         print(f"⏭️  {skipped_scans} file(s) already scanned, left untouched")
 
-    # Wait for the scanner worker to finish loading models. One retry (clean
-    # up the stalled/dead worker, relaunch, wait again) before falling back
-    # to a subprocess-per-file scan.
-    worker_ready = wait_for_worker_ready(scanner_worker) if pending_scans else False
+    chunk_counts = {md_file: count_gliner_chunks(md_file) for md_file in pending_scans}
+    total_chunks = sum(chunk_counts.values())
+    progress_context["total"] = total_chunks
+
+    # Wait for the scanner worker to finish loading models. One retry cleans up
+    # and relaunches a stalled/dead worker; a second failure aborts the job so
+    # the heavyweight model is never reloaded once per document.
+    worker_ready = wait_for_worker_ready(scanner_worker, progress_context) if pending_scans else False
 
     if pending_scans and not worker_ready:
         print("⚠️  Scanner worker not ready — cleaning up and retrying once", file=sys.stderr)
         stop_scanner_worker(scanner_worker)
-        scanner_worker = start_scanner_worker()
-        worker_ready = wait_for_worker_ready(scanner_worker)
+        resources.scanner_worker = start_scanner_worker()
+        scanner_worker = resources.scanner_worker
+        worker_ready = wait_for_worker_ready(scanner_worker, progress_context)
 
     if pending_scans and not worker_ready:
-        print("⚠️  Scanner worker not available after retry, falling back to subprocess-per-file", file=sys.stderr)
+        raise RuntimeError(
+            "Scanner GLiNER indisponible après deux tentatives ; "
+            "le traitement est arrêté pour éviter de recharger le modèle pour chaque fichier."
+        )
     print()
 
     scan_success_count = 0
@@ -1846,18 +1965,26 @@ def main():
     index_path = state_target.parent / "document-index.json"
     # Raw detections contain PII. They live only in a private OS temporary
     # directory and are deleted individually as each file is merged.
-    scan_workspace = tempfile.TemporaryDirectory(prefix="piecemaker-scans-")
+    resources.scan_workspace = tempfile.TemporaryDirectory(prefix="piecemaker-scans-")
+    scan_workspace = resources.scan_workspace
     scan_output_dir = scan_workspace.name
 
+    emit_chunk_progress(0, total_chunks)
+    completed_chunks = 0
     for i, md_file in enumerate(pending_scans, start=1):
-        print_progress("SCAN", i, len(pending_scans))
         print(f"🔍 [{i}/{len(pending_scans)}] Scanning: {Path(md_file).name}")
 
-        if worker_ready:
-            success = scan_file_via_worker(scanner_worker, md_file, scan_output_dir)
-        else:
-            # Fallback: subprocess per file (old behavior)
-            success = scan_file(md_file, scan_output_dir)
+        progress_context["active_chunks"] = chunk_counts[md_file]
+        success = scan_file_via_worker(
+            scanner_worker,
+            md_file,
+            scan_output_dir,
+            progress_context,
+            completed_chunks,
+        )
+        if success:
+            completed_chunks += chunk_counts[md_file]
+            emit_chunk_progress(completed_chunks, total_chunks)
 
         if not success:
             print(f"   ⚠️  Scan failed, markdown preserved")
@@ -1906,10 +2033,6 @@ def main():
         scan_success_count += 1
         print()
 
-    # Shut down worker
-    stop_scanner_worker(scanner_worker)
-    scan_workspace.cleanup()
-
     print(f"✅ Phase 2 complete: {scan_success_count}/{len(pending_scans)} files scanned")
     print()
 
@@ -1954,6 +2077,14 @@ def main():
         return 1
 
 
+def main():
+    resources = PipelineResources()
+    try:
+        return run_pipeline(resources)
+    finally:
+        resources.close()
+
+
 if __name__ == "__main__":
     try:
         sys.exit(main())