Skip to content

Commit 82300a5

Browse files
jexpclaude
andcommitted
feat: add LiteParseLoader backed by LiteParse (optional extra)
Adds LiteParseLoader to neo4j_graphrag.experimental.components.data_loader as an alternative PDF/document loader using LiteParse — a local, zero-cloud Rust-based parser with optional Tesseract OCR support. - New liteparse optional extra: pip install "neo4j-graphrag[liteparse]" - Lazy-imports liteparse; raises a clear ImportError with install hint when absent - Parser instance cached per loader to avoid repeated Rust init overhead - Supports local FS fast-path (file path) and non-local FS (bytes) via fsspec - 9 unit tests (fully mocked, no liteparse install required) - 5 integration tests (real liteparse; auto-skipped when not installed) - Usage example in examples/customize/build_graph/components/loaders/ Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 7b69354 commit 82300a5

7 files changed

Lines changed: 401 additions & 5 deletions

File tree

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22

33
## Next
44

5+
### Added
6+
7+
- Experimental: `LiteParseLoader` — an alternative PDF/document loader backed by [LiteParse](https://github.com/run-llama/liteparse), a local Rust-based parser with optional OCR support. Install with `pip install "neo4j-graphrag[liteparse]"` and pass it as `file_loader` to `SimpleKGPipeline`.
8+
59
## 1.17.0
610

711
### Added
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
"""Example: using LiteParseLoader as an alternative PDF loader.
2+
3+
LiteParse (https://github.com/run-llama/liteparse) is a local, zero-cloud PDF
4+
parser with optional OCR support. Install the optional extra before running:
5+
6+
pip install "neo4j-graphrag[liteparse]"
7+
8+
Usage::
9+
10+
from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
11+
from neo4j_graphrag.experimental.components.data_loader import LiteParseLoader
12+
13+
pipeline = SimpleKGPipeline(
14+
llm=...,
15+
driver=...,
16+
embedder=...,
17+
file_loader=LiteParseLoader(ocr_enabled=True),
18+
from_file=True,
19+
)
20+
await pipeline.run_async(file_path="document.pdf")
21+
"""
22+
23+
# LiteParseLoader lives in the main package — import it from there.
24+
from neo4j_graphrag.experimental.components.data_loader import LiteParseLoader
25+
26+
__all__ = ["LiteParseLoader"]

pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ examples = [
7171
# only advertise/install the `nlp` extra on Python <3.14.
7272
nlp = ["spacy==3.8.7; python_version < '3.14'"]
7373
fuzzy-matching = ["rapidfuzz>=3.12.2,<4.0.0"]
74+
liteparse = ["liteparse>=2.0.0,<3.0.0"]
7475

7576
[dependency-groups]
7677
dev = [
@@ -115,6 +116,7 @@ exclude = ["docs"]
115116
module = [
116117
"weaviate.*",
117118
"sentence_transformers.*",
119+
"liteparse.*",
118120
"conftest",
119121
]
120122
ignore_missing_imports = true

src/neo4j_graphrag/experimental/components/data_loader.py

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,21 +12,29 @@
1212
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
15-
"""Document loaders: base class, PDF, Markdown, and extension-based file dispatch."""
15+
"""Document loaders: base class, PDF, Markdown, LiteParse, and extension-based dispatch."""
1616

1717
from __future__ import annotations
1818

1919
import io
2020
from abc import abstractmethod
2121
from pathlib import Path
22-
from typing import Dict, Optional, Union, cast
22+
from typing import Any, Dict, Optional, Union, cast
2323

2424
import fsspec
2525
import pypdf
2626
from fsspec import AbstractFileSystem
2727
from fsspec.implementations.local import LocalFileSystem
2828

2929
from neo4j_graphrag.exceptions import MarkdownLoadError, PdfLoaderError
30+
31+
try:
32+
from liteparse import LiteParse as _LiteParse # type: ignore[import]
33+
34+
_LITEPARSE_AVAILABLE = True
35+
except ImportError:
36+
_LiteParse = None # type: ignore[assignment,misc]
37+
_LITEPARSE_AVAILABLE = False
3038
from neo4j_graphrag.experimental.components.types import (
3139
DocumentInfo,
3240
DocumentType,
@@ -137,3 +145,91 @@ async def run(
137145
document_type=DocumentType.MARKDOWN,
138146
),
139147
)
148+
149+
150+
class LiteParseLoader(DataLoader):
151+
"""Loads and parses PDF files using LiteParse (local, no cloud dependency).
152+
153+
LiteParse uses a compiled Rust core and optional Tesseract OCR to extract text
154+
from PDFs, DOCX, XLSX, PPTX, and images. It runs fully offline — no API key
155+
or network access required.
156+
157+
Requires the ``liteparse`` optional extra::
158+
159+
pip install "neo4j-graphrag[liteparse]"
160+
161+
Args:
162+
ocr_enabled: Enable Tesseract OCR for scanned pages.
163+
ocr_server_url: URL of a remote OCR server (optional).
164+
ocr_language: Tesseract language code, e.g. ``"eng"``, ``"fra"``.
165+
dpi: Rendering resolution for OCR (default 300).
166+
target_pages: Page range string, e.g. ``"1-5,10,15-20"``.
167+
password: Password for encrypted PDFs.
168+
"""
169+
170+
def __init__(
171+
self,
172+
ocr_enabled: bool = False,
173+
ocr_server_url: Optional[str] = None,
174+
ocr_language: Optional[str] = None,
175+
dpi: Optional[int] = None,
176+
target_pages: Optional[str] = None,
177+
password: Optional[str] = None,
178+
) -> None:
179+
optional = {
180+
"ocr_server_url": ocr_server_url,
181+
"ocr_language": ocr_language,
182+
"dpi": dpi,
183+
"target_pages": target_pages,
184+
"password": password,
185+
}
186+
self._kwargs: Dict[str, Any] = {
187+
"ocr_enabled": ocr_enabled,
188+
**{k: v for k, v in optional.items() if v is not None},
189+
}
190+
self._parser: Any = None # lazily initialised
191+
192+
def _get_parser(self) -> Any:
193+
if not _LITEPARSE_AVAILABLE:
194+
raise ImportError(
195+
"liteparse is required for LiteParseLoader. "
196+
'Install it with: pip install "neo4j-graphrag[liteparse]"'
197+
)
198+
if self._parser is None:
199+
self._parser = _LiteParse(**self._kwargs)
200+
return self._parser
201+
202+
def load_file(self, filepath: str, fs: AbstractFileSystem) -> str:
203+
"""Parse a document and return the full extracted text."""
204+
parser = self._get_parser() # ImportError propagates; not caught below
205+
try:
206+
if is_default_fs(fs):
207+
result = parser.parse(filepath)
208+
else:
209+
with fs.open(filepath, "rb") as fp:
210+
result = parser.parse(fp.read())
211+
return str(result.text)
212+
except Exception as e:
213+
raise PdfLoaderError(e) from e
214+
215+
async def run(
216+
self,
217+
filepath: Union[str, Path],
218+
metadata: Optional[Dict[str, str]] = None,
219+
fs: Optional[Union[AbstractFileSystem, str]] = None,
220+
) -> LoadedDocument:
221+
if not isinstance(filepath, str):
222+
filepath = str(filepath)
223+
if isinstance(fs, str):
224+
fs = fsspec.filesystem(fs)
225+
elif fs is None:
226+
fs = LocalFileSystem()
227+
text = self.load_file(filepath, fs)
228+
return LoadedDocument(
229+
text=text,
230+
document_info=DocumentInfo(
231+
path=filepath,
232+
metadata=self.get_document_metadata(text, metadata),
233+
document_type=DocumentType.PDF,
234+
),
235+
)
Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
# Copyright (c) "Neo4j"
2+
# Neo4j Sweden AB [https://neo4j.com]
3+
# #
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
# #
8+
# https://www.apache.org/licenses/LICENSE-2.0
9+
# #
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
"""Tests for LiteParseLoader — liteparse is mocked throughout so no install needed."""
17+
18+
import sys
19+
from pathlib import Path
20+
from types import ModuleType, SimpleNamespace
21+
from typing import Generator
22+
from unittest.mock import MagicMock, patch
23+
24+
import pytest
25+
from fsspec.implementations.local import LocalFileSystem
26+
27+
from neo4j_graphrag.exceptions import PdfLoaderError
28+
from neo4j_graphrag.experimental.components.data_loader import LiteParseLoader
29+
from neo4j_graphrag.experimental.components.types import DocumentType
30+
31+
SAMPLE_PDF = str(Path(__file__).parent / "sample_data/lorem_ipsum.pdf")
32+
SAMPLE_TEXT = "Lorem ipsum dolor sit amet."
33+
34+
35+
def _make_liteparse_stub(parse_text: str = SAMPLE_TEXT) -> ModuleType:
36+
result = SimpleNamespace(text=parse_text)
37+
LiteParse = MagicMock(return_value=MagicMock(parse=MagicMock(return_value=result)))
38+
stub = ModuleType("liteparse")
39+
stub.LiteParse = LiteParse # type: ignore[attr-defined]
40+
return stub
41+
42+
43+
@pytest.fixture(autouse=True)
44+
def inject_liteparse_stub() -> Generator[ModuleType, None, None]:
45+
stub = _make_liteparse_stub()
46+
with patch.dict(sys.modules, {"liteparse": stub}), patch(
47+
"neo4j_graphrag.experimental.components.data_loader._LiteParse",
48+
stub.LiteParse,
49+
), patch(
50+
"neo4j_graphrag.experimental.components.data_loader._LITEPARSE_AVAILABLE",
51+
True,
52+
):
53+
yield stub
54+
55+
56+
def test_load_file_local_fs_returns_text(inject_liteparse_stub: ModuleType) -> None:
57+
loader = LiteParseLoader()
58+
text = loader.load_file(SAMPLE_PDF, fs=LocalFileSystem())
59+
assert text == SAMPLE_TEXT
60+
inject_liteparse_stub.LiteParse.return_value.parse.assert_called_once_with(
61+
SAMPLE_PDF
62+
)
63+
64+
65+
def test_load_file_non_local_fs_uses_bytes(inject_liteparse_stub: ModuleType) -> None:
66+
fake_fs = MagicMock()
67+
fake_bytes = b"%PDF fake"
68+
fake_fs.open.return_value.__enter__ = MagicMock(
69+
return_value=MagicMock(read=MagicMock(return_value=fake_bytes))
70+
)
71+
fake_fs.open.return_value.__exit__ = MagicMock(return_value=False)
72+
73+
loader = LiteParseLoader()
74+
with patch(
75+
"neo4j_graphrag.experimental.components.data_loader.is_default_fs",
76+
return_value=False,
77+
):
78+
text = loader.load_file(SAMPLE_PDF, fs=fake_fs)
79+
80+
assert text == SAMPLE_TEXT
81+
inject_liteparse_stub.LiteParse.return_value.parse.assert_called_once_with(
82+
fake_bytes
83+
)
84+
85+
86+
def test_load_file_wraps_parse_error_as_pdf_loader_error(
87+
inject_liteparse_stub: ModuleType,
88+
) -> None:
89+
inject_liteparse_stub.LiteParse.return_value.parse.side_effect = RuntimeError(
90+
"corrupt PDF"
91+
)
92+
loader = LiteParseLoader()
93+
with pytest.raises(PdfLoaderError):
94+
loader.load_file(SAMPLE_PDF, fs=LocalFileSystem())
95+
96+
97+
def test_missing_liteparse_raises_import_error() -> None:
98+
with patch(
99+
"neo4j_graphrag.experimental.components.data_loader._LITEPARSE_AVAILABLE",
100+
False,
101+
):
102+
loader = LiteParseLoader()
103+
with pytest.raises(ImportError, match="pip install"):
104+
loader.load_file(SAMPLE_PDF, fs=LocalFileSystem())
105+
106+
107+
@pytest.mark.asyncio
108+
async def test_run_returns_loaded_document() -> None:
109+
loader = LiteParseLoader()
110+
doc = await loader.run(filepath=SAMPLE_PDF)
111+
assert doc.text == SAMPLE_TEXT
112+
assert doc.document_info.document_type == DocumentType.PDF
113+
assert doc.document_info.path == SAMPLE_PDF
114+
115+
116+
@pytest.mark.asyncio
117+
async def test_run_with_path_object() -> None:
118+
loader = LiteParseLoader()
119+
doc = await loader.run(filepath=Path(SAMPLE_PDF))
120+
assert doc.document_info.path == SAMPLE_PDF
121+
122+
123+
@pytest.mark.asyncio
124+
async def test_run_passes_metadata() -> None:
125+
loader = LiteParseLoader()
126+
meta = {"source": "test", "lang": "en"}
127+
doc = await loader.run(filepath=SAMPLE_PDF, metadata=meta)
128+
assert doc.document_info.metadata == meta
129+
130+
131+
@pytest.mark.asyncio
132+
async def test_run_fs_string_resolves() -> None:
133+
loader = LiteParseLoader()
134+
doc = await loader.run(filepath=SAMPLE_PDF, fs="file")
135+
assert doc.text == SAMPLE_TEXT
136+
137+
138+
def test_constructor_kwargs_forwarded_to_liteparse(
139+
inject_liteparse_stub: ModuleType,
140+
) -> None:
141+
loader = LiteParseLoader(
142+
ocr_enabled=True,
143+
ocr_language="fra",
144+
dpi=300,
145+
target_pages="1-3",
146+
password="s3cr3t",
147+
)
148+
loader.load_file(SAMPLE_PDF, fs=LocalFileSystem())
149+
inject_liteparse_stub.LiteParse.assert_called_once_with(
150+
ocr_enabled=True,
151+
ocr_language="fra",
152+
dpi=300,
153+
target_pages="1-3",
154+
password="s3cr3t",
155+
)

0 commit comments

Comments
 (0)