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
66 lines
2.3 KiB
Python
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
|