docling-studio/document-parser/infra/neo4j/tree_reader.py
Pier-Jean Malandrino c1d3a687ac feat(neo4j): Day 2 — TreeWriter, TreeReader, pipeline wiring
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
2026-04-29 14:00:00 +02:00

66 lines
2.3 KiB
Python

"""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