diff --git a/api/apps/services/provider_api_service.py b/api/apps/services/provider_api_service.py index 0038777f08..d030074b5b 100644 --- a/api/apps/services/provider_api_service.py +++ b/api/apps/services/provider_api_service.py @@ -83,7 +83,7 @@ def list_providers(tenant_id: str, all_available: bool = False): if factory_info["name"] in ["Youdao", "FastEmbed", "BAAI", "Builtin", "siliconflow_intl"]: continue model_types = sorted(set(model_type for llm in factory_info.get("llm", []) for model_type in _factory_model_types(llm))) if factory_info.get("llm", []) else [] - if factory_info["name"] in ["MinerU", "PaddleOCR", "OpenDataLoader"]: + if factory_info["name"] in ["MinerU", "PaddleOCR", "OpenDataLoader", "Mistral OCR"]: model_types.append("ocr") provider = {"model_types": model_types, "name": factory_info["name"], "url": {"default": factory_info.get("url", "")}} if factory_info["name"].lower() == "siliconflow": @@ -105,7 +105,7 @@ def list_providers(tenant_id: str, all_available: bool = False): provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, name) has_instance = bool(provider_obj and TenantModelInstanceService.get_all_by_provider_id(provider_obj.id)) model_types = sorted(set(model_type for llm in factory_info.get("llm", []) for model_type in _factory_model_types(llm))) if factory_info.get("llm", []) else [] - if name in ["MinerU", "PaddleOCR", "OpenDataLoader"]: + if name in ["MinerU", "PaddleOCR", "OpenDataLoader", "Mistral OCR"]: model_types.append("ocr") provider = {"has_instance": has_instance, "model_types": model_types, "name": factory_info["name"], "url": {"default": factory_info.get("url", "")}} diff --git a/api/db/joint_services/tenant_model_service.py b/api/db/joint_services/tenant_model_service.py index 90796dfa4e..46559ffbf7 100644 --- a/api/db/joint_services/tenant_model_service.py +++ b/api/db/joint_services/tenant_model_service.py @@ -24,6 +24,8 @@ from common.constants import ( ModelTypeBinary, MINERU_DEFAULT_CONFIG, MINERU_ENV_KEYS, + MISTRAL_OCR_DEFAULT_CONFIG, + MISTRAL_OCR_ENV_KEYS, OPENDATALOADER_DEFAULT_CONFIG, OPENDATALOADER_ENV_KEYS, PADDLEOCR_DEFAULT_CONFIG, @@ -496,3 +498,12 @@ def ensure_somark_from_env(tenant_id: str) -> str | None: "somark-from-env", _collect_env_config(SOMARK_ENV_KEYS, SOMARK_DEFAULT_CONFIG), ) + + +def ensure_mistral_ocr_from_env(tenant_id: str) -> str | None: + return _ensure_ocr_provider_from_env( + tenant_id, + "Mistral OCR", + "mistral-ocr-latest", + _collect_env_config(MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG), + ) diff --git a/common/constants.py b/common/constants.py index 42d9c9f8a9..139612e42d 100644 --- a/common/constants.py +++ b/common/constants.py @@ -372,3 +372,15 @@ SOMARK_DEFAULT_CONFIG = { "SOMARK_ENABLE_IMAGE_UNDERSTANDING": 1, "SOMARK_KEEP_HEADER_FOOTER": 0, } +MISTRAL_OCR_ENV_KEYS = [ + "MISTRAL_OCR_BASE_URL", + "MISTRAL_OCR_API_KEY", + "MISTRAL_OCR_TABLE_FORMAT", + "MISTRAL_OCR_KEEP_HEADER_FOOTER", +] +MISTRAL_OCR_DEFAULT_CONFIG = { + "MISTRAL_OCR_BASE_URL": "https://api.mistral.ai/v1", + "MISTRAL_OCR_API_KEY": "", + "MISTRAL_OCR_TABLE_FORMAT": "html", + "MISTRAL_OCR_KEEP_HEADER_FOOTER": 0, +} diff --git a/common/parser_config_utils.py b/common/parser_config_utils.py index df7daafe1b..f37fd18c3c 100644 --- a/common/parser_config_utils.py +++ b/common/parser_config_utils.py @@ -39,5 +39,10 @@ def normalize_layout_recognizer(layout_recognizer_raw: Any) -> tuple[Any, str | # expects all three segments to locate the provider/instance row. parser_model_name = layout_recognizer_raw layout_recognizer = "SoMark" + elif lowered.endswith("@mistral ocr"): + # Separate OCR-only factory (never the multi-type "Mistral" factory), + # so this suffix cannot collide with pixtral vision models. + parser_model_name = layout_recognizer_raw + layout_recognizer = "Mistral OCR" return layout_recognizer, parser_model_name diff --git a/conf/llm_factories.json b/conf/llm_factories.json index 0f52ded6dc..e3436ee11b 100644 --- a/conf/llm_factories.json +++ b/conf/llm_factories.json @@ -2288,15 +2288,6 @@ "chat" ] }, - { - "llm_name": "mistral-ocr-latest", - "tags": "LLM,IMAGE2TEXT,131k", - "max_tokens": 131000, - "model_type": [ - "image2text", - "chat" - ] - }, { "llm_name": "open-mistral-nemo", "tags": "LLM,CHAT,131k", @@ -7630,6 +7621,21 @@ } ] }, + { + "name": "Mistral OCR", + "logo": "", + "tags": "OCR", + "status": "1", + "rank": "930", + "llm": [ + { + "llm_name": "mistral-ocr-latest", + "tags": "OCR", + "max_tokens": 8192, + "model_type": "ocr" + } + ] + }, { "name": "n1n", "logo": "", diff --git a/deepdoc/parser/mistral_parser.py b/deepdoc/parser/mistral_parser.py new file mode 100644 index 0000000000..a7d5c014a7 --- /dev/null +++ b/deepdoc/parser/mistral_parser.py @@ -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 diff --git a/rag/app/book.py b/rag/app/book.py index a9718eb3ba..6377902f2e 100644 --- a/rag/app/book.py +++ b/rag/app/book.py @@ -110,6 +110,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, **kwargs, ) diff --git a/rag/app/laws.py b/rag/app/laws.py index b2c518ff1f..88222276f4 100644 --- a/rag/app/laws.py +++ b/rag/app/laws.py @@ -202,6 +202,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, **kwargs, ) diff --git a/rag/app/manual.py b/rag/app/manual.py index 766055080e..3206250440 100644 --- a/rag/app/manual.py +++ b/rag/app/manual.py @@ -167,6 +167,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, parse_method="manual", **kwargs, diff --git a/rag/app/naive.py b/rag/app/naive.py index 5062864f4c..0ae8cc801f 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -350,6 +350,64 @@ def by_somark( return None, None, None +def by_mistral_ocr( + filename, + binary=None, + from_page=0, + to_page=MAXIMUM_PAGE_NUMBER, + lang="Chinese", + callback=None, + pdf_cls=None, + parse_method: str = "raw", + mistral_ocr_llm_name: str | None = None, + tenant_id: str | None = None, + **kwargs, +): + pdf_parser = None + if tenant_id: + if not mistral_ocr_llm_name: + try: + from api.db.joint_services.tenant_model_service import ensure_mistral_ocr_from_env + + mistral_ocr_llm_name = ensure_mistral_ocr_from_env(tenant_id) + except Exception as e: + logging.warning(f"fallback to env mistral ocr: {e}") + + if mistral_ocr_llm_name: + try: + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, mistral_ocr_llm_name) + ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) + pdf_parser = ocr_model.mdl + # Best-effort figure description: hand the parser the tenant's + # vision model so Mistral OCR's extracted images get VLM captions + # (parity with MinerU/deepdoc). Skip silently if none is configured. + if "vision_model" not in kwargs: + try: + vision_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.VISION) + kwargs["vision_model"] = LLMBundle(tenant_id=tenant_id, model_config=vision_model_config, lang=lang) + except Exception as vlm_err: + logging.info(f"[Mistral OCR] no vision model for tenant; skipping figure description: {vlm_err}") + sections, tables = pdf_parser.parse_pdf( + filepath=filename, + binary=binary, + callback=callback, + parse_method=parse_method, + from_page=from_page, + to_page=to_page, + **kwargs, + ) + return sections, tables, pdf_parser + except Exception as e: + logging.error(f"Failed to parse pdf via LLMBundle Mistral OCR ({mistral_ocr_llm_name}): {e}") + if callback: + callback(-1, f"Failed to parse pdf via Mistral OCR ({mistral_ocr_llm_name}): {e}") + return None, None, None + + if callback: + callback(-1, "Mistral OCR not found.") + return None, None, None + + def by_plaintext(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, callback=None, **kwargs): layout_recognizer = (kwargs.get("layout_recognizer") or "").strip() if (not layout_recognizer) or (layout_recognizer == "Plain Text"): @@ -378,6 +436,7 @@ PARSERS = { "tcadp parser": by_tcadp, "paddleocr": by_paddleocr, "somark": by_somark, + "mistral ocr": by_mistral_ocr, "plaintext": by_plaintext, # default } @@ -1005,6 +1064,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= paddleocr_llm_name=parser_model_name, opendataloader_llm_name=opendataloader_llm_name, somark_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, **kwargs, ) sections = _normalize_section_text_for_rtl_presentation_forms(sections) @@ -1015,7 +1075,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= if table_context_size or image_context_size: tables = append_context2table_image4pdf(sections, tables, image_context_size) - if name in ["tcadp", "docling", "mineru", "paddleocr", "opendataloader", "somark"]: + if name in ["tcadp", "docling", "mineru", "paddleocr", "opendataloader", "somark", "mistral ocr"]: if int(parser_config.get("chunk_token_num", 0)) <= 0: parser_config["chunk_token_num"] = 0 diff --git a/rag/app/one.py b/rag/app/one.py index ba81a1d2a0..97ac790583 100644 --- a/rag/app/one.py +++ b/rag/app/one.py @@ -106,6 +106,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, **kwargs, ) diff --git a/rag/app/paper.py b/rag/app/paper.py index 5f30bc0da3..23b15161e9 100644 --- a/rag/app/paper.py +++ b/rag/app/paper.py @@ -164,6 +164,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, parse_method="paper", **kwargs, ) diff --git a/rag/app/presentation.py b/rag/app/presentation.py index b2875ee167..d631795c22 100644 --- a/rag/app/presentation.py +++ b/rag/app/presentation.py @@ -214,6 +214,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, + mistral_ocr_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, **kwargs, ) diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index 5374f4c460..0aabc89f4f 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -257,7 +257,7 @@ class ParserParam(ProcessParamBase): pdf_parse_method = pdf_config.get("parse_method", "") self.check_empty(pdf_parse_method, "Parse method abnormal.") - if pdf_parse_method.lower() not in ["deepdoc", "plain_text", "mineru", "docling", "opendataloader", "tcadp parser", "paddleocr", "somark"]: + if pdf_parse_method.lower() not in ["deepdoc", "plain_text", "mineru", "docling", "opendataloader", "tcadp parser", "paddleocr", "somark", "mistral ocr"]: self.check_empty(pdf_config.get("lang", ""), "PDF VLM language") pdf_output_format = pdf_config.get("output_format", "") @@ -369,6 +369,9 @@ class Parser(ProcessBase): # downstream requires all three segments. parser_model_name = raw_parse_method parse_method = "SoMark" + elif lowered.endswith("@mistral ocr"): + parser_model_name = raw_parse_method + parse_method = "Mistral OCR" # DeepDOC returns structured page boxes directly. if parse_method.lower() == "deepdoc": @@ -565,6 +568,52 @@ class Parser(ProcessBase): box["image"] = image bboxes.append(box) + elif parse_method.lower() == "mistral ocr": + + def resolve_mistral_ocr_llm_name(): + configured = parser_model_name or conf.get("mistral_ocr_llm_name") + if configured: + return configured + tenant_id = self._canvas._tenant_id + if not tenant_id: + return None + from api.db.joint_services.tenant_model_service import ensure_mistral_ocr_from_env + + return ensure_mistral_ocr_from_env(tenant_id) + + parser_model_name = resolve_mistral_ocr_llm_name() + if not parser_model_name: + raise RuntimeError("Mistral OCR model not configured. Please add Mistral OCR in Model Providers or set MISTRAL_OCR_* env.") + + tenant_id = self._canvas._tenant_id + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, parser_model_name) + ocr_model = LLMBundle(tenant_id, ocr_model_config) + pdf_parser = ocr_model.mdl + + lines, _ = pdf_parser.parse_pdf( + filepath=name, + binary=blob, + callback=self.callback, + parse_method="pipeline", + ) + bboxes = [] + for item in lines or []: + if not isinstance(item, tuple) or len(item) < 3: + continue + text, layout_type, poss = item[0], item[1], item[2] + box = { + "text": text, + "layout_type": layout_type or "text", + } + if isinstance(poss, str) and poss: + positions = [[pos[0][-1] + 1, *pos[1:]] for pos in pdf_parser.extract_positions(poss)] + if positions: + box["positions"] = positions + image = pdf_parser.crop(poss, 1) + if image is not None: + box["image"] = image + bboxes.append(box) + elif parse_method.lower() == "tcadp parser": # ADP is a document parsing tool using Tencent Cloud API table_result_type = conf.get("table_result_type", "1") diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index 00874b7f06..603601c4b3 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -395,6 +395,15 @@ class xAICV(GptV4): super().__init__(key, model_name, lang=lang, base_url=base_url, **kwargs) +class MistralCV(GptV4): + _FACTORY_NAME = "Mistral" + + def __init__(self, key, model_name="pixtral-12b-2409", lang="Chinese", base_url=None, **kwargs): + if not base_url: + base_url = "https://api.mistral.ai/v1" + super().__init__(key, model_name, lang=lang, base_url=base_url, **kwargs) + + class QWenCV(GptV4): _FACTORY_NAME = "Tongyi-Qianwen" diff --git a/rag/llm/ocr_model.py b/rag/llm/ocr_model.py index c89e606f65..80b077d57e 100644 --- a/rag/llm/ocr_model.py +++ b/rag/llm/ocr_model.py @@ -19,8 +19,10 @@ import os from typing import Any, Optional from deepdoc.parser.mineru_parser import MinerUParser +from deepdoc.parser.mistral_parser import MistralParser from deepdoc.parser.opendataloader_parser import OpenDataLoaderParser from deepdoc.parser.paddleocr_parser import PaddleOCRParser +from deepdoc.parser.pdf_parser import MAXIMUM_PAGE_NUMBER from deepdoc.parser.somark_parser import SoMarkParser @@ -325,3 +327,71 @@ class SoMarkOcrModel(Base, SoMarkParser): **kwargs, ) return sections, tables + + +class MistralOcrModel(Base, MistralParser): + _FACTORY_NAME = "Mistral OCR" + + def __init__(self, key: str | dict, model_name: str, **kwargs): + Base.__init__(self, key, model_name, **kwargs) + raw_config: dict = {} + if isinstance(key, dict): + raw_config = key + elif key: + try: + raw_config = json.loads(key) + except Exception: + raw_config = {} + + # Only unwrap a nested {"api_key": {...}} config object; a flat config + # whose "api_key" is a string must be preserved so the key is not lost. + nested_config = raw_config.get("api_key") if isinstance(raw_config, dict) else None + config = nested_config if isinstance(nested_config, dict) else raw_config + if not isinstance(config, dict): + config = {} + + key_as_secret = key if isinstance(key, str) and key and not key.lstrip().startswith("{") else "" + + def _resolve(ui_key: str, env_key: str, default=""): + return config.get(ui_key, config.get(env_key, kwargs.get(ui_key, kwargs.get(env_key, os.environ.get(env_key, default))))) + + base_url = _resolve("mistral_ocr_base_url", "MISTRAL_OCR_BASE_URL", kwargs.get("base_url") or "https://api.mistral.ai/v1") + api_key = _resolve("api_key", "MISTRAL_OCR_API_KEY", key_as_secret) + table_format = _resolve("mistral_ocr_table_format", "MISTRAL_OCR_TABLE_FORMAT", "html") + keep_hf = _resolve("mistral_ocr_keep_header_footer", "MISTRAL_OCR_KEEP_HEADER_FOOTER", 0) + + # Redact sensitive config keys before logging + redacted_config = {} + for k, v in config.items(): + if any(s in k.lower() for s in ("key", "password", "token", "secret")): + redacted_config[k] = "[REDACTED]" + else: + redacted_config[k] = v + logging.info(f"Parsed Mistral OCR config (sensitive fields redacted): {redacted_config}") + + MistralParser.__init__( + self, + base_url=base_url, + api_key=api_key, + model=model_name, + table_format=table_format, + keep_header_footer=str(keep_hf).strip().lower() in {"1", "true", "yes", "on"}, + ) + + def check_available(self) -> tuple[bool, str]: + return self.check_installation() + + def parse_pdf(self, filepath, binary=None, callback=None, parse_method: str = "raw", from_page: int = 0, to_page: int = MAXIMUM_PAGE_NUMBER, **kwargs): + ok, reason = self.check_available() + if not ok: + raise RuntimeError(f"Mistral OCR not accessible: {reason}") + return MistralParser.parse_pdf( + self, + filepath=filepath, + binary=binary, + callback=callback, + parse_method=parse_method, + from_page=from_page, + to_page=to_page, + **kwargs, + ) diff --git a/test/unit_test/common/test_mistral_ocr_env.py b/test/unit_test/common/test_mistral_ocr_env.py new file mode 100644 index 0000000000..5a26c91f03 --- /dev/null +++ b/test/unit_test/common/test_mistral_ocr_env.py @@ -0,0 +1,24 @@ +def test_env_keys_and_defaults_present(): + from common.constants import MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG + + assert "MISTRAL_OCR_API_KEY" in MISTRAL_OCR_ENV_KEYS + assert "MISTRAL_OCR_BASE_URL" in MISTRAL_OCR_ENV_KEYS + assert MISTRAL_OCR_DEFAULT_CONFIG["MISTRAL_OCR_BASE_URL"] == "https://api.mistral.ai/v1" + + +def test_collect_env_config_returns_none_without_env(monkeypatch): + from common.constants import MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG + from api.db.joint_services.tenant_model_service import _collect_env_config + + for k in MISTRAL_OCR_ENV_KEYS: + monkeypatch.delenv(k, raising=False) + assert _collect_env_config(MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG) is None + + +def test_collect_env_config_populated_when_key_set(monkeypatch): + from common.constants import MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG + from api.db.joint_services.tenant_model_service import _collect_env_config + + monkeypatch.setenv("MISTRAL_OCR_API_KEY", "sk-live") + cfg = _collect_env_config(MISTRAL_OCR_ENV_KEYS, MISTRAL_OCR_DEFAULT_CONFIG) + assert cfg["MISTRAL_OCR_API_KEY"] == "sk-live" diff --git a/test/unit_test/common/test_mistral_ocr_factory.py b/test/unit_test/common/test_mistral_ocr_factory.py new file mode 100644 index 0000000000..b6d30dec5f --- /dev/null +++ b/test/unit_test/common/test_mistral_ocr_factory.py @@ -0,0 +1,25 @@ +import json +from pathlib import Path + + +def test_mistral_ocr_factory_present_and_ocr_tagged(): + repo_root = Path(__file__).resolve().parents[3] + data = json.loads((repo_root / "conf" / "llm_factories.json").read_text()) + factories = {f["name"]: f for f in data["factory_llm_infos"]} + assert "Mistral OCR" in factories, "Mistral OCR factory missing" + fac = factories["Mistral OCR"] + assert "OCR" in fac["tags"] + # ships a default OCR model (like the other OCR factories) so it is usable + # without manually adding one through the model provider page; the llm_name + # is the real Mistral API id because the parser POSTs it verbatim to /v1/ocr. + models = {m["llm_name"]: m for m in fac["llm"]} + assert "mistral-ocr-latest" in models + assert models["mistral-ocr-latest"]["model_type"] == "ocr" + assert "OCR" in models["mistral-ocr-latest"]["tags"] + + +def test_mistral_ocr_factory_distinct_from_mistral(): + repo_root = Path(__file__).resolve().parents[3] + data = json.loads((repo_root / "conf" / "llm_factories.json").read_text()) + names = [f["name"] for f in data["factory_llm_infos"]] + assert "Mistral" in names and "Mistral OCR" in names diff --git a/test/unit_test/common/test_parser_config_utils_mistral.py b/test/unit_test/common/test_parser_config_utils_mistral.py new file mode 100644 index 0000000000..65117ab078 --- /dev/null +++ b/test/unit_test/common/test_parser_config_utils_mistral.py @@ -0,0 +1,23 @@ +from common.parser_config_utils import normalize_layout_recognizer + + +def test_mistral_ocr_suffix_normalized(): + raw = "mistral-ocr-latest@inst@Mistral OCR" + layout, model_name = normalize_layout_recognizer(raw) + assert layout == "Mistral OCR" + assert model_name == raw + + +def test_mistral_ocr_suffix_case_insensitive(): + raw = "mistral-ocr-latest@inst@MISTRAL OCR" + layout, model_name = normalize_layout_recognizer(raw) + assert layout == "Mistral OCR" + assert model_name == raw + + +def test_plain_mistral_not_captured(): + # a pixtral vision model on the multi-type Mistral factory must NOT be + # dirverted to OCR + layout, model_name = normalize_layout_recognizer("pixtral-large-latest@inst@Mistral") + assert layout == "pixtral-large-latest@inst@Mistral" + assert model_name is None diff --git a/test/unit_test/deepdoc/parser/test_mistral_ocr_model.py b/test/unit_test/deepdoc/parser/test_mistral_ocr_model.py new file mode 100644 index 0000000000..5cb803927d --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_mistral_ocr_model.py @@ -0,0 +1,33 @@ +def test_mistral_ocr_model_registers_in_registry(): + from rag.llm import OcrModel + + assert "Mistral OCR" in OcrModel + + +def test_mistral_ocr_model_reads_flat_env_config(): + from rag.llm.ocr_model import MistralOcrModel + import json + + key = json.dumps({"MISTRAL_OCR_API_KEY": "sk-x", "MISTRAL_OCR_BASE_URL": "https://api.mistral.ai/v1"}) + mdl = MistralOcrModel(key=key, model_name="mistral-ocr-latest") + assert mdl.api_key == "sk-x" + assert mdl.model == "mistral-ocr-latest" + + +def test_mistral_ocr_model_reads_raw_secret_key(): + from rag.llm.ocr_model import MistralOcrModel + + mdl = MistralOcrModel(key="sk-raw", model_name="mistral-ocr-latest") + assert mdl.api_key == "sk-raw" + + +def test_mistral_ocr_model_preserves_flat_string_api_key(): + # regression: a flat config whose "api_key" is a plain string (not a nested + # config object) must keep the key rather than discard it. + from rag.llm.ocr_model import MistralOcrModel + import json + + key = json.dumps({"api_key": "sk-flat", "mistral_ocr_base_url": "https://api.mistral.ai/v1"}) + mdl = MistralOcrModel(key=key, model_name="mistral-ocr-latest") + assert mdl.api_key == "sk-flat" + assert mdl.base_url.rstrip("/").endswith("api.mistral.ai/v1") diff --git a/test/unit_test/deepdoc/parser/test_mistral_parser.py b/test/unit_test/deepdoc/parser/test_mistral_parser.py new file mode 100644 index 0000000000..78840f58e6 --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_mistral_parser.py @@ -0,0 +1,567 @@ +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + +import pytest + + +def _load_mistral_parser(monkeypatch): + """Load mistral_parser.py directly, bypassing deepdoc/__init__.py's + beartype_this_package() and the heavy deepdoc dependency chain. + Mirrors test_somark_parser.py.""" + repo_root = Path(__file__).resolve().parents[4] + + deepdoc_mod = ModuleType("deepdoc") + deepdoc_mod.__path__ = [str(repo_root / "deepdoc")] + monkeypatch.setitem(sys.modules, "deepdoc", deepdoc_mod) + + parser_mod = ModuleType("deepdoc.parser") + parser_mod.__path__ = [str(repo_root / "deepdoc" / "parser")] + monkeypatch.setitem(sys.modules, "deepdoc.parser", parser_mod) + + pdf_parser_mod = ModuleType("deepdoc.parser.pdf_parser") + + class _RAGFlowPdfParser: + pass + + pdf_parser_mod.RAGFlowPdfParser = _RAGFlowPdfParser + pdf_parser_mod.MAXIMUM_PAGE_NUMBER = 100000 + monkeypatch.setitem(sys.modules, "deepdoc.parser.pdf_parser", pdf_parser_mod) + + utils_mod = ModuleType("deepdoc.parser.utils") + utils_mod.extract_pdf_outlines = lambda *_a, **_k: [] + monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", utils_mod) + + module_name = "test_mistral_parser_unit_module" + module_path = repo_root / "deepdoc" / "parser" / "mistral_parser.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _make_parser(m, api_key="", **kwargs): + """Build a MistralParser without any network call. __init__ only sets attributes.""" + return m.MistralParser(base_url="https://api.mistral.ai/v1", api_key=api_key, **kwargs) + + +def test_resolve_internal_type_maps_known_types(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p._resolve_internal_type("text") == "text" + assert p._resolve_internal_type("title") == "text" + assert p._resolve_internal_type("list") == "text" + assert p._resolve_internal_type("image") == "image" + assert p._resolve_internal_type("table") == "table" + + +def test_resolve_internal_type_drops_header_footer_by_default(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p._resolve_internal_type("header") is None + assert p._resolve_internal_type("footer") is None + + +def test_resolve_internal_type_keeps_header_footer_when_flagged(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, keep_header_footer=True) + assert p._resolve_internal_type("header") == "text" + assert p._resolve_internal_type("footer") == "text" + + +def test_resolve_internal_type_unknown_falls_back_to_text(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p._resolve_internal_type("brand_new_type") == "text" + + +def test_resolve_internal_type_none_falls_back_to_text(monkeypatch): + # Regression: a block with missing/null "type" yields None here; the + # annotation must accept it (beartype runs in production) and fall back + # to text instead of aborting the whole parse. + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p._resolve_internal_type(None) == "text" + + +def test_block_text_image_returns_empty(monkeypatch): + m = _load_mistral_parser(monkeypatch) + assert m.MistralParser._block_text({"content": "x"}, "image") == "" + + +def test_block_text_strips_and_returns_content(monkeypatch): + m = _load_mistral_parser(monkeypatch) + assert m.MistralParser._block_text({"content": " hi "}, "text") == "hi" + + +def _ocr_response(): + """A real-shaped /v1/ocr response (from a live probe): two pages, a header + (dropped), text, title, an image with bbox, and a table block. dimensions + differ per page to guard against a constant rescale factor.""" + return { + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 2, "doc_size_bytes": 211291}, + "pages": [ + { + "index": 0, + "dimensions": {"dpi": 87, "width": 720, "height": 1018}, + "markdown": "# Title\n\nhello", + "images": [{"id": "img-0.jpeg", "top_left_x": 251, "top_left_y": 72, "bottom_right_x": 311, "bottom_right_y": 145}], + "tables": [], + "blocks": [ + {"type": "header", "content": "Cofinanziato", "top_left_x": 135, "top_left_y": 72, "bottom_right_x": 213, "bottom_right_y": 145}, + {"type": "title", "content": "Title", "top_left_x": 40, "top_left_y": 160, "bottom_right_x": 400, "bottom_right_y": 190}, + {"type": "text", "content": "hello world", "top_left_x": 40, "top_left_y": 200, "bottom_right_x": 500, "bottom_right_y": 230}, + {"type": "image", "content": "", "top_left_x": 251, "top_left_y": 72, "bottom_right_x": 311, "bottom_right_y": 145}, + ], + }, + { + "index": 1, + "dimensions": {"dpi": 144, "width": 1021, "height": 681}, + "markdown": "|a|b|", + "images": [], + "tables": [{"id": "tbl-0.md", "content": "
a
", "format": "html", "word_confidence_scores": None}], + "blocks": [ + {"type": "table", "content": "
a
", "table_id": "tbl-0.md", "top_left_x": 49, "top_left_y": 103, "bottom_right_x": 960, "bottom_right_y": 597}, + ], + }, + ], + } + + +def test_normalize_pages_maps_bbox_and_page_size(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + pages = p._normalize_pages(_ocr_response()) + assert [pg["page_num"] for pages_ in [pages] for pg in pages_] == [0, 1] + assert pages[0]["page_size"] == {"w": 720, "h": 1018} + assert pages[1]["page_size"] == {"w": 1021, "h": 681} + first_text = pages[0]["blocks"][2] + assert first_text["bbox"] == [40, 200, 500, 230] # [x0, top, x1, bott] + + +def test_normalize_pages_bbox_none_when_no_coords(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + resp = {"pages": [{"index": 0, "dimensions": {"width": 100, "height": 100}, "blocks": [{"type": "image", "content": ""}]}]} + pages = p._normalize_pages(resp) + assert pages[0]["blocks"][0]["bbox"] is None + # geometry-less image must be skipped, not emitted as a zero-area crop + assert p._transfer_to_sections(pages) == [] + + +def test_normalize_pages_null_index_coerced_to_zero(monkeypatch): + # Regression: a present-but-null "index" must coerce to 0, not None + # (dict.get's default only applies when the key is absent), so the + # downstream page_idx + 1 in _line_tag does not raise TypeError. + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + resp = { + "pages": [ + { + "index": None, + "dimensions": {"width": 100, "height": 100}, + "blocks": [{"type": "text", "content": "hi", "top_left_x": 1, "top_left_y": 2, "bottom_right_x": 3, "bottom_right_y": 4}], + } + ] + } + pages = p._normalize_pages(resp) + assert pages[0]["page_num"] == 0 + secs = p._transfer_to_sections(pages) # must not raise + # naive path: text block -> (text, line_tag); tag stamped page 1 (1-based) + assert any(s[0] == "hi" and s[1].startswith("@@1\t") for s in secs) + + +def test_transfer_naive_path_returns_2_tuples_without_header(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + pages = p._normalize_pages(_ocr_response()) + secs = p._transfer_to_sections(pages) # parse_method None -> naive + assert all(isinstance(s, tuple) and len(s) == 2 for s in secs) + joined = " ".join(s[0] for s in secs) + assert "Cofinanziato" not in joined # header dropped + assert "hello world" in joined + assert "" in joined # table inlined as text + + +def test_transfer_pipeline_path_returns_typed_3_tuples(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + pages = p._normalize_pages(_ocr_response()) + secs = p._transfer_to_sections(pages, parse_method="pipeline") + assert all(isinstance(s, tuple) and len(s) == 3 for s in secs) + layout_types = {s[1] for s in secs} + assert {"text", "image", "table"} <= layout_types + + +def test_transfer_naive_image_carries_caption_and_tag(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + pages = p._normalize_pages(_ocr_response()) + secs = p._transfer_to_sections(pages) + img = [s for s in secs if "@@" in s[0] and "##" in s[0]] + assert len(img) == 1 # exactly the image block, tag embedded in text + + +def test_line_tag_rescales_per_page_dimensions(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + + class _Img: + def __init__(self, size): + self.size = size + + # Distinct scale factors per page and per axis so a swapped-axis or + # wrong-page-index bug cannot pass: page 0 scales x by 3.0, y by 4.0; + # page 1 scales x by 4.0, y by 6.0. + p.page_images = [_Img((300, 800)), _Img((200, 600))] + + tag0 = p._line_tag({"page_idx": 0, "bbox": [10, 20, 40, 60], "page_size": {"w": 100, "h": 200}}) + assert tag0 == "@@1\t30.0\t120.0\t80.0\t240.0##" + + tag1 = p._line_tag({"page_idx": 1, "bbox": [5, 10, 25, 40], "page_size": {"w": 50, "h": 100}}) + assert tag1 == "@@2\t20.0\t100.0\t60.0\t240.0##" + + +def test_line_tag_clamps_negative_coords_and_roundtrips(monkeypatch): + # Regression: a bbox bleeding past the page origin rescales to negative + # coords; extract_positions' regex excludes '-', so an unclamped tag is + # silently dropped (crop returns None). Clamp to >= 0 so the tag survives. + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + + class _Img: + def __init__(self, size): + self.size = size + + p.page_images = [_Img((300, 800))] + tag = p._line_tag({"page_idx": 0, "bbox": [-10, -20, 40, 60], "page_size": {"w": 100, "h": 200}}) + assert tag == "@@1\t0.0\t120.0\t0.0\t240.0##" # negatives clamped to 0.0 + poss = m.MistralParser.extract_positions(tag) + assert len(poss) == 1 # the regex matches; block is not dropped + _pn, left, _right, top, _bottom = poss[0] + assert left == 0.0 and top == 0.0 + + +def test_transfer_to_tables_is_empty(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + pages = p._normalize_pages(_ocr_response()) + assert p._transfer_to_tables(pages) == [] + + +def test_crop_returns_none_without_positions(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p.crop("plain text with no tag") is None + + +def test_crop_reads_page_images_for_tagged_text(monkeypatch): + from PIL import Image + + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + p.page_images = [Image.new("RGB", (200, 300), "white")] + tag = "@@1\t10.0\t100.0\t20.0\t60.0##" + out = p.crop("caption" + tag, need_position=True) + assert isinstance(out, tuple) and len(out) == 2 + + +class _Resp: + def __init__(self, status, payload=None, text=""): + self.status_code = status + self._payload = payload or {} + self.text = text + + def json(self): + return self._payload + + +def test_check_installation_requires_api_key(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="") + ok, reason = p.check_installation() + assert ok is False and "key" in reason.lower() + + +def test_call_ocr_inline_posts_pages_selector(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + captured = {} + + def fake_post(url, headers=None, json=None, timeout=None, **kw): + captured["url"] = url + captured["json"] = json + return _Resp(200, {"pages": [{"index": 5, "markdown": "x", "blocks": []}], "usage_info": {"pages_processed": 1}}) + + monkeypatch.setattr(m.requests, "post", fake_post) + out = p._call_ocr(b"%PDF-1.4 fake", "f.pdf", pages=[5]) + assert captured["url"].endswith("/ocr") + assert captured["json"]["pages"] == [5] + assert captured["json"]["include_blocks"] is True + assert out["pages"][0]["index"] == 5 + + +def test_call_ocr_omits_pages_when_none(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + captured = {} + monkeypatch.setattr(m.requests, "post", lambda url, headers=None, json=None, timeout=None, **kw: captured.update(json=json) or _Resp(200, {"pages": []})) + p._call_ocr(b"%PDF fake", "f.pdf", pages=None) + assert "pages" not in captured["json"] + + +def test_call_ocr_raises_on_http_error(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-bad") + monkeypatch.setattr(m.requests, "post", lambda *a, **k: _Resp(401, text="Unauthorized")) + with pytest.raises(RuntimeError, match="401"): + p._call_ocr(b"%PDF fake", "f.pdf", pages=None) + + +def test_call_ocr_uploads_when_over_inline_limit(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test", inline_max_bytes=4) + calls = {"post": [], "get": [], "delete": []} + + def fake_post(url, headers=None, json=None, data=None, files=None, timeout=None, **kw): + calls["post"].append(url) + if url.endswith("/files"): + return _Resp(200, {"id": "file-1"}) + return _Resp(200, {"pages": [{"index": 0, "blocks": []}]}) + + monkeypatch.setattr(m.requests, "post", fake_post) + monkeypatch.setattr(m.requests, "get", lambda url, headers=None, params=None, timeout=None, **kw: calls["get"].append(url) or _Resp(200, {"url": "https://signed"})) + monkeypatch.setattr(m.requests, "delete", lambda url, headers=None, timeout=None, **kw: calls["delete"].append(url) or _Resp(200, {})) + + out = p._call_ocr(b"%PDF-too-big", "big.pdf", pages=None) + assert any(u.endswith("/files") for u in calls["post"]) + assert calls["get"] and calls["delete"] # signed-url fetch + cleanup + assert out["pages"][0]["index"] == 0 + + +def test_call_ocr_upload_files_error_no_delete(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test", inline_max_bytes=4) + calls = {"delete": []} + monkeypatch.setattr(m.requests, "post", lambda url, headers=None, json=None, data=None, files=None, timeout=None, **kw: _Resp(500, text="boom")) + monkeypatch.setattr(m.requests, "delete", lambda url, headers=None, timeout=None, **kw: calls["delete"].append(url) or _Resp(200, {})) + with pytest.raises(RuntimeError, match="500"): + p._call_ocr(b"%PDF-too-big", "big.pdf", pages=None) + assert calls["delete"] == [] # file_id never set -> nothing to clean up + + +def test_call_ocr_signed_url_error_triggers_cleanup(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test", inline_max_bytes=4) + calls = {"delete": []} + monkeypatch.setattr(m.requests, "post", lambda url, headers=None, json=None, data=None, files=None, timeout=None, **kw: _Resp(200, {"id": "file-1"})) + monkeypatch.setattr(m.requests, "get", lambda url, headers=None, params=None, timeout=None, **kw: _Resp(403, text="nope")) + monkeypatch.setattr(m.requests, "delete", lambda url, headers=None, timeout=None, **kw: calls["delete"].append(url) or _Resp(200, {})) + with pytest.raises(RuntimeError, match="403"): + p._call_ocr(b"%PDF-too-big", "big.pdf", pages=None) + assert any("file-1" in u for u in calls["delete"]) # cleanup ran + + +def test_call_ocr_post_upload_ocr_error_triggers_cleanup(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test", inline_max_bytes=4) + calls = {"delete": []} + + def fake_post(url, headers=None, json=None, data=None, files=None, timeout=None, **kw): + if url.endswith("/files"): + return _Resp(200, {"id": "file-9"}) + return _Resp(422, text="bad ocr") + + monkeypatch.setattr(m.requests, "post", fake_post) + monkeypatch.setattr(m.requests, "get", lambda url, headers=None, params=None, timeout=None, **kw: _Resp(200, {"url": "https://signed"})) + monkeypatch.setattr(m.requests, "delete", lambda url, headers=None, timeout=None, **kw: calls["delete"].append(url) or _Resp(200, {})) + with pytest.raises(RuntimeError, match="422"): + p._call_ocr(b"%PDF-too-big", "big.pdf", pages=None) + assert any("file-9" in u for u in calls["delete"]) # finally cleanup ran despite OCR failure + + +def test_call_ocr_delete_failure_does_not_mask_error(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test", inline_max_bytes=4) + + def fake_post(url, headers=None, json=None, data=None, files=None, timeout=None, **kw): + if url.endswith("/files"): + return _Resp(200, {"id": "file-x"}) + return _Resp(422, text="bad ocr") + + def boom_delete(url, headers=None, timeout=None, **kw): + raise RuntimeError("delete network error") + + monkeypatch.setattr(m.requests, "post", fake_post) + monkeypatch.setattr(m.requests, "get", lambda url, headers=None, params=None, timeout=None, **kw: _Resp(200, {"url": "https://signed"})) + monkeypatch.setattr(m.requests, "delete", boom_delete) + # the real OCR error, not the swallowed delete error + with pytest.raises(RuntimeError, match="422"): + p._call_ocr(b"%PDF-too-big", "big.pdf", pages=None) + + +def _patch_render(m, p, n_pages): + from PIL import Image + + p.page_images = [Image.new("RGB", (100, 140), "white") for _ in range(n_pages)] + # neuter __images__ so parse_pdf doesn't touch pdfplumber + p.__images__ = lambda *a, **k: None + + +def test_parse_pdf_whole_document_omits_pages(monkeypatch, tmp_path): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + seen = {} + p._call_ocr = lambda pdf_bytes, filename, pages, callback=None: seen.update(pages=pages) or _ocr_response() + _patch_render(m, p, 2) + pdf = tmp_path / "x.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + secs, tables = p.parse_pdf(str(pdf)) + assert seen["pages"] is None # whole doc -> no selector + assert tables == [] + assert any("hello world" in s[0] for s in secs) + + +def test_parse_pdf_restricted_range_sends_absolute_pages(monkeypatch, tmp_path): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + seen = {} + p._call_ocr = lambda pdf_bytes, filename, pages, callback=None: seen.update(pages=pages) or {"pages": []} + _patch_render(m, p, 10) + pdf = tmp_path / "x.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + p.parse_pdf(str(pdf), from_page=4, to_page=8) + assert seen["pages"] == [4, 5, 6, 7] # 0-based, half-open, absolute + + +def test_parse_pdf_pipeline_returns_3_tuples(monkeypatch, tmp_path): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + p._call_ocr = lambda *a, **k: _ocr_response() + _patch_render(m, p, 2) + pdf = tmp_path / "x.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + secs, _ = p.parse_pdf(str(pdf), parse_method="pipeline") + assert all(len(s) == 3 for s in secs) + + +def test_parse_pdf_empty_range_short_circuits(monkeypatch, tmp_path): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + called = {"ocr": 0} + + def fake_ocr(pdf_bytes, filename, pages, callback=None): + called["ocr"] += 1 + return {"pages": []} + + p._call_ocr = fake_ocr + _patch_render(m, p, 3) # only 3 pages rendered + pdf = tmp_path / "x.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + # from_page beyond the rendered page count -> empty selector after clamp + secs, tables = p.parse_pdf(str(pdf), from_page=5, to_page=10) + assert (secs, tables) == ([], []) + assert called["ocr"] == 0 # short-circuited, no OCR call made + + +def test_parse_pdf_binary_bytes_path(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + p._call_ocr = lambda pdf_bytes, filename, pages, callback=None: _ocr_response() + _patch_render(m, p, 2) + secs, tables = p.parse_pdf("x.pdf", binary=b"%PDF-1.4 minimal") + assert any("hello world" in s[0] for s in secs) + assert tables == [] + + +def test_parse_pdf_binary_stream_normalized(monkeypatch): + from io import BytesIO + + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + seen = {} + p._call_ocr = lambda pdf_bytes, filename, pages, callback=None: seen.update(pdf_bytes=pdf_bytes) or _ocr_response() + _patch_render(m, p, 2) + p.parse_pdf("x.pdf", binary=BytesIO(b"%PDF-1.4 stream")) + assert seen["pdf_bytes"] == b"%PDF-1.4 stream" # BytesIO normalized to raw bytes + + +def test_image_description_injected_when_vision_model_present(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + p.vision_model = object() # truthy -> enrichment path taken + p._describe_image = lambda line_tag: "A scatter plot of X versus Y" + pages = p._normalize_pages(_ocr_response()) + secs = p._transfer_to_sections(pages) + img = [s for s in secs if "@@" in s[0] and "scatter plot" in s[0]] + assert len(img) == 1 # VLM caption injected into the image chunk text + + +def test_image_description_skipped_without_vision_model(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + assert p.vision_model is None # default: no enrichment + calls = {"n": 0} + + def _boom(line_tag): + calls["n"] += 1 + return "SHOULD NOT APPEAR" + + p._describe_image = _boom + pages = p._normalize_pages(_ocr_response()) + secs = p._transfer_to_sections(pages) + assert calls["n"] == 0 # _describe_image not called when vision_model is None + assert not any("SHOULD NOT APPEAR" in s[0] for s in secs) + # the image chunk still exists, labelled by fallback (empty caption -> "image N") + assert any("@@" in s[0] and "image 1" in s[0] for s in secs) + + +def test_describe_image_returns_empty_when_crop_none(monkeypatch): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + p.vision_model = object() + p.crop = lambda *a, **k: None # no image to crop + assert p._describe_image("@@1\t0\t0\t0\t0##") == "" + + +def test_parse_pdf_consumes_vision_model_kwarg(monkeypatch, tmp_path): + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m, api_key="sk-test") + seen = {} + p._call_ocr = lambda pdf_bytes, filename, pages, callback=None: seen.setdefault("pages", pages) or _ocr_response() + _patch_render(m, p, 2) + pdf = tmp_path / "x.pdf" + pdf.write_bytes(b"%PDF-1.4 minimal") + p.parse_pdf(str(pdf), vision_model="VM") + assert p.vision_model == "VM" # popped from kwargs into self, not forwarded to _call_ocr + + +def test_describe_image_passes_pil_image_not_bytes(monkeypatch): + # Regression: vision_llm_chunk expects a PIL Image (it calls img.size/img.save), + # so _describe_image must pass the crop directly, not its bytes. + import sys as _sys + from types import ModuleType + from PIL import Image + + m = _load_mistral_parser(monkeypatch) + p = _make_parser(m) + p.vision_model = object() + p.crop = lambda *a, **k: Image.new("RGB", (40, 40), "white") + + captured = {} + for name in ("rag", "rag.app", "rag.prompts"): + if name not in _sys.modules: + monkeypatch.setitem(_sys.modules, name, ModuleType(name)) + pic = ModuleType("rag.app.picture") + pic.vision_llm_chunk = lambda binary, vision_model, prompt=None, callback=None: (captured.update(kind=type(binary).__name__), "a white square")[1] + gen = ModuleType("rag.prompts.generator") + gen.vision_llm_figure_describe_prompt = lambda: "describe" + monkeypatch.setitem(_sys.modules, "rag.app.picture", pic) + monkeypatch.setitem(_sys.modules, "rag.prompts.generator", gen) + + out = p._describe_image("@@1\t0\t0\t40\t40##") + assert out == "a white square" + assert captured["kind"] == "Image" # PIL Image, not 'bytes' diff --git a/test/unit_test/llm/test_mistral_cv.py b/test/unit_test/llm/test_mistral_cv.py new file mode 100644 index 0000000000..9a83d30f89 --- /dev/null +++ b/test/unit_test/llm/test_mistral_cv.py @@ -0,0 +1,12 @@ +def test_mistral_cv_registers_in_registry(): + from rag.llm import CvModel + + assert "Mistral" in CvModel, sorted(CvModel.keys()) + + +def test_mistral_cv_defaults_to_pixtral_and_mistral_base_url(): + from rag.llm.cv_model import MistralCV + + cv = MistralCV(key="sk-test") + assert cv.model_name == "pixtral-12b-2409" + assert cv.base_url.rstrip("/").endswith("api.mistral.ai/v1")