Files
legal-ai/mcp-server/src/legal_mcp/services/extractor.py
Chaim 9ae7304d44
All checks were successful
G12 Leak-Guard / leak-guard (pull_request) Successful in 4s
Lint — undefined names / undefined-names (pull_request) Successful in 10s
feat(ocr): replace Google Vision with Mistral OCR as PDF fallback
Switches the scanned-PDF fallback from Google Cloud Vision to
Mistral OCR (mistral-ocr-latest) for better Hebrew accuracy and
robustness against broken embedded OCR layers (e.g. case 1044-03-26
which returned English garbage through Vision).

Routing strategy (document-level, not per-page):
- PyMuPDF extracts all pages; pages that pass _text_quality_ok()
  use PyMuPDF output directly (free, ~50ms).
- If ANY page fails quality → Mistral OCR called once for the whole
  PDF, returning per-page Markdown for all pages (consistent format,
  no plain-text/Markdown mix within a document).

Markdown output preserved: Mistral returns ## headers and |tables|;
chunker updated to recognise ATX Markdown headers (##/###) as section
boundaries in _split_into_sections().

Config: GOOGLE_CLOUD_VISION_API_KEY → MISTRAL_API_KEY; allowlist
updated vision.googleapis.com → api.mistral.ai.
MISTRAL_API_KEY added to Coolify container env.

Invariants: G1 (single OCR fallback path, not parallel), G2 (no
duplicate extractor route), INV-AH (Mistral handles gershayim
natively; quote-fix only on PyMuPDF path).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-27 10:11:45 +00:00

448 lines
15 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Text extraction from PDF, DOCX, DOC, and RTF files.
Primary PDF extraction: PyMuPDF direct text (for born-digital PDFs).
Fallback: Mistral OCR (for scanned documents or broken OCR layers).
Routing logic (document-level, not per-page):
1. PyMuPDF extracts text from every page.
2. Pages are quality-checked via _text_quality_ok().
3. If ALL pages pass → use PyMuPDF output (free, ~50ms, no API call).
4. If ANY page fails → call Mistral OCR once for the entire PDF.
Mistral returns per-page Markdown; page_offsets are computed from it.
DOC files: converted to DOCX via LibreOffice before extraction.
Post-processing: Hebrew abbreviation quote fixer (PyMuPDF path only;
Mistral handles gershayim natively).
"""
from __future__ import annotations
import asyncio
import base64
import io
import logging
import re
import subprocess
import tempfile
from pathlib import Path
import fitz # PyMuPDF
import httpx
from PIL import Image
from docx import Document as DocxDocument
from striprtf.striprtf import rtf_to_text
from legal_mcp import config
from legal_mcp.services import storage
logger = logging.getLogger(__name__)
# ── Mistral OCR ───────────────────────────────────────────────────
_MISTRAL_OCR_URL = "https://api.mistral.ai/v1/ocr"
_MISTRAL_OCR_MODEL = "mistral-ocr-latest"
async def _call_mistral_ocr(path: Path) -> list[str]:
"""Call Mistral OCR API on a PDF. Returns per-page Markdown text list.
The Mistral response contains ``pages[i].markdown`` for each page.
If the response has fewer pages than the PDF, trailing pages are
padded with empty strings by the caller.
"""
if not config.MISTRAL_API_KEY:
raise RuntimeError(
"MISTRAL_API_KEY not configured — cannot OCR scanned PDF. "
"Set the env var in Coolify."
)
pdf_b64 = base64.b64encode(path.read_bytes()).decode()
async with httpx.AsyncClient(timeout=300.0) as client:
resp = await client.post(
_MISTRAL_OCR_URL,
headers={
"Authorization": f"Bearer {config.MISTRAL_API_KEY}",
"Content-Type": "application/json",
},
json={
"model": _MISTRAL_OCR_MODEL,
"document": {
"type": "document_url",
"document_url": f"data:application/pdf;base64,{pdf_b64}",
},
"include_image_base64": False,
},
)
if resp.status_code != 200:
raise RuntimeError(
f"Mistral OCR returned {resp.status_code}: {resp.text[:400]}"
)
pages = resp.json().get("pages", [])
return [p.get("markdown", "") for p in pages]
# ── Hebrew text quality detection ────────────────────────────────
_HEBREW_RE = re.compile(r'[֐-׿]')
_WORD_RE = re.compile(r'\S+')
def _text_quality_ok(text: str) -> bool:
"""Check if PyMuPDF-extracted text is genuine Hebrew legal content.
Returns True if text appears to be real content.
Broken OCR layers from scanned PDFs often have:
- Very short words / single-character fragments
- Each word on its own line (high words-per-line ratio)
- Non-Hebrew characters mixed in
"""
words = _WORD_RE.findall(text)
if len(words) < 10:
return False
avg_len = sum(len(w) for w in words) / len(words)
if avg_len < 2.5:
return False
single_char_pct = sum(1 for w in words if len(w) == 1) / len(words)
if single_char_pct > 0.4:
return False
lines = [l for l in text.split("\n") if l.strip()]
if lines and len(words) / len(lines) < 3.0:
return False
letters = re.findall(r'[a-zA-Z֐-׿]', text)
if letters:
hebrew_pct = sum(1 for c in letters if _HEBREW_RE.match(c)) / len(letters)
if hebrew_pct < 0.5:
return False
return True
# ── Hebrew abbreviation quote fixer (PyMuPDF path only) ──────────
_HEBREW_ABBREV_FIXES: dict[str, str] = {
'עוהייד': 'עוה"ד',
'עוייד': 'עו"ד',
'הנייל': 'הנ"ל',
'מצייב': 'מצ"ב',
'ביהמייש': 'ביהמ"ש',
'תייז': 'ת"ז',
'עייי': 'ע"י',
'אחייכ': 'אח"כ',
'סייק': 'ס"ק',
'דייר': 'ד"ר',
'כדוייח': 'כדו"ח',
'חווייד': 'חוו"ד',
'מייר': 'מ"ר',
'יחייד': 'יח"ד',
'בייכ': 'ב"כ',
# Patterns where double-yod (יי) substitutes for gershayim (״) in born-digital PDFs
'בליימ': 'בל"מ',
'תמייא': 'תמ"א',
}
_ABBREV_PATTERN = re.compile(
'|'.join(re.escape(k) for k in sorted(_HEBREW_ABBREV_FIXES, key=len, reverse=True))
)
_HEBREW_YEAR_RE = re.compile(r'(תש[א-ת]+)יי([א-ת])')
def _fix_hebrew_quotes(text: str) -> str:
"""Fix gershayim encoded as double-yod in born-digital PDFs."""
text = _ABBREV_PATTERN.sub(lambda m: _HEBREW_ABBREV_FIXES[m.group()], text)
text = _HEBREW_YEAR_RE.sub(r'\1"\2', text)
return text
# ── Page joining ──────────────────────────────────────────────────
# Separator used when joining per-page text. Constant so chunker /
# retrofit can reproduce the join when computing page offsets.
PAGE_SEPARATOR = "\n\n"
def _join_pages(pages_text: list[str]) -> tuple[str, list[int]]:
"""Join per-page text with PAGE_SEPARATOR while recording start offsets."""
offsets: list[int] = []
parts: list[str] = []
cursor = 0
for i, pg in enumerate(pages_text):
offsets.append(cursor)
parts.append(pg)
cursor += len(pg)
if i < len(pages_text) - 1:
parts.append(PAGE_SEPARATOR)
cursor += len(PAGE_SEPARATOR)
return "".join(parts), offsets
# ── PDF extraction ────────────────────────────────────────────────
async def _extract_pdf(path: Path) -> tuple[str, int, list[int]]:
"""Extract text from PDF using document-level routing.
Stage 1 — PyMuPDF pre-screen (free, ~50ms, no API call):
Run on every page. Collect per-page text; flag pages where
PyMuPDF returns < 50 chars or _text_quality_ok() fails
(scanned, blank, or broken embedded OCR layer).
Stage 2 — Mistral OCR (triggered when any page fails Stage 1):
Send the entire PDF to Mistral once. Use its per-page Markdown
for ALL pages — consistent source, no mixed plain/Markdown formats.
Mistral handles gershayim natively; no quote-fix applied.
Page offsets are always computed so the chunker can attribute each
chunk to its source page number (multimodal hybrid retrieval).
"""
doc = fitz.open(str(path))
page_count = len(doc)
# Stage 1: PyMuPDF pre-screen
pymupdf_pages: list[str] = []
failed: list[int] = []
for i in range(page_count):
text = doc[i].get_text().strip()
if len(text) > 50 and _text_quality_ok(text):
pymupdf_pages.append(_fix_hebrew_quotes(text))
else:
pymupdf_pages.append("")
failed.append(i)
doc.close()
if not failed:
logger.debug(
"PDF %s: all %d pages digital — PyMuPDF only", path.name, page_count
)
joined, offsets = _join_pages(pymupdf_pages)
return joined, page_count, offsets
# Stage 2: Mistral OCR for entire document
logger.info(
"PDF %s: %d/%d pages failed quality check → Mistral OCR",
path.name, len(failed), page_count,
)
mistral_pages = await _call_mistral_ocr(path)
# Pad if Mistral returns fewer pages than PyMuPDF counted
while len(mistral_pages) < page_count:
mistral_pages.append("")
joined, offsets = _join_pages(mistral_pages[:page_count])
return joined, page_count, offsets
def page_at_offset(offset: int, page_offsets: list[int]) -> int:
"""Return the 1-based page number containing a given char offset.
page_offsets[i] is the start of page (i+1) in the joined text.
"""
if not page_offsets:
return 1
page = 1
for i, start in enumerate(page_offsets):
if start <= offset:
page = i + 1
else:
break
return page
# ── Public entry point ────────────────────────────────────────────
async def extract_text(file_path: str) -> tuple[str, int, list[int] | None]:
"""Extract text from a document file.
Returns:
``(text, page_count, page_offsets)`` where:
- ``text``: extracted text. Plain text for PyMuPDF path;
Markdown for Mistral path (tables, ``##`` headers preserved).
- ``page_count``: number of pages (0 for non-PDF)
- ``page_offsets``: char start of each page inside ``text``,
or ``None`` for non-PDF formats
"""
path = Path(file_path)
suffix = path.suffix.lower()
if suffix == ".pdf":
return await _extract_pdf(path)
elif suffix == ".docx":
return _extract_docx(path), 0, None
elif suffix == ".doc":
return _extract_doc(path), 0, None
elif suffix == ".rtf":
return _extract_rtf(path), 0, None
elif suffix in (".txt", ".md"):
return path.read_text(encoding="utf-8"), 0, None
else:
raise ValueError(f"Unsupported file type: {suffix}")
# ── Non-PDF formats ───────────────────────────────────────────────
def _extract_doc(path: Path) -> str:
"""Extract text from legacy .doc via LibreOffice → DOCX conversion."""
with tempfile.TemporaryDirectory() as tmp_dir:
# Isolate the LibreOffice user profile per call: headless soffice
# locks a single shared profile, so concurrent .doc conversions would
# otherwise fail with a profile-lock error.
result = subprocess.run(
[
"libreoffice",
f"-env:UserInstallation=file://{tmp_dir}/lo-profile",
"--headless", "--convert-to", "docx", str(path), "--outdir", tmp_dir,
],
capture_output=True, text=True, timeout=120,
)
if result.returncode != 0:
raise RuntimeError(f"LibreOffice conversion failed: {result.stderr}")
docx_path = Path(tmp_dir) / f"{path.stem}.docx"
if not docx_path.exists():
raise FileNotFoundError(f"Converted file not found: {docx_path}")
return _extract_docx(docx_path)
def _extract_docx(path: Path) -> str:
"""Extract text from DOCX file."""
doc = DocxDocument(str(path))
paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
return "\n\n".join(paragraphs)
def _extract_rtf(path: Path) -> str:
"""Extract text from RTF file."""
rtf_content = path.read_text(encoding="utf-8", errors="replace")
return rtf_to_text(rtf_content)
# ── Multimodal page rendering (V9) ───────────────────────────────
# Unchanged — multimodal embedding always uses PyMuPDF-rendered images
# regardless of whether text extraction used PyMuPDF or Mistral.
def _pixmap_to_pil(pix: fitz.Pixmap) -> Image.Image:
"""Convert a PyMuPDF pixmap to PIL.Image (RGB)."""
if pix.alpha:
pix = fitz.Pixmap(pix, 0)
return Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
def render_pages_for_multimodal(
pdf_path: str | Path,
embed_dpi: int,
thumb_dpi: int | None = None,
thumbnail_dir: Path | None = None,
) -> list[tuple[Image.Image, Path | None]]:
"""Render each PDF page as PIL.Image at ``embed_dpi`` for the
multimodal embedder, and optionally save JPEG thumbnails.
Returns ``[(pil_image, thumb_path_or_None), ...]`` in page order.
"""
src = Path(pdf_path)
if not src.is_file():
raise FileNotFoundError(f"PDF not found: {src}")
if thumbnail_dir is not None:
thumbnail_dir.mkdir(parents=True, exist_ok=True)
out: list[tuple[Image.Image, Path | None]] = []
doc = fitz.open(str(src))
try:
for page_idx, page in enumerate(doc):
page_num = page_idx + 1
pix = page.get_pixmap(dpi=embed_dpi)
img = _pixmap_to_pil(pix)
thumb_path: Path | None = None
if thumbnail_dir is not None and thumb_dpi:
thumb_path = thumbnail_dir / f"p{page_num:03d}.jpg"
ratio = thumb_dpi / embed_dpi
thumb_size = (
max(1, int(img.width * ratio)),
max(1, int(img.height * ratio)),
)
thumb = img.resize(thumb_size, Image.Resampling.LANCZOS)
_tbuf = io.BytesIO()
thumb.save(_tbuf, "JPEG", quality=75, optimize=True)
try:
_tkey = thumb_path.resolve().relative_to(
Path(config.DATA_DIR).resolve()).as_posix()
storage.put_bytes_sync(
_tkey, _tbuf.getvalue(), bucket=storage.Bucket.DERIVED,
content_type="image/jpeg")
except ValueError:
thumb.save(thumb_path, "JPEG", quality=75, optimize=True)
out.append((img, thumb_path))
finally:
doc.close()
return out
# ── Nevo preamble stripping ───────────────────────────────────────
_NEVO_MARKERS = ("ספרות:", "חקיקה שאוזכרה:", "מיני-רציו:", "פסקי דין שאוזכרו:",
"כתבי עת:", "הועתק מנבו")
_DECISION_START = re.compile(
r"^[ \t>*_#]{0,6}(?:"
r"בפנינו|לפנינו|לפניי|הערר שבנדון|ועדת הערר לתכנון|רקע עובדתי|עסקינן|"
r"פסק[ \t\-]{0,3}די(?:ן|נו)|"
r"(?:כב(?:וד)?['׳\"]?\s*)?(?:ה?שופט[ת]?|ה?נשיא[ה]?|המשנה לנשיא)\s+[^\n,]{1,40}:"
r")",
re.MULTILINE,
)
def strip_nevo_preamble(text: str) -> str:
"""Remove Nevo database preamble (bibliography, legislation, mini-ratio).
Returns the original text unchanged if no preamble is detected.
Works on both plain text (PyMuPDF) and Markdown (Mistral) since
_DECISION_START already tolerates leading ``[ \t>*_#]{0,6}``.
"""
head = text[:1500]
if not any(marker in head for marker in _NEVO_MARKERS):
return text
m = _DECISION_START.search(text)
if m and m.start() > 50:
stripped = text[m.start():]
logger.debug("Stripped %d chars of Nevo preamble", m.start())
return stripped
return text
_RATIO_MARKER = "מיני-רציו:"
def extract_nevo_ratio(text: str) -> str:
"""Return the Nevo מיני-רציו block (editorial holdings summary), or ''."""
if not text:
return ""
start = text.find(_RATIO_MARKER)
if start == -1:
return ""
body = text[start + len(_RATIO_MARKER):]
end = len(body)
dm = _DECISION_START.search(body)
if dm:
end = min(end, dm.start())
for marker in _NEVO_MARKERS:
if marker == _RATIO_MARKER:
continue
pos = body.find(marker)
if pos != -1:
end = min(end, pos)
return body[:end].strip()