zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-25 09:53:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Mistral OCR (/v1/ocr) as a document parser; fix "Can't find model" mis-tag (#5782, #7075) (#17057)

Browse Source 
### What problem does this PR solve?

Adds first-class support for **Mistral OCR** (`POST /v1/ocr`) as a
document parser, and fixes the long-standing bug where selecting
`mistral-ocr-latest` fails with `Can't find model for
<tenant>/image2text/mistral-ocr-latest`.

`mistral-ocr-latest` is Mistral's dedicated document-OCR endpoint, not a
vision-chat (`image2text`) model, but the catalog tagged it `image2text`
— so it resolved to the `CvModel` registry, which has no `Mistral`
entry, and there was no `OcrModel` entry either. This PR registers it
correctly and wires it end to end.

Closes #17056
Closes #5782
Closes #7075

**What it does**

1. **`MistralParser` + `MistralOcrModel`**
(`deepdoc/parser/mistral_parser.py`, `rag/llm/ocr_model.py`) — a proper
`OcrModel` factory `Mistral OCR`, mirroring the SoMark cloud-OCR
template. Tables stay inline as HTML; the page range maps to Mistral's
native `pages` selector (absolute page indices, billed per selected
page, so multi-task documents do not re-OCR the whole file); documents
over the inline limit go through the `/v1/files` signed-URL flow with
cleanup.
2. **Removes the `image2text` mis-tag** for `mistral-ocr-latest` from
the `Mistral` factory in `conf/llm_factories.json` (it now lives only in
the `Mistral OCR` factory, typed `ocr`). This is what closes the `Can't
find model` path.
3. **`MistralCV`** (`rag/llm/cv_model.py`) — a thin `GptV4` subclass
over Mistral's OpenAI-compatible endpoint, registering a `Mistral` entry
in the `CvModel` registry so Mistral vision models (`pixtral-*`) become
usable as `image2text` at all.
4. **Figure description** — Mistral-OCR-extracted figures are captioned
using the tenant's configured `image2text` model (any provider),
matching MinerU/deepdoc behaviour.
5. **Wires the parser into every chunking method** (`naive`, `paper`,
`book`, `laws`, `manual`, `one`, `presentation`) and the `rag/flow` DAG
path. This also fixes a related latent gap where those chunkers
forwarded only `mineru_llm_name`, so any model-based OCR provider
selected on a non-`naive` method silently fell through.

**Notes on the API contract** (verified against the live Mistral API):
`pages` is a selector (returns absolute `index`, bills only the
requested pages); `include_blocks: true` returns per-block bounding
boxes usable for chunk highlighting and figure cropping; large files use
`POST /v1/files` → signed URL → OCR → `DELETE`.

**Testing**: new unit tests cover the response→sections contract (both
the 2-tuple `naive` path and the typed 3-tuple DAG path), the
position-tag rescale, the HTTP client incl. upload failure/cleanup
paths, `parse_pdf` page-range threading, registry registration, env
config, the suffix normalization, the factory catalog entry, `MistralCV`
registration, and figure-description injection. Verified end to end
against the live Mistral API on real PDFs (table extraction,
page-selector cost avoidance, figure captioning).

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Xavierando
2026-07-24 15:07:48 +02:00
committed by GitHub
parent bbd0dc5463
commit 08332501a8
22 changed files with 1412 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

487
deepdoc/parser/mistral_parser.py Normal file
View File

@@ -0,0 +1,487 @@
#
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Mistral OCR PDF parser.
Calls Mistral's ``POST /v1/ocr`` endpoint (a dedicated document-OCR API, not
chat-completions), then normalizes the per-block response into the same shape
SoMark produces so the proven section-building contract is reused verbatim.
"""
import base64
import logging
import re
from io import BytesIO
from os import PathLike
from pathlib import Path
from typing import Optional
import numpy as np
import pdfplumber
import requests
from PIL import Image
from deepdoc.parser.pdf_parser import MAXIMUM_PAGE_NUMBER, RAGFlowPdfParser
from deepdoc.parser.utils import extract_pdf_outlines
# Mistral block "type" -> RAGFlow internal layout type.
# header/footer are resolved at runtime via keep_header_footer.
MISTRAL_TYPE_TO_RAGFLOW = {
"text": "text",
"title": "text",
"list": "text",
"table": "table",
"image": "image",
"equation": "equation",
"code": "code",
}
class MistralParser(RAGFlowPdfParser):
"""Parse a PDF through Mistral OCR into (sections, tables)."""
def __init__(
self,
base_url: str,
api_key: str = "",
*,
model: str = "mistral-ocr-latest",
table_format: str = "html",
keep_header_footer: bool = False,
timeout: int = 600,
inline_max_bytes: int = 20 * 1024 * 1024,
):
self.base_url = (base_url or "https://api.mistral.ai/v1").strip().rstrip("/")
self.api_key = api_key
self.model = model
self.table_format = table_format.strip().lower() if table_format else "html"
self.keep_header_footer = bool(keep_header_footer)
self.timeout = int(timeout)
self.inline_max_bytes = int(inline_max_bytes)
self.outlines: list = []
self.logger = logging.getLogger(self.__class__.__name__)
# RAGFlowPdfParser.__init__ sets this default; MistralParser skips
# that heavy __init__, so crop() needs it set before __images__ runs.
self.page_from = 0
# Optional tenant vision model (LLMBundle) for figure description;
# set from parse_pdf's vision_model kwarg. None -> no enrichment.
self.vision_model = None
# ------------------------------------------------------------------
# Page image rendering
# ------------------------------------------------------------------
def __images__(self, fnm, zoomin: int = 1, page_from: int = 0, page_to: int = MAXIMUM_PAGE_NUMBER, callback=None):
self.page_from = page_from
self.page_to = page_to
try:
ctx = pdfplumber.open(fnm) if isinstance(fnm, (str, PathLike)) else pdfplumber.open(BytesIO(fnm))
with ctx as pdf:
self.pdf = pdf
self.page_images = [p.to_image(resolution=72 * zoomin, antialias=True).original for _, p in enumerate(self.pdf.pages[page_from:page_to])]
except Exception as exc:
self.page_images = None
self.total_page = 0
self.logger.exception(exc)
# ------------------------------------------------------------------
# Block classification
# ------------------------------------------------------------------
def _resolve_internal_type(self, block_type: Optional[str]) -> Optional[str]:
"""Return the RAGFlow internal layout type, or None to discard.
header/footer obey keep_header_footer; unknown types fall back to text
to avoid silent loss.
"""
if block_type in ("header", "footer"):
return "text" if self.keep_header_footer else None
return MISTRAL_TYPE_TO_RAGFLOW.get(block_type, "text")
@staticmethod
def _block_text(block: dict, internal_type: str) -> str:
"""Textual payload for a block. image blocks contribute no text."""
if internal_type == "image":
return ""
return (block.get("content") or "").strip()
# ------------------------------------------------------------------
# Response -> normalized pages (SoMark-compatible block shape)
# ------------------------------------------------------------------
def _normalize_pages(self, response: dict) -> list[dict]:
"""Map a /v1/ocr response into per-page dicts whose blocks carry a
[x0, top, x1, bott] bbox and a {w,h} page_size, so downstream section
building is identical to SoMark's."""
out: list[dict] = []
for page in (response or {}).get("pages") or []:
dims = page.get("dimensions") or {}
page_size = {"w": dims.get("width") or 1, "h": dims.get("height") or 1}
blocks = []
for b in page.get("blocks") or []:
coord_keys = ("top_left_x", "top_left_y", "bottom_right_x", "bottom_right_y")
if any(k in b for k in coord_keys):
bbox = [
b.get("top_left_x", 0),
b.get("top_left_y", 0),
b.get("bottom_right_x", 0),
b.get("bottom_right_y", 0),
]
else:
bbox = None
blocks.append(
{
"type": b.get("type"),
"content": b.get("content"),
"bbox": bbox,
}
)
out.append(
{
"page_num": page.get("index") or 0,
"page_size": page_size,
"blocks": blocks,
}
)
return out
# ------------------------------------------------------------------
# Position tagging (compatible with RAGFlow extract_positions/crop)
# ------------------------------------------------------------------
def _line_tag(self, bx: dict) -> str:
"""Build a ``@@page\tx0\tx1\ttop\tbott##`` tag (page 1-based, order
x0,x1,top,bott), rescaling bbox from Mistral's dimensions space to the
locally rendered page pixels."""
page_idx = bx.get("page_idx", 0)
pn = [page_idx + 1]
bbox = bx.get("bbox") or [0, 0, 0, 0]
if len(bbox) != 4:
bbox = [0, 0, 0, 0]
x0, top, x1, bott = bbox
page_size = bx.get("page_size") or {}
src_w = page_size.get("w") or 1
src_h = page_size.get("h") or 1
if x0 > x1:
x0, x1 = x1, x0
if top > bott:
top, bott = bott, top
if getattr(self, "page_images", None) and len(self.page_images) > page_idx:
page_width, page_height = self.page_images[page_idx].size
x0 = max(0.0, (x0 / src_w) * page_width)
x1 = max(0.0, (x1 / src_w) * page_width)
top = max(0.0, (top / src_h) * page_height)
bott = max(0.0, (bott / src_h) * page_height)
return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##".format("-".join(str(x) for x in pn), x0, x1, top, bott)
@staticmethod
def extract_positions(txt: str):
poss = []
for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", txt):
pn, left, right, top, bottom = tag.strip("#").strip("@").split("\t")
left, right, top, bottom = float(left), float(right), float(top), float(bottom)
poss.append(([int(p) - 1 for p in pn.split("-")], left, right, top, bottom))
return poss
def crop(self, text, ZM=1, need_position=False):
"""Image crop based on tags."""
imgs = []
poss = self.extract_positions(text)
if not poss:
return (None, None) if need_position else None
if not getattr(self, "page_images", None):
self.logger.warning("[Mistral OCR] crop called without page images; skip.")
return (None, None) if need_position else None
page_count = len(self.page_images)
filtered = []
for pns, left, right, top, bottom in poss:
valid_pns = [p for p in pns if 0 <= p < page_count]
if valid_pns:
filtered.append((valid_pns, left, right, top, bottom))
if not filtered:
return (None, None) if need_position else None
poss = filtered
max_width = max(np.max([right - left for (_, left, right, _, _) in poss]), 6)
GAP = 6
first = poss[0]
poss.insert(
0,
(
[first[0][0]],
first[1],
first[2],
max(0, first[3] - 120),
max(first[3] - GAP, 0),
),
)
last = poss[-1]
last_pn = last[0][-1]
if not (0 <= last_pn < page_count):
return (None, None) if need_position else None
last_h = self.page_images[last_pn].size[1]
poss.append(
(
[last_pn],
last[1],
last[2],
min(last_h, last[4] + GAP),
min(last_h, last[4] + 120),
)
)
positions = []
for ii, (pns, left, right, top, bottom) in enumerate(poss):
right = left + max_width
if bottom <= top:
bottom = top + 2
for pn in pns[1:]:
if 0 <= pn - 1 < page_count:
bottom += self.page_images[pn - 1].size[1]
if not (0 <= pns[0] < page_count):
continue
base_img = self.page_images[pns[0]]
x0, y0, x1, y1 = (
int(left),
int(top),
int(right),
int(min(bottom, base_img.size[1])),
)
if x0 > x1:
x0, x1 = x1, x0
if y0 > y1:
y0, y1 = y1, y0
if x1 <= x0 or y1 <= y0:
continue
imgs.append(base_img.crop((x0, y0, x1, y1)))
if 0 < ii < len(poss) - 1:
positions.append((pns[0] + self.page_from, x0, x1, y0, y1))
bottom -= base_img.size[1]
for pn in pns[1:]:
if not (0 <= pn < page_count):
continue
page = self.page_images[pn]
x0, y0, x1, y1 = (
int(left),
0,
int(right),
int(min(bottom, page.size[1])),
)
if x0 > x1:
x0, x1 = x1, x0
if y0 > y1:
y0, y1 = y1, y0
if x1 <= x0 or y1 <= y0:
bottom -= page.size[1]
continue
imgs.append(page.crop((x0, y0, x1, y1)))
if 0 < ii < len(poss) - 1:
positions.append((pn + self.page_from, x0, x1, y0, y1))
bottom -= page.size[1]
if not imgs:
return (None, None) if need_position else None
height = sum(img.size[1] + GAP for img in imgs)
width = int(np.max([i.size[0] for i in imgs]))
pic = Image.new("RGB", (width, int(height)), (245, 245, 245))
offset = 0
for ii, img in enumerate(imgs):
if ii == 0 or ii + 1 == len(imgs):
img = img.convert("RGBA")
overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
overlay.putalpha(128)
img = Image.alpha_composite(img, overlay).convert("RGB")
pic.paste(img, (0, int(offset)))
offset += img.size[1] + GAP
return (pic, positions) if need_position else pic
# ------------------------------------------------------------------
# Sections / tables
# ------------------------------------------------------------------
def _describe_image(self, line_tag: str) -> str:
"""Best-effort VLM caption of the figure at ``line_tag`` using the tenant
vision model (``self.vision_model``). Crops the figure from the locally
rendered page and describes it. Returns "" on any failure so figure
enrichment never breaks parsing (mirrors MinerU's image enhancement)."""
try:
img = self.crop("figure" + line_tag, ZM=1)
if img is None:
return ""
from rag.app.picture import vision_llm_chunk
from rag.prompts.generator import vision_llm_figure_describe_prompt
# vision_llm_chunk expects a PIL Image (it calls img.size / img.save),
# not raw bytes — pass the crop directly.
desc = vision_llm_chunk(binary=img, vision_model=self.vision_model, prompt=vision_llm_figure_describe_prompt())
return (desc or "").strip()
except Exception:
self.logger.info("[Mistral OCR] figure description skipped", exc_info=True)
return ""
def _transfer_to_sections(self, pages: list[dict], parse_method: Optional[str] = None) -> list[tuple]:
"""manual/pipeline (rag/flow DAG) want typed 3-tuples
(text, layout_type, line_tag); every other caller (naive.py) wants
2-tuples (text, line_tag) that naive_merge consumes directly."""
typed = parse_method in {"manual", "pipeline"}
sections: list[tuple] = []
image_seq = 0
for page in pages or []:
page_idx = page.get("page_num", 0)
page_size = page.get("page_size") or {}
for block in page.get("blocks") or []:
internal = self._resolve_internal_type(block.get("type"))
if internal is None:
continue
bbox = block.get("bbox")
tag_input = {"page_idx": page_idx, "bbox": bbox, "page_size": page_size}
if internal == "image":
if not bbox or len(bbox) != 4:
continue # no geometry -> nothing to crop
line_tag = self._line_tag(tag_input)
image_seq += 1
caption = (block.get("content") or "").strip()
description = self._describe_image(line_tag) if self.vision_model is not None else ""
label = description or caption or f"image {image_seq}"
if typed:
sections.append((label, internal, line_tag))
else:
# chunk id is hash(content + doc_id): empty image text would
# collide across figures. Keep a unique caption; append the
# tag so crop() recovers the figure, remove_tag() strips it.
sections.append((f"{label}{line_tag}", ""))
continue
text = self._block_text(block, internal)
if not text:
continue
line_tag = self._line_tag(tag_input)
if typed:
sections.append((text, internal, line_tag))
else:
sections.append((text, line_tag))
return sections
def _transfer_to_tables(self, pages: list[dict]) -> list:
# Tables are inlined as HTML in section text; no separate extraction.
return []
# ------------------------------------------------------------------
# HTTP client
# ------------------------------------------------------------------
def _headers(self) -> dict:
return {"Authorization": f"Bearer {self.api_key}"}
def check_installation(self) -> tuple[bool, str]:
if not self.api_key:
return False, "Mistral API key is not configured."
return True, ""
def _raise_for_status(self, resp, action: str) -> None:
if resp.status_code != 200:
raise RuntimeError(f"Mistral {action} failed: {resp.status_code} {resp.text[:300]}")
def _ocr_payload(self, document: dict, pages: Optional[list[int]]) -> dict:
payload = {
"model": self.model,
"document": document,
"include_blocks": True,
"table_format": self.table_format,
"include_image_base64": False,
}
if pages:
payload["pages"] = pages
return payload
def _call_ocr(self, pdf_bytes: bytes, filename: str, pages: Optional[list[int]], callback=None) -> dict:
ok, reason = self.check_installation()
if not ok:
raise RuntimeError(reason)
if len(pdf_bytes) <= self.inline_max_bytes:
b64 = base64.b64encode(pdf_bytes).decode()
document = {"type": "document_url", "document_url": f"data:application/pdf;base64,{b64}"}
return self._post_ocr(self._ocr_payload(document, pages))
# Large file: upload -> signed url -> ocr -> delete.
file_id = None
try:
r = requests.post(f"{self.base_url}/files", headers=self._headers(), files={"file": (filename, pdf_bytes, "application/pdf")}, data={"purpose": "ocr"}, timeout=self.timeout)
self._raise_for_status(r, "/files upload")
file_id = r.json().get("id")
r = requests.get(f"{self.base_url}/files/{file_id}/url", headers=self._headers(), params={"expiry": 24}, timeout=self.timeout)
self._raise_for_status(r, "signed-url fetch")
signed = r.json().get("url")
document = {"type": "document_url", "document_url": signed}
return self._post_ocr(self._ocr_payload(document, pages))
finally:
if file_id:
try:
resp = requests.delete(f"{self.base_url}/files/{file_id}", headers=self._headers(), timeout=self.timeout)
if not 200 <= resp.status_code < 300:
self.logger.warning("failed to delete uploaded file %s: %s %s", file_id, resp.status_code, resp.text[:200])
except Exception:
self.logger.warning("failed to delete uploaded file %s", file_id)
def _post_ocr(self, payload: dict) -> dict:
r = requests.post(f"{self.base_url}/ocr", headers=self._headers(), json=payload, timeout=self.timeout)
self._raise_for_status(r, "OCR")
return r.json()
# ------------------------------------------------------------------
# Public entry point
# ------------------------------------------------------------------
def parse_pdf(self, filepath: str | PathLike[str], binary=None, callback=None, parse_method: str = "raw", from_page: int = 0, to_page: int = MAXIMUM_PAGE_NUMBER, **kwargs) -> tuple[list, list]:
# Optional tenant vision model for figure description (best-effort).
self.vision_model = kwargs.pop("vision_model", None)
# Load bytes.
if binary is not None:
pdf_bytes = binary.getvalue() if hasattr(binary, "getvalue") else bytes(binary)
else:
pdf_bytes = Path(filepath).read_bytes()
self.outlines = extract_pdf_outlines(pdf_bytes)
# Render the WHOLE document locally: _line_tag/crop index page_images in
# absolute page order, so a sliced render would crop the wrong page.
self.__images__(pdf_bytes, zoomin=1)
# pages is a selector: include only when the caller restricted the range.
pages: Optional[list[int]] = None
if from_page > 0 or to_page < MAXIMUM_PAGE_NUMBER:
# Bound the selector by the real page count, never by the sentinel
# to_page (MAXIMUM_PAGE_NUMBER) — otherwise a failed render would
# build a ~100k-element list. Fall back to the PDF page count.
rendered = len(self.page_images) if getattr(self, "page_images", None) else 0
total = rendered or self.total_page_number(Path(filepath).name, pdf_bytes)
end = min(to_page, total) if total else from_page
pages = list(range(from_page, end))
if not pages:
return [], []
if callback:
callback(0.15, "[Mistral OCR] submitting document")
response = self._call_ocr(pdf_bytes, Path(filepath).name, pages, callback=callback)
norm = self._normalize_pages(response)
if callback:
n_blocks = sum(len(p.get("blocks") or []) for p in norm)
callback(0.75, f"[Mistral OCR] parsed {n_blocks} blocks across {len(norm)} pages")
sections = self._transfer_to_sections(norm, parse_method)
tables = self._transfer_to_tables(norm)
return sections, tables