danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat(neo4j): Day 2 — TreeWriter, TreeReader, pipeline wiring

Browse source 
Serialize a DoclingDocument to a Neo4j graph: Document + Page + Element
nodes with dynamic specific labels (SectionHeader, Paragraph, Table,
Figure, …), plus HAS_ROOT / PARENT_OF / NEXT / ON_PAGE edges. Replace-on-
write for idempotent re-ingestion.

The reader returns the verbatim document_json stored on the Document
node — reconstruction from graph nodes is deferred to v0.6.

Wired into AnalysisService._finalize_analysis: runs after conversion,
degrades gracefully by default, fails fast when neo4j_required is set.

Refs #186
This commit is contained in:
Pier-Jean Malandrino 2026-04-17 15:28:28 +02:00
parent 25a8794a0f
commit c1d3a687ac
6 changed files with 628 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
document-parser/infra/neo4j/__init__.py
View file

@ -6,5 +6,21 @@ walkers between DoclingDocument and the graph model.
from infra.neo4j.driver import Neo4jDriver, close_driver, get_driver
from infra.neo4j.schema import bootstrap_schema
from infra.neo4j.tree_reader import (
delete_document,
document_exists,
read_document_json,
)
from infra.neo4j.tree_writer import TreeWriteResult, write_document
__all__ = ["Neo4jDriver", "bootstrap_schema", "close_driver", "get_driver"]
__all__ = [
"Neo4jDriver",
"TreeWriteResult",
"bootstrap_schema",
"close_driver",
"delete_document",
"document_exists",
"get_driver",
"read_document_json",
"write_document",
]

66
document-parser/infra/neo4j/tree_reader.py Normal file
View file

@ -0,0 +1,66 @@
"""TreeReader — fetch a DoclingDocument back from Neo4j.
v0.5.0 implementation relies on the verbatim `document_json` property stored
on the Document node by TreeWriter. Reconstruction by walking Element nodes
is deferred to v0.6 (EnrichmentWriter prerequisite), where we may need to
rebuild the DoclingDocument after enrichments have been patched on graph
nodes directly.
"""
from __future__ import annotations
import logging
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__)
async def read_document_json(neo: Neo4jDriver, doc_id: str) -> str | None:
"""Return the stored DoclingDocument JSON for `doc_id`, or None if absent."""
async with neo.driver.session(database=neo.database) as session:
result = await session.run(
"MATCH (d:Document {id: $doc_id}) RETURN d.document_json AS json",
doc_id=doc_id,
)
record = await result.single()
if record is None:
return None
return record["json"]
async def document_exists(neo: Neo4jDriver, doc_id: str) -> bool:
async with neo.driver.session(database=neo.database) as session:
result = await session.run(
"MATCH (d:Document {id: $doc_id}) RETURN count(d) AS n",
doc_id=doc_id,
)
record = await result.single()
return bool(record and record["n"] > 0)
async def delete_document(neo: Neo4jDriver, doc_id: str) -> int:
"""Wipe everything related to a doc_id. Returns nodes removed."""
async with neo.driver.session(database=neo.database) as session:
result = await session.run(
"""
MATCH (d:Document {id: $doc_id})
OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n)
WITH d, collect(DISTINCT n) AS children
DETACH DELETE d
WITH children
UNWIND children AS c
DETACH DELETE c
RETURN size(children) + 1 AS removed
""",
doc_id=doc_id,
)
record = await result.single()
# Also clean up orphan elements and pages tagged with this doc_id.
await session.run(
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id
)
await session.run(
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
)
return int(record["removed"]) if record else 0

337
document-parser/infra/neo4j/tree_writer.py Normal file
View file

