8103460e9c
Adds the `docling-agent` reasoning-trace viewer as a Studio tunnel, per `docs/design/reasoning-trace.md`. Users pick an analyzed document, import a RAGResult JSON, and the iterations are overlaid on the document graph. Graph source is decoupled from Neo4j: a new pure builder (`infra/docling_graph.build_graph_payload`) reads `document_json` from SQLite and emits the same Cytoscape-shaped payload that `fetch_graph` returns from Neo4j. Neo4j stays exclusive to the Maintain ingestion pipeline. Shared DoclingDocument helpers live in `infra/docling_tree.py` so TreeWriter and the builder can't drift on label taxonomy or tree walks. Also removes the Cytoscape minimap (cytoscape-navigator) from GraphView: second render instance hurt perf on large documents for no UX win. Backend - new `GET /api/documents/:id/reasoning-graph` (SQLite-only) - new `infra/docling_tree.py`, `infra/docling_graph.py` - `analysis_repo.find_latest_completed_by_document` - tests: `test_docling_graph.py` (builder), `test_graph_api.py` (endpoint) Frontend - `features/reasoning/` — store, overlay, types, panel, import dialog, workspace, doc picker - new `ReasoningPage` + `/reasoning` and `/reasoning/:docId` routes - `GraphView` gains a `fetcher` prop so reasoning can inject the SQLite-backed fetcher while Maintain keeps using the Neo4j one - drops minimap (nav container, dep, CSS) - legend filters + section parenting extracted for reuse - i18n base strings (FR + EN)
295 lines
10 KiB
Python
295 lines
10 KiB
Python
"""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 contextlib
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from infra.docling_tree import (
|
|
dfs_order,
|
|
element_label,
|
|
iter_items,
|
|
iter_pages,
|
|
iter_provs,
|
|
parent_ref,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from infra.neo4j.driver import Neo4jDriver
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TreeWriteResult:
|
|
doc_id: str
|
|
elements_written: int
|
|
pages_written: int
|
|
provenances_written: int = 0
|
|
|
|
|
|
def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]:
|
|
"""Properties stored on the `:Element` node itself.
|
|
|
|
Provenance (page + bbox) is NOT here anymore — see `_iter_provs` and the
|
|
`:Provenance` nodes. Keeping it out of the element matches DoclingDocument's
|
|
own model (`prov` is a list of objects, not a scalar).
|
|
"""
|
|
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 "",
|
|
}
|
|
# 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.
|
|
with contextlib.suppress(TypeError, ValueError):
|
|
props["cells_json"] = json.dumps(item["data"])
|
|
return props
|
|
|
|
|
|
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=UTC).isoformat()
|
|
|
|
elements: list[dict[str, Any]] = []
|
|
# Parallel list: one row per Provenance — each refers back to its owner
|
|
# element via `self_ref`, so we can batch MATCH-and-link after both node
|
|
# sets are created.
|
|
provenances: 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),
|
|
}
|
|
)
|
|
for prov in iter_provs(item):
|
|
provenances.append({"doc_id": doc_id, "self_ref": ref, **prov})
|
|
|
|
pages: list[dict[str, Any]] = [{"doc_id": doc_id, **p} for p in iter_pages(doc_data)]
|
|
|
|
reading_order = dfs_order(doc_data)
|
|
|
|
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,
|
|
)
|
|
# Orphan sweep — covers Provenance/Element/Page/Chunk that may linger
|
|
# from an interrupted write or a pre-refactor schema.
|
|
await tx.run("MATCH (pv:Provenance {doc_id: $doc_id}) DETACH DELETE pv", doc_id=doc_id)
|
|
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:<specific> 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,
|
|
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. Provenance nodes — one per (element, prov-entry) pair. Mirrors
|
|
# Docling's `item.prov = list[ProvenanceItem]` 1:1 so a single item
|
|
# that spans page breaks (or appears twice in the layout) keeps every
|
|
# (page, bbox, charspan) without losing data.
|
|
if provenances:
|
|
await tx.run(
|
|
"""
|
|
UNWIND $rows AS r
|
|
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
|
|
CREATE (pv:Provenance {
|
|
doc_id: r.doc_id,
|
|
element_ref: r.self_ref,
|
|
prov_order: r.order,
|
|
page_no: r.page_no,
|
|
bbox_l: r.bbox_l,
|
|
bbox_t: r.bbox_t,
|
|
bbox_r: r.bbox_r,
|
|
bbox_b: r.bbox_b,
|
|
coord_origin: r.coord_origin,
|
|
charspan_start: r.charspan_start,
|
|
charspan_end: r.charspan_end
|
|
})
|
|
CREATE (e)-[:HAS_PROV {order: r.order}]->(pv)
|
|
""",
|
|
rows=provenances,
|
|
)
|
|
# ON_PAGE now attaches the Provenance to its Page — lets downstream
|
|
# queries ("what's on page 3?") stay simple without walking through
|
|
# the Element. A Provenance with no page_no (rare) yields no edge.
|
|
await tx.run(
|
|
"""
|
|
UNWIND $rows AS r
|
|
WITH r WHERE r.page_no IS NOT NULL
|
|
MATCH (pv:Provenance {
|
|
doc_id: r.doc_id,
|
|
element_ref: r.self_ref,
|
|
prov_order: r.order
|
|
})
|
|
MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no})
|
|
MERGE (pv)-[:ON_PAGE]->(p)
|
|
""",
|
|
rows=provenances,
|
|
)
|
|
|
|
# 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, %d provenances)",
|
|
doc_id,
|
|
len(elements),
|
|
len(pages),
|
|
len(provenances),
|
|
)
|
|
return TreeWriteResult(
|
|
doc_id=doc_id,
|
|
elements_written=len(elements),
|
|
pages_written=len(pages),
|
|
provenances_written=len(provenances),
|
|
)
|