diff --git a/document-parser/infra/neo4j/__init__.py b/document-parser/infra/neo4j/__init__.py index 23be98d..5655afa 100644 --- a/document-parser/infra/neo4j/__init__.py +++ b/document-parser/infra/neo4j/__init__.py @@ -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", +] diff --git a/document-parser/infra/neo4j/tree_reader.py b/document-parser/infra/neo4j/tree_reader.py new file mode 100644 index 0000000..34fd264 --- /dev/null +++ b/document-parser/infra/neo4j/tree_reader.py @@ -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 diff --git a/document-parser/infra/neo4j/tree_writer.py b/document-parser/infra/neo4j/tree_writer.py new file mode 100644 index 0000000..9d93482 --- /dev/null +++ b/document-parser/infra/neo4j/tree_writer.py @@ -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: 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)) diff --git a/document-parser/main.py b/document-parser/main.py index 1c41921..911d0f4 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -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 diff --git a/document-parser/services/analysis_service.py b/document-parser/services/analysis_service.py index 9c2f4f8..301b8a5 100644 --- a/document-parser/services/analysis_service.py +++ b/document-parser/services/analysis_service.py @@ -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, diff --git a/document-parser/tests/neo4j/test_tree_writer.py b/document-parser/tests/neo4j/test_tree_writer.py new file mode 100644 index 0000000..5c74a46 --- /dev/null +++ b/document-parser/tests/neo4j/test_tree_writer.py @@ -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