diff --git a/DOCKER.md b/DOCKER.md index c3190898..18671aba 100644 --- a/DOCKER.md +++ b/DOCKER.md @@ -118,6 +118,30 @@ COMPOSE_PROJECT_NAME=morphik # Set once before the first start `docker-compose.run.yml` mounts `./morphik.toml` read-only into both the API and worker containers. Edit that file and run `./start-morphik.sh` again to apply changes. +### 5. Text in images + +By default, Docling reads text in uploaded images with RapidOCR on the worker's CPU. RapidOCR downloads its model files +from www.modelscope.cn the first time it runs, and again whenever the worker container is recreated. + +To send images to a vision model instead, point `image_ocr_model` at a key in `[registered_models]`: + +```toml +[registered_models] +openai_gpt4-1-mini = { model_name = "gpt-4.1-mini" } + +[parser] +image_ocr_model = "openai_gpt4-1-mini" +``` + +PNG, JPEG, GIF, WebP, BMP and TIFF files then go to that model, which returns the text as Markdown, or a one or two +sentence description when an image has little text. Images are converted to PNG and scaled to at most 2048 px on the +longest side before sending. PDFs and Office files still use Docling. Rate limits, timeouts and provider outages are +retried like other transient ingestion errors; an invalid key fails the document. With `use_colpali=true`, images go to +ColPali and skip this step. + +The setting only applies to new ingests. To re-read images already indexed with local OCR, requeue them with +`POST /ingest/requeue`. + ## Accessing services - Morphik API: http://localhost:8000 diff --git a/core/config.py b/core/config.py index e927f5a4..da3a483e 100644 --- a/core/config.py +++ b/core/config.py @@ -86,6 +86,7 @@ class Settings(BaseSettings): CHUNK_SIZE: int CHUNK_OVERLAP: int USE_CONTEXTUAL_CHUNKING: bool = False + IMAGE_OCR_MODEL: Optional[str] = None PARSER_XML: ParserXMLSettings = ParserXMLSettings() # Reranker configuration @@ -274,6 +275,10 @@ def get_settings() -> Settings: "USE_CONTEXTUAL_CHUNKING": config["parser"].get("use_contextual_chunking", False), } ) + image_ocr_model = config["parser"].get("image_ocr_model") or None + if image_ocr_model and image_ocr_model not in settings_dict.get("REGISTERED_MODELS", {}): + raise ValueError(f"[parser] image_ocr_model '{image_ocr_model}' is not a key in [registered_models]") + settings_dict["IMAGE_OCR_MODEL"] = image_ocr_model # Load parser XML config if "xml" in config["parser"]: diff --git a/core/parser/morphik_parser.py b/core/parser/morphik_parser.py index 7e781f7e..315d272b 100644 --- a/core/parser/morphik_parser.py +++ b/core/parser/morphik_parser.py @@ -1,3 +1,4 @@ +import base64 import io import logging import os @@ -8,11 +9,13 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple +import litellm import openpyxl from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import PdfPipelineOptions from docling.document_converter import DocumentConverter, PdfFormatOption from httpx import AsyncClient, Timeout +from PIL import Image from core.config import get_settings from core.models.chunk import Chunk @@ -26,6 +29,17 @@ logger = logging.getLogger(__name__) +# Raster formats a vision model can read once converted to PNG. SVG is vector and is left to Docling. +OCR_IMAGE_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp", "image/bmp", "image/tiff"} +# OpenAI scales high-detail images to fit 2048 px anyway; downscaling first only saves upload bytes. +OCR_MAX_SIDE = 2048 +OCR_PROMPT = ( + "Transcribe all text in this image exactly as written, in reading order. Keep the structure with " + "Markdown: headings, lists, and tables for tabular content such as forms and grids. Do not translate, " + "summarize, or correct anything. If the image has little or no text, describe what it shows in one or " + "two sentences instead. Reply with the transcription only." +) + class BaseChunker(ABC): """Base class for text chunking strategies""" @@ -207,6 +221,7 @@ def __init__( anthropic_api_key: Optional[str] = None, frame_sample_rate: int = 1, use_contextual_chunking: bool = False, + image_ocr_model: Optional[str] = None, ): # Initialize basic configuration self._assemblyai_api_key = assemblyai_api_key @@ -225,6 +240,16 @@ def __init__( # Initialize logger self.logger = logging.getLogger(__name__) + # Images go to this registered vision model instead of Docling's OCR when set + self._image_ocr_model = image_ocr_model + self._image_ocr_config: Dict[str, Any] = {} + if image_ocr_model: + registered = getattr(self.settings, "REGISTERED_MODELS", {}) or {} + if image_ocr_model not in registered: + raise ValueError(f"image_ocr_model '{image_ocr_model}' is not a key in registered_models") + self._image_ocr_config = registered[image_ocr_model] + self.logger.info(f"Image OCR uses vision model {image_ocr_model}") + # Setup for API mode parsing self._parse_api_endpoints: Optional[List[str]] = None self._parse_api_key: Optional[str] = None @@ -262,6 +287,60 @@ def _is_video_file(self, file: bytes, filename: str) -> bool: logging.error(f"Error detecting file type: {str(e)}") return filename.lower().endswith(".mp4") + def _is_ocr_image(self, file: bytes, filename: str) -> bool: + """True when images should go to the vision model and this file is a raster image.""" + if not self._image_ocr_model: + return False + try: + return detect_content_type(content=file, filename=filename) in OCR_IMAGE_TYPES + except Exception: + return False + + @staticmethod + def _image_to_png(file: bytes) -> bytes: + """First frame as RGB(A) PNG, longest side at most OCR_MAX_SIDE. Covers formats the API rejects (BMP, TIFF).""" + with Image.open(io.BytesIO(file)) as image: + image.seek(0) + frame = image.convert("RGBA" if "A" in image.getbands() else "RGB") + if max(frame.size) > OCR_MAX_SIDE: + frame.thumbnail((OCR_MAX_SIDE, OCR_MAX_SIDE), Image.LANCZOS) + output = io.BytesIO() + frame.save(output, format="PNG") + return output.getvalue() + + async def _transcribe_image(self, file: bytes, filename: str) -> str: + """Read an image's text with the configured vision model.""" + image_url = "data:image/png;base64," + base64.b64encode(self._image_to_png(file)).decode() + params: Dict[str, Any] = { + "model": self._image_ocr_config["model_name"], + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": OCR_PROMPT}, + {"type": "image_url", "image_url": {"url": image_url, "detail": "high"}}, + ], + } + ], + "max_tokens": 4096, + } + params.update({k: v for k, v in self._image_ocr_config.items() if k not in ("model_name", "vision")}) + try: + response = await litellm.acompletion(**params) + except ( + litellm.RateLimitError, + litellm.Timeout, + litellm.APIConnectionError, + litellm.InternalServerError, + litellm.ServiceUnavailableError, + ) as e: + # A later retry can fix these. ConnectionError makes the ingestion worker retry the job. + raise ConnectionError(f"Image OCR call to {self._image_ocr_model} failed: {e}") from e + text = response.choices[0].message.content or "" + if not text.strip(): + self.logger.warning(f"Image OCR returned no text for {filename}") + return text + def _is_xml_file(self, filename: str, content_type: Optional[str] = None) -> bool: """Check if the file is an XML file.""" if filename and filename.lower().endswith(".xml"): @@ -601,6 +680,9 @@ async def _parse_document(self, file: bytes, filename: str) -> Tuple[Dict[str, A except Exception as e: self.logger.warning(f"Fast Excel parser failed for {filename}: {e}, falling back to Docling") + if self._is_ocr_image(file, filename): + return {}, await self._transcribe_image(file, filename) + # For complex formats, use API if configured, otherwise local Docling if self._parse_api_endpoints: try: @@ -629,6 +711,9 @@ async def parse_file_to_text_deep(self, file: bytes, filename: str) -> Tuple[Dic """Run an expensive parse fallback after normal parsing produced no usable chunks.""" if self._is_plain_text_file(filename): return await self.parse_file_to_text(file, filename) + if self._is_ocr_image(file, filename): + # The vision model already read this image. Docling would only rerun local OCR. + return {"parse_fallback": "deep"}, "" parse_file = file parse_filename = filename diff --git a/core/services_init.py b/core/services_init.py index 66bd2d98..b74a9b96 100644 --- a/core/services_init.py +++ b/core/services_init.py @@ -89,6 +89,7 @@ assemblyai_api_key=settings.ASSEMBLYAI_API_KEY, anthropic_api_key=settings.ANTHROPIC_API_KEY, use_contextual_chunking=settings.USE_CONTEXTUAL_CHUNKING, + image_ocr_model=settings.IMAGE_OCR_MODEL, ) embedding_model = LiteLLMEmbeddingModel(model_key=settings.EMBEDDING_MODEL) diff --git a/core/tests/unit/test_image_ocr_model.py b/core/tests/unit/test_image_ocr_model.py new file mode 100644 index 00000000..891b621f --- /dev/null +++ b/core/tests/unit/test_image_ocr_model.py @@ -0,0 +1,158 @@ +"""[parser] image_ocr_model sends raster images to a vision model instead of Docling's local OCR.""" + +import base64 +import io +from types import SimpleNamespace + +import litellm +import pytest +from PIL import Image + +from core.parser import morphik_parser +from core.parser.morphik_parser import OCR_MAX_SIDE, MorphikParser + +pytestmark = pytest.mark.unit + +REGISTERED = { + "vision_ocr": {"model_name": "gpt-4.1-mini", "vision": True, "api_base": "https://llm.internal/v1"}, +} + + +def _image(fmt: str, size=(40, 20)) -> bytes: + output = io.BytesIO() + Image.new("RGB", size, "white").save(output, format=fmt) + return output.getvalue() + + +def _sent_image(call) -> Image.Image: + url = call["messages"][0]["content"][1]["image_url"]["url"] + assert url.startswith("data:image/png;base64,") + return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))) + + +@pytest.fixture +def make_parser(monkeypatch): + monkeypatch.setattr( + morphik_parser, "get_settings", lambda: SimpleNamespace(REGISTERED_MODELS=REGISTERED, PARSER_MODE="local") + ) + + def make(image_ocr_model="vision_ocr"): + parser = MorphikParser(image_ocr_model=image_ocr_model) + + async def docling(file, filename): + docling_calls.append(filename) + return "docling text" + + parser._parse_document_local = docling + parser._parse_document_local_deep = docling + return parser + + docling_calls: list = [] + make.docling_calls = docling_calls + return make + + +@pytest.fixture +def completions(monkeypatch): + calls = [] + + async def fake_acompletion(**params): + calls.append(params) + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="Order #4521\nStatus: Open"))]) + + monkeypatch.setattr(morphik_parser.litellm, "acompletion", fake_acompletion) + return calls + + +@pytest.mark.asyncio +async def test_png_is_read_by_the_vision_model(make_parser, completions): + parser = make_parser() + + metadata, text = await parser.parse_file_to_text(_image("PNG"), "image.png") + + assert (metadata, text) == ({}, "Order #4521\nStatus: Open") + assert make_parser.docling_calls == [] + (call,) = completions + assert call["model"] == "gpt-4.1-mini" + assert call["api_base"] == "https://llm.internal/v1" + assert "vision" not in call + assert _sent_image(call).size == (40, 20) + + +@pytest.mark.asyncio +async def test_formats_the_api_rejects_are_sent_as_png(make_parser, completions): + parser = make_parser() + + await parser.parse_file_to_text(_image("BMP"), "scan.bmp") + await parser.parse_file_to_text(_image("TIFF"), "scan.tiff") + + assert [_sent_image(call).format for call in completions] == ["PNG", "PNG"] + + +@pytest.mark.asyncio +async def test_large_images_are_downscaled(make_parser, completions): + parser = make_parser() + + await parser.parse_file_to_text(_image("JPEG", size=(4096, 1024)), "wide.jpg") + + assert _sent_image(completions[0]).size == (OCR_MAX_SIDE, OCR_MAX_SIDE // 4) + + +@pytest.mark.asyncio +async def test_without_the_setting_images_use_docling(make_parser, completions): + parser = make_parser(image_ocr_model=None) + + _, text = await parser.parse_file_to_text(_image("PNG"), "image.png") + + assert text == "docling text" + assert completions == [] + + +@pytest.mark.asyncio +async def test_other_files_still_use_docling(make_parser, completions): + parser = make_parser() + + _, text = await parser.parse_file_to_text(b"%PDF-1.4\n%fake", "spec.pdf") + + assert text == "docling text" + assert completions == [] + + +@pytest.mark.asyncio +async def test_deep_fallback_does_not_rerun_local_ocr_on_images(make_parser, completions): + parser = make_parser() + + metadata, text = await parser.parse_file_to_text_deep(_image("PNG"), "image.png") + + assert (metadata, text) == ({"parse_fallback": "deep"}, "") + assert make_parser.docling_calls == [] + + +@pytest.mark.asyncio +async def test_rate_limits_are_retryable(make_parser, monkeypatch): + async def rate_limited(**params): + raise litellm.RateLimitError("slow down", llm_provider="openai", model="gpt-4.1-mini") + + monkeypatch.setattr(morphik_parser.litellm, "acompletion", rate_limited) + parser = make_parser() + + # The ingestion worker retries ConnectionError instead of failing the document. + with pytest.raises(ConnectionError): + await parser.parse_file_to_text(_image("PNG"), "image.png") + + +@pytest.mark.asyncio +async def test_auth_errors_fail_without_retry(make_parser, monkeypatch): + async def unauthorized(**params): + raise litellm.AuthenticationError("bad key", llm_provider="openai", model="gpt-4.1-mini") + + monkeypatch.setattr(morphik_parser.litellm, "acompletion", unauthorized) + parser = make_parser() + + with pytest.raises(litellm.AuthenticationError): + await parser.parse_file_to_text(_image("PNG"), "image.png") + + +def test_unknown_model_key_is_rejected(make_parser): + with pytest.raises(ValueError, match="registered_models"): + make_parser(image_ocr_model="missing") diff --git a/core/workers/ingestion_worker.py b/core/workers/ingestion_worker.py index fcc9144b..0ab6d7e6 100644 --- a/core/workers/ingestion_worker.py +++ b/core/workers/ingestion_worker.py @@ -1963,6 +1963,7 @@ async def startup(ctx): assemblyai_api_key=settings.ASSEMBLYAI_API_KEY, anthropic_api_key=settings.ANTHROPIC_API_KEY, use_contextual_chunking=settings.USE_CONTEXTUAL_CHUNKING, + image_ocr_model=settings.IMAGE_OCR_MODEL, ) ctx["parser"] = parser diff --git a/morphik.docker.toml b/morphik.docker.toml index 83afabd5..8030ea33 100644 --- a/morphik.docker.toml +++ b/morphik.docker.toml @@ -85,6 +85,9 @@ chunk_size = 6000 chunk_overlap = 300 use_contextual_chunking = false contextual_chunking_model = "openai_gpt4-1-mini" # Reference to a key in registered_models +# Read text in images (PNG, JPEG, GIF, WebP, BMP, TIFF) with this registered vision model instead of +# Docling's local OCR. Leave unset to keep local OCR. +# image_ocr_model = "openai_gpt4-1-mini" [parser.xml] max_tokens = 350 diff --git a/morphik.toml b/morphik.toml index b328b974..37e4a1be 100644 --- a/morphik.toml +++ b/morphik.toml @@ -84,6 +84,9 @@ chunk_size = 6000 chunk_overlap = 300 use_contextual_chunking = false contextual_chunking_model = "ollama_qwen_vision" # Reference to a key in registered_models +# Read text in images (PNG, JPEG, GIF, WebP, BMP, TIFF) with this registered vision model instead of +# Docling's local OCR. Leave unset to keep local OCR. +# image_ocr_model = "openai_gpt4-1-mini" [parser.xml] max_tokens = 350