From fdc7f3ff3c372b024ce83ce7ddc23a13ea6c1756 Mon Sep 17 00:00:00 2001 From: sandeep-selvaraj Date: Thu, 3 Sep 2026 08:51:23 +0200 Subject: [PATCH] feat: OCR scanned PDFs with a vision model via LiteLLM Add ocr="llm" to the Python bindings, alongside ocr="hosted". When a PDF needs OCR, rasterise every page with pypdfium2 and transcribe it to Markdown through litellm.completion, so the document goes only to the provider the user configures. --- README.md | 2 + python/README.md | 20 +++++ python/anydoc/__init__.py | 55 +++++++++--- python/anydoc/_llm_ocr.py | 166 ++++++++++++++++++++++++++++++++++++ python/pyproject.toml | 5 ++ python/tests/test_anydoc.py | 75 +++++++++++++++- 6 files changed, 310 insertions(+), 13 deletions(-) create mode 100644 python/anydoc/_llm_ocr.py diff --git a/README.md b/README.md index 760d3f18..75b47bce 100644 --- a/README.md +++ b/README.md @@ -144,6 +144,8 @@ anydoc reads text-based PDFs locally but does no OCR, so a PDF with scanned or i Only documents that need OCR leave the machine, and the whole document goes, since Parse has no page selection. If Parse cannot convert it, Node rejects with `code: 'hosted'` and Python raises `HostedError`. `--api-url`, `apiUrl` and `api_url`, else `FIRECRAWL_API_URL`, point at another Parse deployment. The Rust crate has no `ocr` option and never makes network calls. +The Python package also offers `ocr="llm"`: it rasterises the scanned pages and transcribes them with a vision model of your choosing through [LiteLLM](https://docs.litellm.ai), so the document goes only to the provider you configure. It needs the `llm` extra (`pip install "firecrawl-anydoc[llm]"`) and is configured from `ANYDOC_OCR_*` environment variables (`ANYDOC_OCR_MODEL` is required). Failures raise `LlmOcrError`. See [python/README.md](python/README.md#scanned-pages). + ## Features - **One output for every format.** Each format parses into a shared document model and renders through a single Markdown serializer, so escaping, tables, heading anchors, and footnotes behave identically whether the input was a `.doc` from 2003 or a `.pptx` from yesterday. diff --git a/python/README.md b/python/README.md index 9d038bed..bfbf6d9c 100644 --- a/python/README.md +++ b/python/README.md @@ -52,6 +52,25 @@ anydoc converts locally and does not do OCR, so a PDF with scanned or image-only markdown = anydoc.to_markdown("scan.pdf", ocr="hosted") ``` +Or keep the document on your own infrastructure with `ocr="llm"`: it rasterises +the pages and transcribes them with a vision model of your choosing through +[LiteLLM](https://docs.litellm.ai). Install the extra and point it at a model: + +```bash +pip install "firecrawl-anydoc[llm]" +export ANYDOC_OCR_MODEL=openai/gpt-4o-mini # any LiteLLM model string +export OPENAI_API_KEY=... # or ANYDOC_OCR_API_KEY +``` + +```python +markdown = anydoc.to_markdown("scan.pdf", ocr="llm") +``` + +Configuration is read from `ANYDOC_OCR_*`: `MODEL` (required), `API_BASE`, +`API_KEY`, `PROMPT`, `DPI` (default 200), `MAX_PAGES` (default 100), +`PAGE_CONCURRENCY` (default 4), `TIMEOUT` (seconds, default 120). `model=` +overrides `ANYDOC_OCR_MODEL` per call. Failures raise `LlmOcrError`. + ## Errors A conversion raises only when no complete Markdown could come out of the file. The exception type names what went wrong: @@ -74,6 +93,7 @@ except (anydoc.EncryptedError, anydoc.UnsupportedError) as error: | `ResourceLimitError` | Crossed a fixed safety limit (decompression, nesting, node count) | | `MissingPartError` | A part required for any meaningful output is absent | | `HostedError` | `ocr="hosted"` could not get the document through Firecrawl Parse | +| `LlmOcrError` | `ocr="llm"` could not transcribe it (missing extra/config, or model failure) | | `OSError` | The file could not be read, from `to_markdown` only | Every conversion failure subclasses `anydoc.ConvertError`, so catching that handles all of them at once. `MalformedError.part` and `MissingPartError.part` name the package part at fault, `ResourceLimitError.limit` names the limit crossed, and `str(error)` carries the whole message. A `format` argument naming no supported format raises `ValueError`. diff --git a/python/anydoc/__init__.py b/python/anydoc/__init__.py index c1ba6560..f956909c 100644 --- a/python/anydoc/__init__.py +++ b/python/anydoc/__init__.py @@ -45,23 +45,32 @@ variants that share a parser (`.docm`, `.xlsm`, `.ppsx`, ...) map onto these via `format_from_bytes` or `format_from_extension`.""" -Ocr = Literal["reject", "hosted"] +Ocr = Literal["reject", "hosted", "llm"] """What happens to a PDF whose pages need OCR. `reject` (the default) raises `NeedsOcrError` naming the pages. `hosted` sends the whole document to -Firecrawl Parse instead, keyless unless a key is given. Documents anydoc -converts itself never leave the machine.""" +Firecrawl Parse. `llm` rasterises the pages and transcribes them with a +vision model through LiteLLM, configured from `ANYDOC_OCR_*` environment +variables and needing the `llm` extra (`pip install firecrawl-anydoc[llm]`). +Documents anydoc converts itself never leave the machine.""" class HostedError(ConvertError): """`ocr="hosted"` could not get the document through Firecrawl Parse.""" +class LlmOcrError(ConvertError): + """`ocr="llm"` could not transcribe the document: the `llm` extra is not + installed, the `ANYDOC_OCR_*` settings are missing or invalid, or the + model call failed.""" + + def to_markdown( path: "str | os.PathLike[str]", *, ocr: Ocr = "reject", api_key: "str | None" = None, api_url: "str | None" = None, + model: "str | None" = None, ) -> str: """Convert a document file to Markdown. The format is detected from the file content; the extension is the fallback for signature-less formats @@ -69,14 +78,18 @@ def to_markdown( For `ocr="hosted"`, `api_key` falls back to `FIRECRAWL_API_KEY`, then keyless; `api_url` to `FIRECRAWL_API_URL`, then - `https://api.firecrawl.dev`.""" + `https://api.firecrawl.dev`. For `ocr="llm"`, `model` overrides + `ANYDOC_OCR_MODEL`; the rest of the configuration is read from + `ANYDOC_OCR_*`.""" try: return _to_markdown(path) except NeedsOcrError: - if ocr != "hosted": - raise - path = Path(path) - return _parse_hosted(path.read_bytes(), path.name, api_key, api_url) + if ocr == "hosted": + path = Path(path) + return _parse_hosted(path.read_bytes(), path.name, api_key, api_url) + if ocr == "llm": + return _parse_llm(Path(path).read_bytes(), model) + raise def to_markdown_bytes( @@ -86,17 +99,20 @@ def to_markdown_bytes( ocr: Ocr = "reject", api_key: "str | None" = None, api_url: "str | None" = None, + model: "str | None" = None, ) -> str: """Convert an in-memory document to Markdown. Without a format, it is detected from the content, which signature-less formats (CSV) have to - name explicitly. `ocr`, `api_key` and `api_url` are as for + name explicitly. `ocr`, `api_key`, `api_url` and `model` are as for `to_markdown`.""" try: return _to_markdown_bytes(data, format) except NeedsOcrError: - if ocr != "hosted": - raise - return _parse_hosted(bytes(data), "document.pdf", api_key, api_url) + if ocr == "hosted": + return _parse_hosted(bytes(data), "document.pdf", api_key, api_url) + if ocr == "llm": + return _parse_llm(bytes(data), model) + raise _API_URL = "https://api.firecrawl.dev" @@ -180,6 +196,20 @@ def _version() -> str: return "unknown" +def _parse_llm(data: bytes, model: "str | None") -> str: + """Rasterise the PDF and transcribe every page with a vision model through + LiteLLM. Setup problems and model-call failures alike surface as + `LlmOcrError`.""" + from anydoc import _llm_ocr + + try: + return _llm_ocr.parse_llm(data, model) + except _llm_ocr.LlmOcrConfigError as error: + raise LlmOcrError(str(error)) from error + except Exception as error: # noqa: BLE001 - any litellm failure becomes LlmOcrError + raise LlmOcrError(f"llm OCR: {error}") from error + + __all__ = [ "Asset", "Block", @@ -195,6 +225,7 @@ def _version() -> str: "LinkTarget", "List", "ListItem", + "LlmOcrError", "MalformedError", "MissingPartError", "NeedsOcrError", diff --git a/python/anydoc/_llm_ocr.py b/python/anydoc/_llm_ocr.py new file mode 100644 index 00000000..78b3e506 --- /dev/null +++ b/python/anydoc/_llm_ocr.py @@ -0,0 +1,166 @@ +"""VLM OCR for scanned PDFs, via LiteLLM. + +Loaded only when ``to_markdown(..., ocr="llm")`` reaches a PDF whose pages need +OCR, so the optional dependencies (``litellm``, ``pypdfium2``, +``pydantic-settings``) stay out of the base install. Rasterise every page to a +PNG and ask a vision model to transcribe it to Markdown, then stitch the pages +back together. + +Configuration comes from ``ANYDOC_OCR_*`` environment variables (see +``_build_settings_class``); the API key is held as a ``SecretStr`` and only +unwrapped at the ``litellm.completion`` call site. +""" + +import base64 +import io +from concurrent.futures import ThreadPoolExecutor + +DEFAULT_PROMPT = ( + "Transcribe this page to clean GitHub-Flavored Markdown. Preserve headings, " + "lists, tables, and reading order. Do not add commentary or code fences " + "around the whole answer; output only the Markdown." +) + + +class LlmOcrConfigError(Exception): + """A setup problem for ``ocr="llm"``: the extra is not installed, or the + ``ANYDOC_OCR_*`` settings are missing or invalid. The public wrapper turns + this into ``LlmOcrError``.""" + + +def _require_deps(): + """Import the optional dependencies, or explain the extra.""" + try: + import litellm + import pypdfium2 as pdfium + except ImportError as error: + raise LlmOcrConfigError( + "ocr='llm' needs the 'llm' extra: pip install firecrawl-anydoc[llm]" + ) from error + return litellm, pdfium + + +def _build_settings_class(): + """Construct the settings model lazily: importing pydantic-settings at + module load would defeat the optional-dependency split.""" + try: + from pydantic import Field, SecretStr + from pydantic_settings import BaseSettings, SettingsConfigDict + except ImportError as error: + raise LlmOcrConfigError( + "ocr='llm' needs the 'llm' extra: pip install firecrawl-anydoc[llm]" + ) from error + + class LlmOcrSettings(BaseSettings): + """Environment-driven configuration for ``ocr="llm"``. + + Every field is read from ``ANYDOC_OCR_``. ``model`` is required; + it is a LiteLLM model string such as ``openai/gpt-4o-mini`` or + ``anthropic/claude-sonnet-4-5``. ``api_key`` is optional here because + LiteLLM also reads the provider's own variable (``OPENAI_API_KEY``, + ``GEMINI_API_KEY``, ...).""" + + model_config = SettingsConfigDict(env_prefix="ANYDOC_OCR_", extra="ignore") + + model: str + api_base: "str | None" = None + api_key: "SecretStr | None" = None + prompt: str = DEFAULT_PROMPT + dpi: int = Field(default=200, gt=0) + max_pages: int = Field(default=100, gt=0) + page_concurrency: int = Field(default=4, gt=0) + timeout: float = Field(default=120.0, gt=0) + + return LlmOcrSettings + + +def _load_settings(model_override): + from pydantic import ValidationError + + build = _build_settings_class() + try: + settings = build() + except ValidationError as error: + missing = [ + "ANYDOC_OCR_" + "_".join(str(p) for p in err["loc"]).upper() + for err in error.errors() + if err["type"] == "missing" + ] + if missing: + raise LlmOcrConfigError( + f"ocr='llm' needs {', '.join(missing)} set (a LiteLLM model string, " + "e.g. ANYDOC_OCR_MODEL=openai/gpt-4o-mini)" + ) from error + raise LlmOcrConfigError(f"invalid ANYDOC_OCR_* settings: {error}") from error + if model_override: + settings.model = model_override + return settings + + +def _rasterise(pdfium, data: bytes, dpi: int, max_pages: int) -> "list[bytes]": + """Every page of the PDF as PNG bytes.""" + pdf = pdfium.PdfDocument(data) + try: + if len(pdf) > max_pages: + raise LlmOcrConfigError( + f"PDF has {len(pdf)} pages, over the ocr='llm' limit of {max_pages} " + "(raise ANYDOC_OCR_MAX_PAGES to allow it)" + ) + pages = [] + for page in pdf: + bitmap = page.render(scale=dpi / 72) + image = bitmap.to_pil() + buffer = io.BytesIO() + image.save(buffer, format="PNG") + pages.append(buffer.getvalue()) + bitmap.close() + page.close() + return pages + finally: + pdf.close() + + +def _transcribe_page(litellm, settings, png: bytes) -> str: + data_uri = "data:image/png;base64," + base64.b64encode(png).decode("ascii") + kwargs = {"timeout": settings.timeout} + if settings.api_base: + kwargs["api_base"] = settings.api_base + if settings.api_key is not None: + kwargs["api_key"] = settings.api_key.get_secret_value() + response = litellm.completion( + model=settings.model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": settings.prompt}, + {"type": "image_url", "image_url": {"url": data_uri}}, + ], + } + ], + **kwargs, + ) + return response.choices[0].message.content or "" + + +def parse_llm(data: bytes, model_override: "str | None") -> str: + """Transcribe a scanned PDF to Markdown with a vision model. Raises + ``LlmOcrConfigError`` for setup problems and lets ``litellm`` exceptions + propagate; the public wrapper turns both into ``LlmOcrError``.""" + litellm, pdfium = _require_deps() + settings = _load_settings(model_override) + pages = _rasterise(pdfium, data, settings.dpi, settings.max_pages) + + if settings.page_concurrency == 1 or len(pages) == 1: + transcribed = [_transcribe_page(litellm, settings, png) for png in pages] + else: + workers = min(settings.page_concurrency, len(pages)) + with ThreadPoolExecutor(max_workers=workers) as pool: + transcribed = list( + pool.map(lambda png: _transcribe_page(litellm, settings, png), pages) + ) + + markdown = "\n\n".join(part.strip() for part in transcribed if part.strip()) + if not markdown: + raise LlmOcrConfigError("the model returned no Markdown") + return markdown + "\n" diff --git a/python/pyproject.toml b/python/pyproject.toml index 81885146..f3fcfbce 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -21,6 +21,11 @@ classifiers = [ # The version comes from Cargo.toml `package.version`. dynamic = ["version"] +# `llm` powers `to_markdown(..., ocr="llm")`: LiteLLM calls the vision model, +# pypdfium2 rasterises the pages, pydantic-settings reads `ANYDOC_OCR_*`. +[project.optional-dependencies] +llm = ["litellm>=1.0", "pypdfium2>=4", "pydantic-settings>=2"] + [project.urls] Homepage = "https://github.com/firecrawl/anydoc#readme" Repository = "https://github.com/firecrawl/anydoc" diff --git a/python/tests/test_anydoc.py b/python/tests/test_anydoc.py index 8cff940e..fd1db112 100644 --- a/python/tests/test_anydoc.py +++ b/python/tests/test_anydoc.py @@ -1,6 +1,7 @@ """Smoke test: the bindings load and every entry point round-trips a fixture.""" import ast +import importlib.util import io import json import os @@ -10,6 +11,8 @@ from contextlib import contextmanager from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path +from types import SimpleNamespace +from unittest import mock import anydoc @@ -61,6 +64,38 @@ def log_message(self, *args): os.environ[name] = value +_LLM_EXTRA = all( + importlib.util.find_spec(name) for name in ("litellm", "pypdfium2", "pydantic_settings") +) + + +@contextmanager +def llm_env(**values): + """Set `ANYDOC_OCR_*` for the block and restore whatever was there.""" + names = {f"ANYDOC_OCR_{key.upper()}": value for key, value in values.items()} + saved = {name: os.environ.get(name) for name in names} + # Clear every ANYDOC_OCR_* so a developer's shell cannot leak in. + for name in list(os.environ): + if name.startswith("ANYDOC_OCR_"): + saved.setdefault(name, os.environ[name]) + del os.environ[name] + os.environ.update(names) + try: + yield + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _completion_reply(markdown): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=markdown))] + ) + + class AnydocTest(unittest.TestCase): def test_to_markdown_detects_the_format_from_the_file_content(self): markdown = anydoc.to_markdown(OUTLINE) @@ -146,6 +181,41 @@ def test_the_keyless_limit_says_to_set_an_api_key(self): with self.assertRaisesRegex(anydoc.HostedError, "set FIRECRAWL_API_KEY"): anydoc.to_markdown_bytes(MIXED.read_bytes(), ocr="hosted") + @unittest.skipUnless(_LLM_EXTRA, "the 'llm' extra is not installed") + def test_ocr_llm_transcribes_every_scanned_page_and_nothing_else(self): + with mock.patch("litellm.completion") as completion, llm_env(model="test/model"): + completion.side_effect = lambda **kw: _completion_reply( + f"# Page {completion.call_count}" + ) + markdown = anydoc.to_markdown(MIXED, ocr="llm") + self.assertIn("# Page 1", markdown) + # handmade-mixed.pdf is two pages, so both are sent. + self.assertEqual(completion.call_count, 2) + content = completion.call_args_list[0].kwargs["messages"][0]["content"] + self.assertEqual(content[1]["type"], "image_url") + self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/png;base64,")) + + # A document that converts locally never reaches the model. + with mock.patch("litellm.completion") as completion, llm_env(model="test/model"): + self.assertRegex(anydoc.to_markdown(OUTLINE, ocr="llm"), r"(?m)^# ") + completion.assert_not_called() + + @unittest.skipUnless(_LLM_EXTRA, "the 'llm' extra is not installed") + def test_ocr_llm_without_a_model_names_the_missing_setting(self): + with llm_env(): + with self.assertRaises(anydoc.LlmOcrError) as caught: + anydoc.to_markdown(MIXED, ocr="llm") + self.assertIsInstance(caught.exception, anydoc.ConvertError) + self.assertIn("ANYDOC_OCR_MODEL", str(caught.exception)) + + @unittest.skipUnless(_LLM_EXTRA, "the 'llm' extra is not installed") + def test_ocr_llm_wraps_model_failures(self): + with mock.patch("litellm.completion", side_effect=RuntimeError("boom")), llm_env( + model="test/model" + ): + with self.assertRaisesRegex(anydoc.LlmOcrError, "boom"): + anydoc.to_markdown_bytes(MIXED.read_bytes(), ocr="llm") + def test_unreadable_files_and_bad_arguments_raise_the_python_exception(self): with self.assertRaises(FileNotFoundError): anydoc.to_markdown("no-such-file.docx") @@ -162,7 +232,10 @@ def test_the_stubs_cover_the_module(self): exported = {name for name in dir(anydoc._anydoc) if not name.startswith("_")} self.assertEqual(stubbed, exported) # __init__.py re-exports the whole module, plus what it adds itself. - self.assertEqual(set(anydoc.__all__), exported | {"Format", "HostedError", "Ocr"}) + self.assertEqual( + set(anydoc.__all__), + exported | {"Format", "HostedError", "LlmOcrError", "Ocr"}, + ) if __name__ == "__main__":