bd73c4bbfd
Pre-existing ruff violations surfaced by CI on PR #188: - TC001: move runtime `Neo4jDriver` imports into `TYPE_CHECKING` blocks (queries.py, chunk_writer.py, schema.py, tree_reader.py, tree_writer.py) - SIM117: combine nested `async with` in chunk_writer.write_chunks and tree_writer.write_document - SIM105: replace try/except/pass with `contextlib.suppress` in tree_writer._element_props - F401: remove unused `Neo4jDriver` import in main._init_neo4j - RUF100: remove unused `# noqa: E402` in tests/neo4j/conftest.py - I001: sort imports in tests/neo4j/test_chunk_writer.py and test_tree_writer.py Zero behaviour change. `ruff check .` now passes cleanly.
64 lines
2.4 KiB
Python
64 lines
2.4 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 typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
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
|