@ -0,0 +1,337 @@
"""TreeWriter — persist a DoclingDocument as a graph in Neo4j.
v0.5.0 strategy: replace-on-write. For a given doc_id, all existing
Document/Element/Page/Chunk nodes are wiped before re-ingestion. The full
serialized `DoclingDocument` JSON is stored as a property on the Document
node so that `TreeReader` can round-trip it verbatim — reconstruction from
graph nodes is deferred to v0.6 (see docs/design/neo4j-integration.md §2).
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__)
# Docling label → specific Neo4j label. Every node also carries :Element.
_LABEL_MAP: dict[str, str] = {
"section_header": "SectionHeader",
"title": "SectionHeader",
"paragraph": "Paragraph",
"text": "Paragraph",
"list_item": "ListItem",
"list": "ListItem",
"table": "Table",
"picture": "Figure",
"formula": "Formula",
"code": "Code",
"caption": "Caption",
"footnote": "Footnote",
"page_header": "PageHeader",
"page_footer": "PageFooter",
}
_DEFAULT_LABEL = "TextElement"
def _element_label(docling_label: str) -> str:
return _LABEL_MAP.get(docling_label.lower(), _DEFAULT_LABEL)
@dataclass
class TreeWriteResult:
doc_id: str
elements_written: int
pages_written: int
def _iter_items(doc_data: dict[str, Any]):
"""Yield every item from texts/tables/pictures/groups with its source list."""
for key in ("texts", "tables", "pictures", "groups"):
for item in doc_data.get(key, []) or []:
yield key, item
def _first_prov(item: dict[str, Any]) -> tuple[int | None, list[float] | None]:
prov = item.get("prov") or []
if not prov:
return None, None
p0 = prov[0]
bbox = p0.get("bbox")
bbox_list: list[float] | None = None
if isinstance(bbox, dict):
bbox_list = [bbox.get("l", 0.0), bbox.get("t", 0.0), bbox.get("r", 0.0), bbox.get("b", 0.0)]
elif isinstance(bbox, list):
bbox_list = list(bbox)
return p0.get("page_no"), bbox_list
def _parent_ref(item: dict[str, Any]) -> str | None:
parent = item.get("parent")
if isinstance(parent, dict):
return parent.get("$ref") or parent.get("cref")
return None
def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]:
page, bbox = _first_prov(item)
props: dict[str, Any] = {
"doc_id": doc_id,
"self_ref": item.get("self_ref") or "",
"docling_label": (item.get("label") or "").lower(),
"text": item.get("text") or "",
"prov_page": page,
"prov_bbox": bbox,
}
# Type-specific extras.
if "level" in item:
props["level"] = item.get("level")
if "caption" in item and isinstance(item.get("caption"), str):
props["caption"] = item.get("caption")
if item.get("data") and isinstance(item["data"], dict):
# Tables carry cell layout under data; stringify to keep the schema flat.
try:
props["cells_json"] = json.dumps(item["data"])
except (TypeError, ValueError):
pass
return props
def _dfs_order(doc_data: dict[str, Any]) -> list[str]:
"""Return self_refs in reading order (DFS pre-order from body)."""
by_ref: dict[str, dict[str, Any]] = {}
for _, item in _iter_items(doc_data):
ref = item.get("self_ref")
if ref:
by_ref[ref] = item
body = doc_data.get("body") or {}
order: list[str] = []
def walk(children: list[dict[str, Any]] | None) -> None:
if not children:
return
for ch in children:
ref = ch.get("$ref") or ch.get("cref")
if not ref:
continue
order.append(ref)
child = by_ref.get(ref)
if child:
walk(child.get("children"))
walk(body.get("children"))
return order
async def write_document(
neo: Neo4jDriver,
*,
doc_id: str,
filename: str,
document_json: str,
tenant_id: str = "default",
source_uri: str | None = None,
docling_version: str | None = None,
) -> TreeWriteResult:
"""Persist the full DoclingDocument tree to Neo4j.
Idempotent: wipes any existing graph for doc_id before writing.
Fails fast (exception propagates) if Neo4j is unavailable — per design §8.5.
"""
doc_data = json.loads(document_json)
ingested_at = datetime.now(tz=timezone.utc).isoformat()
elements: list[dict[str, Any]] = []
for _, item in _iter_items(doc_data):
ref = item.get("self_ref")
if not ref:
continue
specific = _element_label(item.get("label") or "")
elements.append(
{
"specific_label": specific,
"parent_ref": _parent_ref(item),
**_element_props(item, doc_id),
}
)
pages: list[dict[str, Any]] = []
for page_no_str, page_obj in (doc_data.get("pages") or {}).items():
try:
page_no = int(page_no_str)
except (TypeError, ValueError):
continue
size = page_obj.get("size") or {}
pages.append(
{
"doc_id": doc_id,
"page_no": page_no,
"width": size.get("width"),
"height": size.get("height"),
}
)
reading_order = _dfs_order(doc_data)
async with neo.driver.session(database=neo.database) as session:
async with await session.begin_transaction() as tx:
# 1. Wipe existing graph for this doc_id (replace strategy).
await tx.run(
"MATCH (d:Document {id: $doc_id}) "
"OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n) "
"DETACH DELETE d, n",
doc_id=doc_id,
)
# Also wipe orphan elements/chunks that may still reference this doc.
await tx.run(
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id
)
await tx.run(
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
)
# 2. Document node (carries the verbatim JSON for TreeReader).
await tx.run(
"""
CREATE (d:Document {
id: $doc_id,
title: $title,
source_uri: $source_uri,
ingested_at: datetime($ingested_at),
docling_version: $docling_version,
stages_applied: ['tree'],
last_tree_write: datetime($ingested_at),
tenant_id: $tenant_id,
document_json: $document_json
})
""",
doc_id=doc_id,
title=filename,
source_uri=source_uri or "",
ingested_at=ingested_at,
docling_version=docling_version or "",
tenant_id=tenant_id,
document_json=document_json,
)
# 3. Page nodes.
if pages:
await tx.run(
"UNWIND $pages AS p "
"CREATE (:Page {doc_id: p.doc_id, page_no: p.page_no, "
"width: p.width, height: p.height})",
pages=pages,
)
# 4. Element nodes — use dynamic :Element:<specific> labels via APOC-free trick.
# We split by specific label so the CREATE statement is static (no APOC).
by_specific: dict[str, list[dict[str, Any]]] = {}
for e in elements:
by_specific.setdefault(e["specific_label"], []).append(e)
for specific, batch in by_specific.items():
await tx.run(
f"""
UNWIND $batch AS e
CREATE (n:Element:{specific} {{
doc_id: e.doc_id,
self_ref: e.self_ref,
docling_label: e.docling_label,
text: e.text,
prov_page: e.prov_page,
prov_bbox: e.prov_bbox,
level: e.level,
caption: e.caption,
cells_json: e.cells_json
}})
""",
batch=batch,
)
# 5. PARENT_OF relations (tree structure). Order tracked inline.
parent_rows = [
{
"doc_id": doc_id,
"parent_ref": e["parent_ref"],
"child_ref": e["self_ref"],
"order": idx,
}
for idx, e in enumerate(elements)
if e["parent_ref"] and e["parent_ref"] != "#/body"
]
if parent_rows:
await tx.run(
"""
UNWIND $rows AS r
MATCH (p:Element {doc_id: r.doc_id, self_ref: r.parent_ref})
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
MERGE (p)-[rel:PARENT_OF]->(c)
SET rel.order = r.order
""",
rows=parent_rows,
)
# 6. HAS_ROOT for top-level children of the document body.
root_rows = [
{"doc_id": doc_id, "child_ref": e["self_ref"]}
for e in elements
if e["parent_ref"] == "#/body"
]
if root_rows:
await tx.run(
"""
UNWIND $rows AS r
MATCH (d:Document {id: r.doc_id})
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
MERGE (d)-[:HAS_ROOT]->(c)
""",
rows=root_rows,
)
# 7. ON_PAGE from first provenance.
on_page_rows = [
{"doc_id": doc_id, "self_ref": e["self_ref"], "page_no": e["prov_page"]}
for e in elements
if e["prov_page"] is not None
]
if on_page_rows:
await tx.run(
"""
UNWIND $rows AS r
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no})
MERGE (e)-[:ON_PAGE]->(p)
""",
rows=on_page_rows,
)
# 8. NEXT chain in DFS pre-order.
if len(reading_order) > 1:
pairs = [
{"doc_id": doc_id, "a": reading_order[i], "b": reading_order[i + 1]}
for i in range(len(reading_order) - 1)
]
await tx.run(
"""
UNWIND $pairs AS p
MATCH (a:Element {doc_id: p.doc_id, self_ref: p.a})
MATCH (b:Element {doc_id: p.doc_id, self_ref: p.b})
MERGE (a)-[:NEXT]->(b)
""",
pairs=pairs,
)
await tx.commit()
logger.info(
"Neo4j: wrote doc %s (%d elements, %d pages)",
doc_id,
len(elements),
len(pages),
)
return TreeWriteResult(doc_id=doc_id, elements_written=len(elements), pages_written=len(pages))

8
document-parser/main.py
View file

@ -75,6 +75,7 @@ def _build_repos() -> tuple[SqliteDocumentRepository, SqliteAnalysisRepository]:
def _build_analysis_service(
document_repo: SqliteDocumentRepository,
analysis_repo: SqliteAnalysisRepository,
neo4j_driver=None,
) -> AnalysisService:
converter = _build_converter()
chunker = _build_chunker()
@ -90,6 +91,7 @@ def _build_analysis_service(
conversion_timeout=settings.conversion_timeout,
max_concurrent=settings.max_concurrent_analyses,
config=config,
neo4j_driver=neo4j_driver,
)
@ -165,14 +167,16 @@ def _build_document_service(
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
await init_db()
document_repo, analysis_repo = _build_repos()
app.state.analysis_service = _build_analysis_service(document_repo, analysis_repo)
app.state.neo4j = await _init_neo4j()
app.state.analysis_service = _build_analysis_service(
document_repo, analysis_repo, neo4j_driver=app.state.neo4j
)
app.state.document_service = _build_document_service(document_repo, analysis_repo)
ingestion_service = _build_ingestion_service()
app.state.ingestion_service = ingestion_service
if ingestion_service is not None:
app.include_router(ingestion_router)
logger.info("Ingestion router mounted")
app.state.neo4j = await _init_neo4j()
logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine)
try:
yield

27
document-parser/services/analysis_service.py
View file

@ -69,6 +69,7 @@ class AnalysisConfig:
default_table_mode: str = "accurate"
batch_page_size: int = 0
neo4j_required: bool = False # if True, ingestion fails when Neo4j write fails
class AnalysisService:
@ -83,6 +84,7 @@ class AnalysisService:
conversion_timeout: int = 600,
max_concurrent: int = _DEFAULT_MAX_CONCURRENT,
config: AnalysisConfig | None = None,
neo4j_driver=None,
):
self._converter = converter
self._chunker = chunker
@ -93,6 +95,7 @@ class AnalysisService:
self._running_tasks: dict[str, asyncio.Task] = {}
self._background_tasks: set[asyncio.Task] = set()
self._config = config or AnalysisConfig()
self._neo4j = neo4j_driver
async def create(
self,
@ -386,8 +389,32 @@ class AnalysisService:
if result.page_count:
await self._document_repo.update_page_count(job.document_id, result.page_count)
await self._write_tree_to_neo4j(job, result.document_json)
logger.info("Analysis completed: %s (%d pages)", job_id, result.page_count)
async def _write_tree_to_neo4j(self, job, document_json: str | None) -> None:
"""Mirror the DoclingDocument tree into Neo4j if configured.
Silent no-op when Neo4j isn't wired in. Logs but does not fail the
analysis when the write fails, unless `config.neo4j_required` is set.
"""
if self._neo4j is None or not document_json:
return
try:
from infra.neo4j import write_document
await write_document(
self._neo4j,
doc_id=job.document_id,
filename=job.document_filename or job.document_id,
document_json=document_json,
)
except Exception:
logger.exception("Neo4j TreeWriter failed for doc %s", job.document_id)
if self._config.neo4j_required:
raise
async def _run_analysis_inner(
self,
job_id: str,

175
document-parser/tests/neo4j/test_tree_writer.py Normal file
View file

@ -0,0 +1,175 @@
"""TreeWriter round-trip + structural sanity checks.
Fixture is a hand-crafted DoclingDocument JSON with: one section containing
two paragraphs and a table, spanning two pages. Tests verify that the graph
mirrors the structure (HAS_ROOT, PARENT_OF, ON_PAGE, NEXT) and that
re-writing the same doc is an idempotent replace.
"""
from __future__ import annotations
import json
from infra.neo4j import read_document_json, write_document
from infra.neo4j.schema import bootstrap_schema
FIXTURE = {
"name": "fixture.pdf",
"pages": {
"1": {"page_no": 1, "size": {"width": 595, "height": 842}},
"2": {"page_no": 2, "size": {"width": 595, "height": 842}},
},
"body": {
"self_ref": "#/body",
"children": [
{"$ref": "#/texts/0"},
{"$ref": "#/texts/1"},
{"$ref": "#/texts/2"},
{"$ref": "#/tables/0"},
],
},
"texts": [
{
"self_ref": "#/texts/0",
"parent": {"$ref": "#/body"},
"label": "section_header",
"text": "Introduction",
"level": 1,
"prov": [{"page_no": 1, "bbox": {"l": 10, "t": 10, "r": 100, "b": 30}}],
},
{
"self_ref": "#/texts/1",
"parent": {"$ref": "#/body"},
"label": "paragraph",
"text": "First paragraph on page 1.",
"prov": [{"page_no": 1, "bbox": {"l": 10, "t": 40, "r": 500, "b": 80}}],
},
{
"self_ref": "#/texts/2",
"parent": {"$ref": "#/body"},
"label": "paragraph",
"text": "Continued on page 2.",
"prov": [{"page_no": 2, "bbox": {"l": 10, "t": 40, "r": 500, "b": 80}}],
},
],
"tables": [
{
"self_ref": "#/tables/0",
"parent": {"$ref": "#/body"},
"label": "table",
"text": "",
"data": {"num_rows": 2, "num_cols": 2, "grid": [[1, 2], [3, 4]]},
"prov": [{"page_no": 2, "bbox": {"l": 10, "t": 90, "r": 500, "b": 200}}],
}
],
"pictures": [],
"groups": [],
}
async def _count(session, cypher: str, **params) -> int:
r = await session.run(cypher, **params)
rec = await r.single()
return int(rec["n"]) if rec else 0
async def test_write_creates_expected_structure(neo4j_driver):
await bootstrap_schema(neo4j_driver)
doc_json = json.dumps(FIXTURE)
result = await write_document(
neo4j_driver,
doc_id="doc-fixture",
filename="fixture.pdf",
document_json=doc_json,
)
assert result.elements_written == 4
assert result.pages_written == 2
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
assert await _count(
s,
"MATCH (d:Document {id: $id}) RETURN count(d) AS n",
id="doc-fixture",
) == 1
assert await _count(
s,
"MATCH (:Document {id: $id})-[:HAS_ROOT]->(e:Element) RETURN count(e) AS n",
id="doc-fixture",
) == 4
assert await _count(
s,
"MATCH (e:Element:SectionHeader {doc_id: $id, self_ref: '#/texts/0'}) "
"RETURN count(e) AS n",
id="doc-fixture",
) == 1
assert await _count(
s,
"MATCH (e:Element:Table {doc_id: $id}) RETURN count(e) AS n",
id="doc-fixture",
) == 1
# Reading-order chain: 3 NEXT edges for 4 elements.
assert await _count(
s,
"MATCH (a:Element {doc_id: $id})-[:NEXT]->(b:Element {doc_id: $id}) "
"RETURN count(*) AS n",
id="doc-fixture",
) == 3
# ON_PAGE: one per element with prov.
assert await _count(
s,
"MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) "
"RETURN count(*) AS n",
id="doc-fixture",
) == 4
async def test_rewrite_is_idempotent_replace(neo4j_driver):
await bootstrap_schema(neo4j_driver)
doc_json = json.dumps(FIXTURE)
await write_document(
neo4j_driver,
doc_id="doc-fixture",
filename="fixture.pdf",
document_json=doc_json,
)
# Second write with the same id must not duplicate anything.
await write_document(
neo4j_driver,
doc_id="doc-fixture",
filename="fixture.pdf",
document_json=doc_json,
)
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
assert await _count(
s, "MATCH (d:Document {id: $id}) RETURN count(d) AS n", id="doc-fixture"
) == 1
assert await _count(
s,
"MATCH (e:Element {doc_id: $id}) RETURN count(e) AS n",
id="doc-fixture",
) == 4
async def test_reader_returns_verbatim_json(neo4j_driver):
await bootstrap_schema(neo4j_driver)
doc_json = json.dumps(FIXTURE, sort_keys=True)
await write_document(
neo4j_driver,
doc_id="doc-fixture",
filename="fixture.pdf",
document_json=doc_json,
)
read_back = await read_document_json(neo4j_driver, "doc-fixture")
assert read_back is not None
assert json.loads(read_back) == json.loads(doc_json)
async def test_reader_missing_doc_returns_none(neo4j_driver):
await bootstrap_schema(neo4j_driver)
assert await read_document_json(neo4j_driver, "no-such-doc") is None