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>
448 lines
15 KiB
Python
448 lines
15 KiB
Python
"""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()
|