diff --git a/document-parser/infra/neo4j/chunk_writer.py b/document-parser/infra/neo4j/chunk_writer.py index 0b249a5..dee8220 100644 --- a/document-parser/infra/neo4j/chunk_writer.py +++ b/document-parser/infra/neo4j/chunk_writer.py @@ -16,9 +16,10 @@ from __future__ import annotations import json import logging from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any -from infra.neo4j.driver import Neo4jDriver +if TYPE_CHECKING: + from infra.neo4j.driver import Neo4jDriver logger = logging.getLogger(__name__) @@ -63,23 +64,23 @@ async def write_chunks( if ref: derived_rows.append({"chunk_id": cid, "doc_id": doc_id, "self_ref": ref}) - async with neo.driver.session(database=neo.database) as session: - async with await session.begin_transaction() as tx: - # Replace existing chunks. - await tx.run( - """ + async with ( + neo.driver.session(database=neo.database) as session, + await session.begin_transaction() as tx, + ): + # Replace existing chunks. + await tx.run( + """ MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk) DETACH DELETE c """, - doc_id=doc_id, - ) - await tx.run( - "MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id - ) + doc_id=doc_id, + ) + await tx.run("MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id) - if chunk_rows: - await tx.run( - """ + if chunk_rows: + await tx.run( + """ MATCH (d:Document {id: $doc_id}) UNWIND $rows AS r CREATE (c:Chunk { @@ -92,33 +93,33 @@ async def write_chunks( }) MERGE (d)-[:HAS_CHUNK]->(c) """, - doc_id=doc_id, - rows=chunk_rows, - ) + doc_id=doc_id, + rows=chunk_rows, + ) - if derived_rows: - await tx.run( - """ + if derived_rows: + await tx.run( + """ UNWIND $rows AS r MATCH (c:Chunk {id: r.chunk_id}) MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref}) MERGE (c)-[:DERIVED_FROM]->(e) """, - rows=derived_rows, - ) + rows=derived_rows, + ) - # Flag the Document with the new stage. - await tx.run( - """ + # Flag the Document with the new stage. + await tx.run( + """ MATCH (d:Document {id: $doc_id}) SET d.stages_applied = [s IN coalesce(d.stages_applied, []) WHERE s <> 'chunks'] + ['chunks'], d.last_chunk_write = datetime() """, - doc_id=doc_id, - ) + doc_id=doc_id, + ) - await tx.commit() + await tx.commit() logger.info( "Neo4j: wrote %d chunks (%d DERIVED_FROM) for doc %s", diff --git a/document-parser/infra/neo4j/queries.py b/document-parser/infra/neo4j/queries.py index fdcd822..3c352af 100644 --- a/document-parser/infra/neo4j/queries.py +++ b/document-parser/infra/neo4j/queries.py @@ -3,9 +3,10 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any -from infra.neo4j.driver import Neo4jDriver +if TYPE_CHECKING: + from infra.neo4j.driver import Neo4jDriver @dataclass diff --git a/document-parser/infra/neo4j/schema.py b/document-parser/infra/neo4j/schema.py index 9307d3b..5fba1eb 100644 --- a/document-parser/infra/neo4j/schema.py +++ b/document-parser/infra/neo4j/schema.py @@ -7,21 +7,21 @@ this multiple times is safe — it's the contract integration tests rely on. from __future__ import annotations import logging +from typing import TYPE_CHECKING -from infra.neo4j.driver import Neo4jDriver +if TYPE_CHECKING: + from infra.neo4j.driver import Neo4jDriver logger = logging.getLogger(__name__) CONSTRAINTS: tuple[str, ...] = ( - "CREATE CONSTRAINT document_id IF NOT EXISTS " - "FOR (d:Document) REQUIRE d.id IS UNIQUE", + "CREATE CONSTRAINT document_id IF NOT EXISTS FOR (d:Document) REQUIRE d.id IS UNIQUE", "CREATE CONSTRAINT element_composite IF NOT EXISTS " "FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE", "CREATE CONSTRAINT page_composite IF NOT EXISTS " "FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE", - "CREATE CONSTRAINT chunk_id IF NOT EXISTS " - "FOR (c:Chunk) REQUIRE c.id IS UNIQUE", + "CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE", ) INDEXES: tuple[str, ...] = ( @@ -30,8 +30,7 @@ INDEXES: tuple[str, ...] = ( ) FULLTEXT_INDEXES: tuple[str, ...] = ( - "CREATE FULLTEXT INDEX element_text IF NOT EXISTS " - "FOR (e:Element) ON EACH [e.text]", + "CREATE FULLTEXT INDEX element_text IF NOT EXISTS FOR (e:Element) ON EACH [e.text]", ) diff --git a/document-parser/infra/neo4j/tree_reader.py b/document-parser/infra/neo4j/tree_reader.py index 34fd264..06048c6 100644 --- a/document-parser/infra/neo4j/tree_reader.py +++ b/document-parser/infra/neo4j/tree_reader.py @@ -10,8 +10,10 @@ nodes directly. from __future__ import annotations import logging +from typing import TYPE_CHECKING -from infra.neo4j.driver import Neo4jDriver +if TYPE_CHECKING: + from infra.neo4j.driver import Neo4jDriver logger = logging.getLogger(__name__) @@ -57,10 +59,6 @@ async def delete_document(neo: Neo4jDriver, doc_id: str) -> int: ) 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 - ) + 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 index 9d93482..2b08ff7 100644 --- a/document-parser/infra/neo4j/tree_writer.py +++ b/document-parser/infra/neo4j/tree_writer.py @@ -9,13 +9,15 @@ graph nodes is deferred to v0.6 (see docs/design/neo4j-integration.md §2). from __future__ import annotations +import contextlib import json import logging from dataclasses import dataclass -from datetime import datetime, timezone -from typing import Any +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any -from infra.neo4j.driver import Neo4jDriver +if TYPE_CHECKING: + from infra.neo4j.driver import Neo4jDriver logger = logging.getLogger(__name__) @@ -96,10 +98,8 @@ def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]: 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: + with contextlib.suppress(TypeError, ValueError): props["cells_json"] = json.dumps(item["data"]) - except (TypeError, ValueError): - pass return props @@ -145,7 +145,7 @@ async def write_document( 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() + ingested_at = datetime.now(tz=UTC).isoformat() elements: list[dict[str, Any]] = [] for _, item in _iter_items(doc_data): @@ -179,26 +179,24 @@ async def write_document( 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 - ) + async with ( + neo.driver.session(database=neo.database) as session, + 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( - """ + # 2. Document node (carries the verbatim JSON for TreeReader). + await tx.run( + """ CREATE (d:Document { id: $doc_id, title: $title, @@ -211,32 +209,32 @@ async def write_document( 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, + 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, ) - # 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""" + # 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, @@ -250,83 +248,83 @@ async def write_document( cells_json: e.cells_json }}) """, - batch=batch, - ) + 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( - """ + # 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, - ) + 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( - """ + # 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, - ) + 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( - """ + # 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, - ) + 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( - """ + # 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, - ) + pairs=pairs, + ) - await tx.commit() + await tx.commit() logger.info( "Neo4j: wrote doc %s (%d elements, %d pages)", diff --git a/document-parser/main.py b/document-parser/main.py index 48fd426..b12f832 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -101,7 +101,7 @@ async def _init_neo4j(): logger.info("Neo4j disabled (NEO4J_URI not set)") return None - from infra.neo4j import Neo4jDriver, bootstrap_schema, get_driver + from infra.neo4j import bootstrap_schema, get_driver try: neo = await get_driver( diff --git a/document-parser/tests/neo4j/conftest.py b/document-parser/tests/neo4j/conftest.py index 8b40230..285d42e 100644 --- a/document-parser/tests/neo4j/conftest.py +++ b/document-parser/tests/neo4j/conftest.py @@ -15,7 +15,7 @@ import pytest # (e.g. local dev without the dependency installed). pytest.importorskip("neo4j") -from infra.neo4j import close_driver, get_driver # noqa: E402 +from infra.neo4j import close_driver, get_driver def _cfg() -> tuple[str, str, str]: diff --git a/document-parser/tests/neo4j/test_chunk_writer.py b/document-parser/tests/neo4j/test_chunk_writer.py index ff9145c..c5d8d97 100644 --- a/document-parser/tests/neo4j/test_chunk_writer.py +++ b/document-parser/tests/neo4j/test_chunk_writer.py @@ -12,7 +12,6 @@ from infra.neo4j import fetch_graph, write_chunks, write_document from infra.neo4j.schema import bootstrap_schema from tests.neo4j.test_tree_writer import FIXTURE - CHUNKS = [ { "text": "Introduction. First paragraph on page 1.", @@ -66,8 +65,7 @@ async def test_write_chunks_and_derived_from(neo4j_driver): for idx, expected in [(0, 2), (1, 1)]: cnt = await ( await s.run( - "MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) " - "RETURN count(e) AS n", + "MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) RETURN count(e) AS n", cid=f"doc-fixture::chunk::{idx}", ) ).single() diff --git a/document-parser/tests/neo4j/test_tree_writer.py b/document-parser/tests/neo4j/test_tree_writer.py index 5c74a46..3885bbc 100644 --- a/document-parser/tests/neo4j/test_tree_writer.py +++ b/document-parser/tests/neo4j/test_tree_writer.py @@ -13,7 +13,6 @@ import json from infra.neo4j import read_document_json, write_document from infra.neo4j.schema import bootstrap_schema - FIXTURE = { "name": "fixture.pdf", "pages": { @@ -89,41 +88,59 @@ async def test_write_creates_expected_structure(neo4j_driver): 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 + 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 + 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 + 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): @@ -145,14 +162,18 @@ async def test_rewrite_is_idempotent_replace(neo4j_driver): ) 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 + 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):