diff --git a/docs/design/reasoning-trace.md b/docs/design/reasoning-trace.md new file mode 100644 index 0000000..db69578 --- /dev/null +++ b/docs/design/reasoning-trace.md @@ -0,0 +1,253 @@ +# Reasoning Trace Viewer — Docling-Studio v0.6.0 (R&D preview) + +Design doc for the `docling-agent` reasoning trace viewer. +Targeted release: **v0.6.0** (R&D branched from `release/0.5.0` in parallel to the +0.5 build, so the Neo4j foundation can be leveraged without blocking the hackathon deliverable). + +Positioning one-liner: +> Studio becomes the **reference viewer for any `docling-agent` run** — not another +> chatbot. The PDF is the debug surface. + +--- + +## 1. Context + +### Upstream trigger +Peter Staar (IBM) suggested surfacing the LLM reasoning trace as `docling-agent` +walks a `DoclingDocument` outline (chunkless RAG, new IBM repo). + +### `docling-agent` in one paragraph +`DoclingRAGAgent(model_id, tools, max_iterations=5).run(task, sources=[doc])` returns +a `RAGResult` with `answer`, `converged`, and `iterations: list[RAGIteration]`. +Each `RAGIteration` carries: `iteration`, `section_ref` (JSON-pointer, e.g. `#/texts/3`), +`reason`, `can_answer`, `response`, `section_text_length`. No bbox — must be resolved +through `DoclingDocument.[i].prov[0].bbox` + `page_no`. Runs on Mellea +(Ollama / OpenAI / HF / WatsonX / LiteLLM / Bedrock). Observability is stdout logs only. + +### What Studio already brings +- **Neo4j graph of the document** (0.5.0, just landed) — every `Element` is keyed by + `(doc_id, self_ref)`, Cytoscape node id is `elem::${self_ref}`. **This is the + killer enabler.** `RAGIteration.section_ref → node` is a string concat, no resolver. +- `GraphView.vue` (Cytoscape + dagre) already handles styles via selectors + (`selector: 'edge[type = "NEXT"]'`, `selector: 'node[kind = "section"]'`) — adding + a `visited` class + `REASONING_NEXT` synthetic edge type is ~20 LOC of style. +- `analysis_jobs.document_json` in SQLite → DoclingDocument available for the sidecar + runner (no PDF re-conversion). Not used by the viewer itself. + +### Personas +- **v1 (this plan)**: dev / integrator of `docling-agent` debugging a run that went wrong. +- **v2 (roadmap)**: live runner with synchronized demo UX. +- v3+ (non-goals here): business analyst for semantic navigation, batch QA. + +--- + +## 2. Scope split — **debug first, demo second** + +Rendering surface pivoted: **the trace is drawn on the Neo4j graph**, not on the PDF. +See §1. This kills the whole bbox-resolution stack from v1. + +| Phase | Value | Runtime deps | Surface | +|---|---|---|---| +| **v1 — Debug (this plan)** | Import externally-produced `RAGResult` JSON, overlay trace on the existing GraphView | **None** server-side. Pure frontend. | GraphView: visited nodes highlighted in order, synthetic `REASONING_NEXT` edges | +| **v2 — Demo (follow-up)** | Run the agent live against a loaded document | Ollama + Mellea + `docling-agent` (new opt-dep group `rag`) | Same GraphView + SSE streaming of iterations, staggered reveal | + +Building v1 first de-risks the **graph-trace UX** on real runs (produced by the +R&D sidecar — see `experiments/reasoning-trace/`) before wiring the live runner. +Code shared between v1 and v2 is the GraphView overlay itself — 100 % reused. + +**Prerequisite for v1**: the target document must have been processed through the +"Maintain" step (Neo4j pipeline). Otherwise the graph is empty and the trace has +nowhere to render — surface an explicit "Run the Maintain step first" empty state. + +--- + +## 3. v1 — Debug mode (frontend-only) + +### 3.1 No backend changes in v1 + +The GraphView already loads nodes keyed by `self_ref` via `GET /api/documents/{doc_id}/graph`. +Iteration `section_ref` → Cytoscape node id is `` `elem::${section_ref}` `` — a client-side +string concat. Nothing to compute server-side. + +Consequences: +- No new router, no new service, no new pydantic model, no new migration. +- No dependency on `docling-agent` in `document-parser/requirements.txt`. +- `RAGResult` JSON (as produced by `experiments/reasoning-trace/`) is consumed + entirely by the frontend. + +### 3.2 Frontend — feature folder + +New `frontend/src/features/reasoning/`: + +``` +reasoning/ +├── store/reasoningStore.ts # Pinia: trace, activeIteration, importDialogOpen +├── ui/ +│ ├── ReasoningPanel.vue # Side panel: query, answer, iteration list +│ ├── IterationCard.vue # Single iteration row (reason + can_answer badge) +│ ├── ImportTraceDialog.vue # Drag-drop / paste RAGResult JSON +│ └── GraphReasoningOverlay.ts # NOT a component — a plugin that decorates cy +└── types.ts # RAGIteration, RAGResult mirror types +``` + +### 3.3 Graph overlay — how it's drawn + +`GraphReasoningOverlay` takes the existing `cy` Cytoscape instance (exposed from +`GraphView.vue` via `defineExpose`) and: + +1. For each `iteration[i].section_ref`, find node `` `elem::${section_ref}` ``. If + missing, tag as `resolution_status: "not_in_graph"` and show a warning in the panel + (common cause: doc not processed through Maintain, or agent returned a ref that + points at a non-Element node). +2. Add class `visited` + data attribute `visitOrder: i+1` on matched nodes. +3. Insert **synthetic edges** between successive visited nodes with `type: "REASONING_NEXT"` + and `data: { order: i }`. These edges are UI-only, never written to Neo4j. +4. On import, fit viewport to the visited subgraph (`cy.fit(cy.$('.visited'), 80)`). +5. On iteration card click → `cy.$(`#elem::${ref}`).flashClass('pulse', 800)` + + centered pan. + +Cytoscape styles (append to the existing stylesheet array in `GraphView.vue`): + +```js +{ selector: 'node.visited', + style: { 'border-color': '#EA580C', 'border-width': 3, 'overlay-opacity': 0 } }, +{ selector: 'node.visited[visitOrder]', + style: { label: 'data(visitOrder)', 'text-valign': 'top', + 'text-background-color': '#EA580C', 'text-background-opacity': 1, + 'color': '#FFFFFF', 'font-weight': 700 } }, +{ selector: 'edge[type = "REASONING_NEXT"]', + style: { 'line-color': '#EA580C', 'target-arrow-color': '#EA580C', + 'target-arrow-shape': 'triangle', 'curve-style': 'bezier', + width: 2, 'z-index': 99 } }, +``` + +Color ramp: single warm color (`#EA580C`) for v1. Gradient cold→warm is v2 polish. + +### 3.4 Integration points + +- `StudioPage.vue` → "Maintain" tab gains an **"Import reasoning trace"** action + (don't add a 3rd mode — the viz lives inside the graph view, not a new workspace). +- `GraphView.vue` → add `defineExpose({ cy })` + a `` + that the parent can populate with ``. +- `ReasoningPanel` appears as a right rail when a trace is loaded; collapsible. + +### 3.5 Empty / error states + +- **Graph empty for this doc** → "Run the Maintain step first. Neo4j has no graph for + this document yet." (the Maintain button is literally next to it.) +- **All `section_ref`s unresolved in graph** → "None of the visited sections exist in + the graph. The agent may have been run against a different document, or the doc was + re-analyzed since. Re-run Maintain or re-run the agent." +- **Some resolved, some not** → show trace with the missing ones greyed out in the panel. + +### 3.6 Tests + +No backend tests in v1 (no backend code). + +Frontend (Vitest): +- `reasoningStore.test.ts` — import trace, active iteration transitions, reset on doc change. +- `graphReasoningOverlay.test.ts` — given a mock `cy` (`cytoscape({ headless: true })`) + with a known node set, verify `visited` class applied to the right ids and the + correct synthetic edges added. +- `ReasoningPanel.test.ts` — empty / loaded / partial-resolution states. + +### 3.7 Out of scope for v1 +- Live agent runner (v2). +- Multi-doc queries — reject import if `RAGResult` was produced against `len(sources) > 1`. +- Phrase-level attribution — `docling-agent` doesn't emit it. +- Persisting traces in Neo4j — see §7. +- PDF highlighting — dropped from v1. Could come back as v2.5 if demand exists. + +--- + +## 4. File inventory (v1) + +**New — R&D sidecar** (already scaffolded on this branch) +- `experiments/reasoning-trace/inspect_doc.py` — self-contained `uv run` script. +- `experiments/reasoning-trace/README.md` +- `experiments/reasoning-trace/.gitignore` + +**New — frontend** +- `frontend/src/features/reasoning/**` (see §3.2) +- Vitest siblings under `**/*.test.ts` + +**Touched** +- `frontend/src/features/analysis/ui/GraphView.vue` — `defineExpose({ cy })` + + `` + 3 new style selectors. +- `frontend/src/pages/StudioPage.vue` — "Import reasoning trace" action in the + Maintain tab rail. + +**Untouched** +- Entire `document-parser/` backend — no new router, service, schema, or dep. +- `pyproject.toml` / `requirements.txt` — **no new runtime dep in v1**. +- Neo4j schema — synthetic edges are client-side Cytoscape only. +- OpenSearch / ingestion — untouched. +- SQLite schema — no migration. + +--- + +## 5. Risks & mitigations + +| Risk | Mitigation | +|---|---| +| `RAGResult` schema drifts in `docling-agent` | `schema_version` discriminator; strict pydantic; one canonical fixture from Peter pinned in CI. | +| `section_ref` variants (`#/texts/3` vs `#/body/texts/3`) | Normalize in parser; regex test matrix. | +| Synthetic groups without `prov` | Documented child-walk fallback + `resolved_via_child` status surfaced in UI. | +| Large `RAGResult` (hundreds of iterations) | Hard-cap `iterations` at 50 in v1 (Peter's agent uses `max_iterations=5` by default) — return 413 above. | +| `document_json` blob large (some docs > 5 MB) | `analysis_repo` already handles it; but **do not** log the blob. Add redaction test. | +| Section ref not in graph (doc not through Maintain, or re-analyzed) | Explicit empty-state in `ReasoningPanel` with a link to the Maintain tab. Partial resolution shown as grey in the trace list. | +| Feature creeping into 0.5.0 | This branch targets **v0.6.0**. Do not merge into `release/0.5.0`. Rebase onto the next release branch when cut. | + +--- + +## 6. Spec anchoring + +Pin the `RAGResult` shape to **docling-agent commit SHA at the time of v1 merge** in +a short ADR `docs/architecture/adrs/ADR-002-rag-result-schema.md`. The schema is +upstream, unversioned, and will move — this doc freezes the contract Studio imports. + +--- + +## 7. v2 preview — demo mode (not in this plan) + +Kept here to constrain v1 interfaces so nothing needs rewriting: +- `POST /api/rag/answer` — server-side runner. Accepts `{doc_id, question, model_id}`. + Streams iterations via SSE. Frontend consumes the stream with the same + `GraphReasoningOverlay` used by v1 import — iterations appear one by one with + staggered reveal (~400 ms) as the SSE stream drips them in. +- Ollama wired through `Mellea` — new optional dep group `rag`. +- Persist traces in Neo4j as `(:ReasoningRun {id, query, converged})-[:VISITED {order, + reason, can_answer}]->(:Element)` for replay + cross-run analytics. Leverages + `TreeWriter` pattern already present. This is where the synthetic UI edges become + real graph edges. +- Cross-run comparison view: overlay multiple runs on the same graph, diff the paths. + +--- + +## 8. Branch & workflow + +- Branch: **`feature/reasoning-trace`** off `origin/release/0.5.0`. +- Merge target: **next release branch (`release/0.6.0`)** once cut — *not* `0.5.0`. +- Until then: live on the feature branch; rebase onto `release/0.5.0` periodically to + absorb Neo4j fixes. +- Issues: one umbrella + one per §4 subsystem (resolver, endpoint, UI panel, overlay, + import dialog, tests). Commit with `Closes #NNN` per project convention. +- PR: opened against `release/0.6.0` when available; draft in the meantime. + +--- + +## 9. Open questions (answered by the sidecar first run) + +1. Are emitted `section_ref`s reachable as `elem::${ref}` in the Neo4j graph built + by `TreeWriter`? I.e. is the `self_ref` the agent sees the same `self_ref` we + wrote to the graph? (Expected yes — both come from the same `DoclingDocument` — + but the sidecar on a real doc from SQLite will confirm in one run.) +2. Hit rate of the agent: with `max_iterations=5` and `granite4:micro-h`, does it + converge, and how many sections does it actually visit? Determines if the overlay + ever has more than 1–2 marked nodes (and whether `REASONING_NEXT` edges are worth + the effort vs just node markers). +3. Quality of `iteration.reason` — is it substantive enough to show in the panel, or + LLM filler we should hide? Sidecar output will tell. +4. Fallback when no section headers exist (`RAGResult(iterations=[], converged=True, + answer=)` — see rag.py): what does the panel show? Probably a degraded + "no trace available, full-doc answer" state. diff --git a/document-parser/api/graph.py b/document-parser/api/graph.py index 9b7b4a5..f38e8e2 100644 --- a/document-parser/api/graph.py +++ b/document-parser/api/graph.py @@ -1,9 +1,11 @@ -"""Graph API — returns a cytoscape-shaped view of the Neo4j graph for a doc. +"""Graph API — returns a cytoscape-shaped view of the document structure. -v0.5 contract: -- Returns the **full** graph for the document (see design §8.4) -- Hard cap at 200 pages; beyond that, HTTP 413 with `truncated: true` -- No pagination (ships in v0.6) +Two endpoints: +- `/graph` — read from Neo4j. Rich graph (elements + chunks + pages + merges). + Requires the Maintain step (IngestionPipeline) to have run for the document. +- `/reasoning-graph` — built on-the-fly from the SQLite `document_json` blob. + No Neo4j dependency. Lighter graph (no chunks) but enough to render the + reasoning-trace overlay on top of `GraphView`. """ from __future__ import annotations @@ -13,6 +15,7 @@ import logging from fastapi import APIRouter, HTTPException, Request from pydantic import BaseModel +from infra.docling_graph import build_graph_payload from infra.neo4j import fetch_graph logger = logging.getLogger(__name__) @@ -74,3 +77,47 @@ async def get_document_graph(doc_id: str, request: Request) -> GraphResponse: truncated=payload.truncated, page_count=payload.page_count, ) + + +@router.get("/{doc_id}/reasoning-graph", response_model=GraphResponse) +async def get_reasoning_graph(doc_id: str, request: Request) -> GraphResponse: + """Graph projection built from SQLite `document_json` — no Neo4j needed. + + Serves the reasoning-trace viewer, which only needs the element/page/edge + structure to overlay iterations onto. + """ + analysis_repo = getattr(request.app.state, "analysis_repo", None) + if analysis_repo is None: + raise HTTPException(status_code=500, detail="AnalysisRepository not wired") + + latest = await analysis_repo.find_latest_completed_by_document(doc_id) + if latest is None or not latest.document_json: + raise HTTPException( + status_code=404, + detail=f"No completed analysis with document_json for {doc_id}", + ) + + payload = build_graph_payload( + latest.document_json, + doc_id=doc_id, + title=latest.document_filename or doc_id, + max_pages=MAX_PAGES, + ) + if payload.truncated: + raise HTTPException( + status_code=413, + detail=( + f"Graph too large: document has {payload.page_count} pages " + f"(cap {MAX_PAGES}). Pagination ships in v0.6." + ), + ) + + return GraphResponse( + doc_id=payload.doc_id, + nodes=[GraphNode(**n) for n in payload.nodes], + edges=[GraphEdge(**e) for e in payload.edges], + node_count=payload.node_count, + edge_count=payload.edge_count, + truncated=payload.truncated, + page_count=payload.page_count, + ) diff --git a/document-parser/infra/docling_graph.py b/document-parser/infra/docling_graph.py new file mode 100644 index 0000000..bd7cdf3 --- /dev/null +++ b/document-parser/infra/docling_graph.py @@ -0,0 +1,159 @@ +"""Build a Cytoscape-shaped graph payload straight from a serialized +`DoclingDocument` (i.e. the `document_json` blob stored in SQLite). + +Mirrors `infra.neo4j.queries.fetch_graph` so the frontend can reuse the same +`GraphView` component — the only intentional difference is the absence of +Chunk nodes / HAS_CHUNK / DERIVED_FROM edges, since chunks are a product of +the Maintain step and don't exist in `document_json` alone. + +Used by the reasoning-trace viewer, which needs the structural graph to +overlay iterations onto but does NOT need (and should not require) Neo4j. +""" + +from __future__ import annotations + +import json +from itertools import pairwise +from typing import Any + +from infra.docling_tree import ( + dfs_order, + element_label, + iter_items, + iter_pages, + iter_provs, + parent_ref, +) +from infra.neo4j.queries import GraphPayload + + +def _element_node(doc_id: str, item: dict[str, Any], provs: list[dict[str, Any]]) -> dict[str, Any]: + first_page = provs[0].get("page_no") if provs else None + return { + "id": f"elem::{item.get('self_ref')}", + "group": "element", + "label": element_label(item.get("label") or ""), + "docling_label": (item.get("label") or "").lower(), + "self_ref": item.get("self_ref"), + "text": (item.get("text") or "")[:200], + "prov_page": first_page, + "provs": provs, + "level": item.get("level"), + "doc_id": doc_id, + } + + +def _page_node(doc_id: str, page: dict[str, Any]) -> dict[str, Any]: + return { + "id": f"page::{page.get('page_no')}", + "group": "page", + "page_no": page.get("page_no"), + "width": page.get("width"), + "height": page.get("height"), + "doc_id": doc_id, + } + + +def _edge(source: str, target: str, edge_type: str, *, order: int | None = None) -> dict[str, Any]: + return { + "id": f"{edge_type}::{source}::{target}", + "source": source, + "target": target, + "type": edge_type, + "order": order, + } + + +def build_graph_payload( + document_json: str, + *, + doc_id: str, + title: str | None = None, + max_pages: int = 200, +) -> GraphPayload: + """Build a `GraphPayload` equivalent to `fetch_graph(neo4j, doc_id)` from + the raw `DoclingDocument` JSON. + + Returns `truncated=True` with empty node/edge lists beyond `max_pages`, so + the caller can mirror the Neo4j endpoint's 413 behavior. + """ + doc_data = json.loads(document_json) + + pages_raw = list(iter_pages(doc_data)) + page_count = len(pages_raw) + if page_count > max_pages: + return GraphPayload( + doc_id=doc_id, + nodes=[], + edges=[], + node_count=0, + edge_count=0, + truncated=True, + page_count=page_count, + ) + + nodes: list[dict[str, Any]] = [] + edges: list[dict[str, Any]] = [] + + doc_node_id = f"doc::{doc_id}" + nodes.append( + { + "id": doc_node_id, + "group": "document", + "doc_id": doc_id, + "title": title, + # `stages_applied` is a Neo4j-only artifact; keep the key present + # for shape parity but leave it empty since SQLite doesn't track it. + "stages_applied": [], + } + ) + + # Page nodes. + for p in pages_raw: + nodes.append(_page_node(doc_id, p)) + + # Element nodes + collect parent/body metadata for edges below. The + # `element_idx` mirrors TreeWriter's `enumerate(elements)` so PARENT_OF + # carries the same `order` the Neo4j projection does. + by_ref: dict[str, dict[str, Any]] = {} + element_idx = 0 + for _, item in iter_items(doc_data): + ref = item.get("self_ref") + if not ref: + continue + by_ref[ref] = item + provs = iter_provs(item) + nodes.append(_element_node(doc_id, item, provs)) + + pref = parent_ref(item) + if pref == "#/body": + edges.append(_edge(doc_node_id, f"elem::{ref}", "HAS_ROOT")) + elif pref: + edges.append(_edge(f"elem::{pref}", f"elem::{ref}", "PARENT_OF", order=element_idx)) + + # ON_PAGE, dedup'd per (element, page) — matches the Neo4j query's + # DISTINCT projection through Provenance. + seen_pages: set[int] = set() + for prov in provs: + page_no = prov.get("page_no") + if page_no is None or page_no in seen_pages: + continue + seen_pages.add(page_no) + edges.append(_edge(f"elem::{ref}", f"page::{page_no}", "ON_PAGE")) + + element_idx += 1 + + # NEXT chain (DFS pre-order from body). + for a, b in pairwise(dfs_order(doc_data)): + if a in by_ref and b in by_ref: + edges.append(_edge(f"elem::{a}", f"elem::{b}", "NEXT")) + + return GraphPayload( + doc_id=doc_id, + nodes=nodes, + edges=edges, + node_count=len(nodes), + edge_count=len(edges), + truncated=False, + page_count=page_count, + ) diff --git a/document-parser/infra/docling_tree.py b/document-parser/infra/docling_tree.py new file mode 100644 index 0000000..dfed869 --- /dev/null +++ b/document-parser/infra/docling_tree.py @@ -0,0 +1,136 @@ +"""Pure helpers over a serialized `DoclingDocument` dict. + +No I/O, no Neo4j. Shared between: +- `infra.neo4j.tree_writer` — persists the tree into Neo4j during the Maintain + step (IngestionPipeline). +- `infra.docling_graph` — builds an in-memory `GraphPayload` from the SQLite + `document_json` blob for the reasoning-trace viewer. + +Keep this module the single source of truth for how we read Docling's own +structure, so the two consumers can't drift. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +# Docling label -> specific Neo4j/Cytoscape label. Every element carries the +# generic :Element tag too. Kept 1:1 with docling-core's label taxonomy so the +# projection is a faithful mirror of the DoclingDocument. +LABEL_MAP: dict[str, str] = { + "section_header": "SectionHeader", + "title": "SectionHeader", + "paragraph": "Paragraph", + "text": "Paragraph", + "list_item": "ListItem", + "list": "List", # distinct from :ListItem — a list is a container + "table": "Table", + "picture": "Figure", + "formula": "Formula", + "code": "Code", + "caption": "Caption", + "footnote": "Footnote", + "page_header": "PageHeader", + "page_footer": "PageFooter", + "key_value_area": "KeyValueArea", + "form_area": "FormArea", + "document_index": "DocumentIndex", +} +DEFAULT_LABEL = "TextElement" + + +def element_label(docling_label: str) -> str: + return LABEL_MAP.get(docling_label.lower(), DEFAULT_LABEL) + + +def iter_items(doc_data: dict[str, Any]) -> Iterator[tuple[str, dict[str, Any]]]: + """Yield every item from texts/tables/pictures/groups with its source list key.""" + for key in ("texts", "tables", "pictures", "groups"): + for item in doc_data.get(key, []) or []: + yield key, item + + +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 iter_provs(item: dict[str, Any]) -> list[dict[str, Any]]: + """Flatten a Docling item's `prov[]` into a list of dict rows. + + A single item may have multiple provs when it spans page breaks or appears + more than once in the layout. The returned dicts carry the original index + under `order` so sequence is preserved. + """ + provs = item.get("prov") or [] + rows: list[dict[str, Any]] = [] + for idx, p in enumerate(provs): + bbox = p.get("bbox") + l_, t_, r_, b_ = 0.0, 0.0, 0.0, 0.0 + if isinstance(bbox, dict): + l_ = float(bbox.get("l", 0.0) or 0.0) + t_ = float(bbox.get("t", 0.0) or 0.0) + r_ = float(bbox.get("r", 0.0) or 0.0) + b_ = float(bbox.get("b", 0.0) or 0.0) + elif isinstance(bbox, (list, tuple)) and len(bbox) >= 4: + l_, t_, r_, b_ = (float(x) for x in bbox[:4]) + coord_origin = (bbox.get("coord_origin") if isinstance(bbox, dict) else None) or "TOPLEFT" + charspan = p.get("charspan") or [] + rows.append( + { + "order": idx, + "page_no": p.get("page_no"), + "bbox_l": l_, + "bbox_t": t_, + "bbox_r": r_, + "bbox_b": b_, + "coord_origin": coord_origin, + "charspan_start": int(charspan[0]) if len(charspan) >= 1 else None, + "charspan_end": int(charspan[1]) if len(charspan) >= 2 else None, + } + ) + return rows + + +def dfs_order(doc_data: dict[str, Any]) -> list[str]: + """Return `self_ref`s 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 + + +def iter_pages(doc_data: dict[str, Any]) -> Iterator[dict[str, Any]]: + """Yield page dicts with `page_no`, `width`, `height` from the `pages` map.""" + 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 or {}).get("size") or {} + yield { + "page_no": page_no, + "width": size.get("width"), + "height": size.get("height"), + } diff --git a/document-parser/infra/neo4j/queries.py b/document-parser/infra/neo4j/queries.py index 3c352af..bfe53ac 100644 --- a/document-parser/infra/neo4j/queries.py +++ b/document-parser/infra/neo4j/queries.py @@ -25,9 +25,31 @@ class GraphPayload: # block contributes a single row — avoids the cartesian product that chained # OPTIONAL MATCH on 6+ edge types would produce (hangs on multi-page docs). # See: https://neo4j.com/developer/kb/using-subqueries-to-control-the-scope-of-aggregations/ +# +# Provenance nodes (post-v0.6 refactor) are NOT returned as top-level graph +# nodes — they're metadata of their owning Element. We aggregate them inline +# per element, and derive a dedup'd ON_PAGE edge set from them. _FETCH_GRAPH = """ MATCH (d:Document {id: $doc_id}) -CALL { WITH d MATCH (e:Element {doc_id: d.id}) RETURN collect(e) AS elements } +CALL { + WITH d + MATCH (e:Element {doc_id: d.id}) + OPTIONAL MATCH (e)-[hp:HAS_PROV]->(pv:Provenance) + WITH e, pv ORDER BY hp.order + WITH e, + collect( + CASE WHEN pv IS NULL THEN NULL ELSE { + order: pv.prov_order, + page_no: pv.page_no, + bbox_l: pv.bbox_l, bbox_t: pv.bbox_t, + bbox_r: pv.bbox_r, bbox_b: pv.bbox_b, + coord_origin: pv.coord_origin, + charspan_start: pv.charspan_start, + charspan_end: pv.charspan_end + } END + ) AS all_provs + RETURN collect({element: e, provs: [p IN all_provs WHERE p IS NOT NULL]}) AS elements +} CALL { WITH d MATCH (p:Page {doc_id: d.id}) RETURN collect(p) AS pages } CALL { WITH d MATCH (c:Chunk {doc_id: d.id}) RETURN collect(c) AS chunks } CALL { @@ -42,7 +64,10 @@ CALL { } CALL { WITH d - MATCH (er:Element {doc_id: d.id})-[:ON_PAGE]->(pr:Page) + // ON_PAGE is stored on Provenance since v0.6; surface it at the Element + // level (dedup'd per Element/Page pair) for the Cytoscape viz. + MATCH (er:Element {doc_id: d.id})-[:HAS_PROV]->(:Provenance)-[:ON_PAGE]->(pr:Page) + WITH DISTINCT er, pr RETURN collect({from: er.self_ref, to: pr.page_no, type: 'ON_PAGE'}) AS on_page_edges } CALL { @@ -66,17 +91,25 @@ RETURN d AS document, elements, pages, chunks, """ -def _element_node(doc_id: str, e: dict[str, Any]) -> dict[str, Any]: +def _element_node( + doc_id: str, e: dict[str, Any], provs: list[dict[str, Any]] | None = None +) -> dict[str, Any]: # Determine the specific element label: Neo4j returns it via labels(e) on the # driver side; when we project nodes via RETURN, the driver wraps them as Node # objects, so we convert below. + first_page: int | None = None + if provs: + # Convenience: the first provenance's page — the old `prov_page` property, + # useful for label rendering in Cytoscape. Full list is in `provs`. + first_page = provs[0].get("page_no") return { "id": f"elem::{e.get('self_ref')}", "group": "element", "docling_label": e.get("docling_label"), "self_ref": e.get("self_ref"), "text": (e.get("text") or "")[:200], - "prov_page": e.get("prov_page"), + "prov_page": first_page, + "provs": provs or [], "level": e.get("level"), "doc_id": doc_id, } @@ -171,11 +204,17 @@ async def fetch_graph( ) # Element nodes, keeping the specific label (:SectionHeader, etc.). - for e in record["elements"] or []: + # Each row is a {element, provs} dict from the CALL above; provs is a list + # of per-provenance dicts in original order. + for row in record["elements"] or []: + if row is None: + continue + e = row.get("element") if isinstance(row, dict) else None if e is None: continue + provs = [p for p in (row.get("provs") or []) if p is not None] labels = [label for label in e.labels if label != "Element"] - node = _element_node(doc_id, dict(e)) + node = _element_node(doc_id, dict(e), provs=provs) node["label"] = labels[0] if labels else "TextElement" nodes.append(node) diff --git a/document-parser/infra/neo4j/schema.py b/document-parser/infra/neo4j/schema.py index 5fba1eb..d35ca54 100644 --- a/document-parser/infra/neo4j/schema.py +++ b/document-parser/infra/neo4j/schema.py @@ -27,6 +27,12 @@ CONSTRAINTS: tuple[str, ...] = ( INDEXES: tuple[str, ...] = ( "CREATE INDEX element_doc IF NOT EXISTS FOR (e:Element) ON (e.doc_id)", "CREATE INDEX chunk_doc IF NOT EXISTS FOR (c:Chunk) ON (c.doc_id)", + # Reasoning tunnel / bbox-highlight: looking up a Provenance by its owner + # element is a hot path (one lookup per visited section). Composite index + # avoids a full scan of every Provenance in the DB. + "CREATE INDEX provenance_element IF NOT EXISTS " + "FOR (pv:Provenance) ON (pv.doc_id, pv.element_ref)", + "CREATE INDEX provenance_page IF NOT EXISTS FOR (pv:Provenance) ON (pv.doc_id, pv.page_no)", ) FULLTEXT_INDEXES: tuple[str, ...] = ( diff --git a/document-parser/infra/neo4j/tree_writer.py b/document-parser/infra/neo4j/tree_writer.py index 2b08ff7..2d94e5f 100644 --- a/document-parser/infra/neo4j/tree_writer.py +++ b/document-parser/infra/neo4j/tree_writer.py @@ -16,80 +16,41 @@ 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__) -# 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 + provenances_written: int = 0 def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]: - page, bbox = _first_prov(item) + """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 "", - "prov_page": page, - "prov_bbox": bbox, } # Type-specific extras. if "level" in item: @@ -103,32 +64,6 @@ def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]: 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, *, @@ -148,36 +83,28 @@ async def write_document( ingested_at = datetime.now(tz=UTC).isoformat() elements: list[dict[str, Any]] = [] - for _, item in _iter_items(doc_data): + # 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 "") + specific = element_label(item.get("label") or "") elements.append( { "specific_label": specific, - "parent_ref": _parent_ref(item), + "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]] = [] - 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"), - } - ) + pages: list[dict[str, Any]] = [{"doc_id": doc_id, **p} for p in iter_pages(doc_data)] - reading_order = _dfs_order(doc_data) + reading_order = dfs_order(doc_data) async with ( neo.driver.session(database=neo.database) as session, @@ -190,7 +117,9 @@ async def write_document( "DETACH DELETE d, n", doc_id=doc_id, ) - # Also wipe orphan elements/chunks that may still reference this doc. + # 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) @@ -241,8 +170,6 @@ async def write_document( 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 @@ -291,21 +218,48 @@ async def write_document( 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: + # 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}) - MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no}) - MERGE (e)-[:ON_PAGE]->(p) + 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=on_page_rows, + 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. @@ -327,9 +281,15 @@ async def write_document( await tx.commit() logger.info( - "Neo4j: wrote doc %s (%d elements, %d pages)", + "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), ) - 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 b12f832..6eae99c 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -167,6 +167,12 @@ def _build_document_service( async def lifespan(app: FastAPI) -> AsyncIterator[None]: await init_db() document_repo, analysis_repo = _build_repos() + # Exposed on app.state so routers that need direct repo access (e.g. the + # reasoning-graph endpoint, which reads `document_json` from SQLite to + # build the graph without touching Neo4j) can reach them without going + # through a service. + app.state.analysis_repo = analysis_repo + app.state.document_repo = document_repo app.state.neo4j = await _init_neo4j() app.state.analysis_service = _build_analysis_service( document_repo, analysis_repo, neo4j_driver=app.state.neo4j diff --git a/document-parser/persistence/analysis_repo.py b/document-parser/persistence/analysis_repo.py index 676679c..7d0d106 100644 --- a/document-parser/persistence/analysis_repo.py +++ b/document-parser/persistence/analysis_repo.py @@ -73,6 +73,23 @@ class SqliteAnalysisRepository: row = await cursor.fetchone() return _row_to_job(row) if row else None + async def find_latest_completed_by_document(self, document_id: str) -> AnalysisJob | None: + """Latest COMPLETED analysis with a non-null `document_json` for a doc. + + Used by the reasoning-trace tunnel to prime Neo4j from an existing + analysis when the graph doesn't yet exist (e.g. analysis ran before + Neo4j was wired in). + """ + async with get_connection() as db: + cursor = await db.execute( + f"{_SELECT_WITH_DOC} WHERE aj.document_id = ? " + "AND aj.status = 'COMPLETED' AND aj.document_json IS NOT NULL " + "ORDER BY aj.completed_at DESC LIMIT 1", + (document_id,), + ) + row = await cursor.fetchone() + return _row_to_job(row) if row else None + async def update_status(self, job: AnalysisJob) -> None: """Persist all mutable fields of an analysis job (status, results, timestamps).""" async with get_connection() as db: diff --git a/document-parser/tests/neo4j/test_chunk_writer.py b/document-parser/tests/neo4j/test_chunk_writer.py index c5d8d97..ad7e6ff 100644 --- a/document-parser/tests/neo4j/test_chunk_writer.py +++ b/document-parser/tests/neo4j/test_chunk_writer.py @@ -102,8 +102,10 @@ async def test_fetch_graph_returns_full_payload(neo4j_driver): assert groups == {"document", "element", "page", "chunk"} edge_types = {e["type"] for e in payload.edges} - # Every edge kind written by TreeWriter and ChunkWriter should be present. - assert {"HAS_ROOT", "PARENT_OF", "NEXT", "ON_PAGE", "HAS_CHUNK", "DERIVED_FROM"} <= edge_types + # Every edge kind this fixture can produce should be present. + # PARENT_OF is intentionally excluded: all FIXTURE items have + # `parent = #/body`, so they're roots (→ HAS_ROOT) with no nested hierarchy. + assert {"HAS_ROOT", "NEXT", "ON_PAGE", "HAS_CHUNK", "DERIVED_FROM"} <= edge_types async def test_fetch_graph_missing_doc_returns_none(neo4j_driver): diff --git a/document-parser/tests/neo4j/test_tree_writer.py b/document-parser/tests/neo4j/test_tree_writer.py index 3885bbc..3094be1 100644 --- a/document-parser/tests/neo4j/test_tree_writer.py +++ b/document-parser/tests/neo4j/test_tree_writer.py @@ -86,6 +86,8 @@ async def test_write_creates_expected_structure(neo4j_driver): assert result.elements_written == 4 assert result.pages_written == 2 + # Every item in FIXTURE has exactly one prov entry → 4 Provenance nodes. + assert result.provenances_written == 4 async with neo4j_driver.driver.session(database=neo4j_driver.database) as s: assert ( @@ -131,16 +133,28 @@ async def test_write_creates_expected_structure(neo4j_driver): ) == 3 ) - # ON_PAGE: one per element with prov. + # Post-v0.6: ON_PAGE attaches Provenance to Page, not Element directly. + # Traverse through the Provenance node. assert ( await _count( s, - "MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) " + "MATCH (:Element {doc_id: $id})-[:HAS_PROV]->" + "(:Provenance)-[:ON_PAGE]->(:Page {doc_id: $id}) " "RETURN count(*) AS n", id="doc-fixture", ) == 4 ) + # Each element has exactly one Provenance here (single-page fixture). + assert ( + await _count( + s, + "MATCH (e:Element {doc_id: $id})-[:HAS_PROV]->(pv:Provenance) " + "RETURN count(pv) AS n", + id="doc-fixture", + ) + == 4 + ) async def test_rewrite_is_idempotent_replace(neo4j_driver): diff --git a/document-parser/tests/test_chunking.py b/document-parser/tests/test_chunking.py index 552f71f..5b37b7e 100644 --- a/document-parser/tests/test_chunking.py +++ b/document-parser/tests/test_chunking.py @@ -501,25 +501,31 @@ class TestRemoteChunkingPath: markdown="# Title\nParagraph text here.", html="

Title

Paragraph text here.

", pages_json="[]", - document_json=json.dumps({ - "schema_name": "DoclingDocument", - "version": "1.0.0", - "name": "test", - "origin": { - "mimetype": "application/pdf", - "filename": "test.pdf", - "binary_hash": 0, - }, - "furniture": {"self_ref": "#/furniture", "children": [], "content_layer": "furniture"}, - "body": {"self_ref": "#/body", "children": [], "content_layer": "body"}, - "groups": [], - "texts": [], - "pictures": [], - "tables": [], - "key_value_items": [], - "form_items": [], - "pages": {}, - }), + document_json=json.dumps( + { + "schema_name": "DoclingDocument", + "version": "1.0.0", + "name": "test", + "origin": { + "mimetype": "application/pdf", + "filename": "test.pdf", + "binary_hash": 0, + }, + "furniture": { + "self_ref": "#/furniture", + "children": [], + "content_layer": "furniture", + }, + "body": {"self_ref": "#/body", "children": [], "content_layer": "body"}, + "groups": [], + "texts": [], + "pictures": [], + "tables": [], + "key_value_items": [], + "form_items": [], + "pages": {}, + } + ), ) analysis_repo.find_by_id = AsyncMock(return_value=job) analysis_repo.update_chunks = AsyncMock(return_value=True) diff --git a/document-parser/tests/test_docling_graph.py b/document-parser/tests/test_docling_graph.py new file mode 100644 index 0000000..e184db9 --- /dev/null +++ b/document-parser/tests/test_docling_graph.py @@ -0,0 +1,238 @@ +"""Tests for `infra.docling_graph.build_graph_payload`. + +The fixture mirrors the DoclingDocument shape used in +`tests/neo4j/test_tree_writer.py` so any structural drift between the two +consumers (TreeWriter -> Neo4j, builder -> SQLite reasoning-graph) surfaces +immediately. +""" + +from __future__ import annotations + +import json + +from infra.docling_graph import build_graph_payload + +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}, + "prov": [{"page_no": 2, "bbox": {"l": 10, "t": 90, "r": 500, "b": 200}}], + } + ], + "pictures": [], + "groups": [], +} + + +def _ids(items, group=None): + return [i["id"] for i in items if group is None or i.get("group") == group] + + +def _types(edges): + return [e["type"] for e in edges] + + +def test_builds_document_page_and_element_nodes(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture", title="fixture.pdf") + assert payload.doc_id == "doc-fixture" + assert payload.page_count == 2 + assert payload.truncated is False + + assert _ids(payload.nodes, group="document") == ["doc::doc-fixture"] + assert set(_ids(payload.nodes, group="page")) == {"page::1", "page::2"} + assert set(_ids(payload.nodes, group="element")) == { + "elem::#/texts/0", + "elem::#/texts/1", + "elem::#/texts/2", + "elem::#/tables/0", + } + + +def test_element_keeps_docling_label_and_specific_label(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture") + by_id = {n["id"]: n for n in payload.nodes} + section = by_id["elem::#/texts/0"] + para = by_id["elem::#/texts/1"] + table = by_id["elem::#/tables/0"] + + # Specific label mirrors TreeWriter's `_LABEL_MAP`. + assert section["label"] == "SectionHeader" + assert para["label"] == "Paragraph" + assert table["label"] == "Table" + # Docling's original label is preserved (lowercased) for filtering. + assert section["docling_label"] == "section_header" + assert para["docling_label"] == "paragraph" + + +def test_provs_are_carried_on_element_nodes(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture") + by_id = {n["id"]: n for n in payload.nodes} + para = by_id["elem::#/texts/1"] + assert para["prov_page"] == 1 + assert len(para["provs"]) == 1 + assert para["provs"][0]["bbox_l"] == 10.0 + assert para["provs"][0]["page_no"] == 1 + + +def test_has_root_parent_next_and_on_page_edges(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture") + types = _types(payload.edges) + # 4 top-level children => 4 HAS_ROOT. + assert types.count("HAS_ROOT") == 4 + # 4 elements in reading order => 3 NEXT edges. + assert types.count("NEXT") == 3 + # 4 elements each on 1 page => 4 ON_PAGE edges. + assert types.count("ON_PAGE") == 4 + # No PARENT_OF in this flat fixture (all parents are #/body). + assert types.count("PARENT_OF") == 0 + + +def test_on_page_edges_point_to_correct_pages(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture") + on_page = [(e["source"], e["target"]) for e in payload.edges if e["type"] == "ON_PAGE"] + assert ("elem::#/texts/0", "page::1") in on_page + assert ("elem::#/texts/2", "page::2") in on_page + assert ("elem::#/tables/0", "page::2") in on_page + + +def test_on_page_dedups_when_element_has_multiple_provs_same_page(): + # Paragraph with two provs on the same page — we expect ONE ON_PAGE edge. + fixture = { + **FIXTURE, + "texts": [ + { + "self_ref": "#/texts/0", + "parent": {"$ref": "#/body"}, + "label": "paragraph", + "text": "Split across two provs same page", + "prov": [ + {"page_no": 1, "bbox": {"l": 0, "t": 0, "r": 10, "b": 10}}, + {"page_no": 1, "bbox": {"l": 0, "t": 20, "r": 10, "b": 30}}, + ], + }, + ], + "tables": [], + "body": {"self_ref": "#/body", "children": [{"$ref": "#/texts/0"}]}, + } + payload = build_graph_payload(json.dumps(fixture), doc_id="doc-split") + on_page = [e for e in payload.edges if e["type"] == "ON_PAGE"] + assert len(on_page) == 1 + + +def test_parent_of_edges_when_items_are_nested(): + fixture = { + "pages": {"1": {"page_no": 1, "size": {"width": 595, "height": 842}}}, + "body": {"self_ref": "#/body", "children": [{"$ref": "#/texts/0"}]}, + "texts": [ + { + "self_ref": "#/texts/0", + "parent": {"$ref": "#/body"}, + "label": "section_header", + "text": "Chapter", + "level": 1, + "children": [{"$ref": "#/texts/1"}], + "prov": [{"page_no": 1, "bbox": {"l": 0, "t": 0, "r": 10, "b": 10}}], + }, + { + "self_ref": "#/texts/1", + "parent": {"$ref": "#/texts/0"}, + "label": "paragraph", + "text": "Body", + "prov": [{"page_no": 1, "bbox": {"l": 0, "t": 20, "r": 10, "b": 30}}], + }, + ], + "tables": [], + "pictures": [], + "groups": [], + } + payload = build_graph_payload(json.dumps(fixture), doc_id="doc-nested") + parents = [(e["source"], e["target"]) for e in payload.edges if e["type"] == "PARENT_OF"] + assert parents == [("elem::#/texts/0", "elem::#/texts/1")] + roots = [e for e in payload.edges if e["type"] == "HAS_ROOT"] + assert len(roots) == 1 # only the section is a direct child of body + # NEXT follows DFS order: #/texts/0 -> #/texts/1 + nexts = [(e["source"], e["target"]) for e in payload.edges if e["type"] == "NEXT"] + assert nexts == [("elem::#/texts/0", "elem::#/texts/1")] + + +def test_truncated_when_page_count_exceeds_cap(): + fixture = { + **FIXTURE, + "pages": {str(i): {"page_no": i, "size": {"width": 1, "height": 1}} for i in range(1, 12)}, + } + payload = build_graph_payload(json.dumps(fixture), doc_id="doc-big", max_pages=10) + assert payload.truncated is True + assert payload.page_count == 11 + assert payload.nodes == [] + assert payload.edges == [] + + +def test_title_is_surfaced_on_document_node(): + payload = build_graph_payload(json.dumps(FIXTURE), doc_id="doc-fixture", title="My Doc.pdf") + doc_node = next(n for n in payload.nodes if n["group"] == "document") + assert doc_node["title"] == "My Doc.pdf" + + +def test_element_text_is_capped_at_200_chars(): + long = "x" * 500 + fixture = { + "pages": {"1": {"page_no": 1, "size": {"width": 1, "height": 1}}}, + "body": {"self_ref": "#/body", "children": [{"$ref": "#/texts/0"}]}, + "texts": [ + { + "self_ref": "#/texts/0", + "parent": {"$ref": "#/body"}, + "label": "paragraph", + "text": long, + "prov": [{"page_no": 1}], + } + ], + "tables": [], + "pictures": [], + "groups": [], + } + payload = build_graph_payload(json.dumps(fixture), doc_id="doc-long") + para = next(n for n in payload.nodes if n.get("self_ref") == "#/texts/0") + assert len(para["text"]) == 200 diff --git a/document-parser/tests/test_graph_api.py b/document-parser/tests/test_graph_api.py new file mode 100644 index 0000000..21f28df --- /dev/null +++ b/document-parser/tests/test_graph_api.py @@ -0,0 +1,114 @@ +"""Tests for `api.graph` — the `/graph` (Neo4j) and `/reasoning-graph` +(SQLite) endpoints. Neo4j itself is not exercised here; `/graph` is covered +by the integration tests under `tests/neo4j/`. This file focuses on the +SQLite-backed reasoning endpoint and the error paths. +""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from api.graph import router +from domain.models import AnalysisJob + +FIXTURE = { + "pages": {"1": {"page_no": 1, "size": {"width": 595, "height": 842}}}, + "body": {"self_ref": "#/body", "children": [{"$ref": "#/texts/0"}]}, + "texts": [ + { + "self_ref": "#/texts/0", + "parent": {"$ref": "#/body"}, + "label": "section_header", + "text": "Hello", + "level": 1, + "prov": [{"page_no": 1, "bbox": {"l": 0, "t": 0, "r": 10, "b": 10}}], + } + ], + "tables": [], + "pictures": [], + "groups": [], +} + + +def _job_with_doc_json() -> AnalysisJob: + job = AnalysisJob(document_id="doc-1") + job.document_filename = "hello.pdf" + job.mark_running() + job.mark_completed( + markdown="# Hello", + html="

Hello

", + pages_json="[]", + document_json=json.dumps(FIXTURE), + chunks_json="[]", + ) + return job + + +@pytest.fixture +def mock_analysis_repo() -> AsyncMock: + repo = AsyncMock() + repo.find_latest_completed_by_document.return_value = _job_with_doc_json() + return repo + + +@pytest.fixture +def client(mock_analysis_repo: AsyncMock) -> TestClient: + app = FastAPI() + app.include_router(router) + app.state.analysis_repo = mock_analysis_repo + app.state.neo4j = None # /reasoning-graph must not need Neo4j + return TestClient(app) + + +class TestReasoningGraph: + def test_returns_payload_built_from_sqlite_json(self, client: TestClient) -> None: + resp = client.get("/api/documents/doc-1/reasoning-graph") + assert resp.status_code == 200 + data = resp.json() + assert data["doc_id"] == "doc-1" + assert data["page_count"] == 1 + assert data["truncated"] is False + + groups = {n["group"] for n in data["nodes"]} + assert groups == {"document", "page", "element"} + + edge_types = {e["type"] for e in data["edges"]} + # HAS_ROOT + ON_PAGE expected; NEXT absent (single element so no chain). + assert edge_types == {"HAS_ROOT", "ON_PAGE"} + + def test_404_when_no_completed_analysis( + self, client: TestClient, mock_analysis_repo: AsyncMock + ) -> None: + mock_analysis_repo.find_latest_completed_by_document.return_value = None + resp = client.get("/api/documents/doc-1/reasoning-graph") + assert resp.status_code == 404 + + def test_404_when_analysis_has_no_document_json( + self, client: TestClient, mock_analysis_repo: AsyncMock + ) -> None: + job = AnalysisJob(document_id="doc-1") + job.mark_running() + job.mark_completed( + markdown="", html="", pages_json="[]", document_json=None, chunks_json="[]" + ) + mock_analysis_repo.find_latest_completed_by_document.return_value = job + resp = client.get("/api/documents/doc-1/reasoning-graph") + assert resp.status_code == 404 + + def test_does_not_need_neo4j(self, client: TestClient) -> None: + # `app.state.neo4j = None` and the endpoint still serves — proves the + # reasoning graph is fully decoupled from the Neo4j provider. + resp = client.get("/api/documents/doc-1/reasoning-graph") + assert resp.status_code == 200 + + +class TestPrimeEndpointRemoved: + def test_graph_prime_endpoint_is_gone(self, client: TestClient) -> None: + # Guardrail — if someone reintroduces /graph/prime we want a failing test. + resp = client.post("/api/documents/doc-1/graph/prime") + assert resp.status_code in (404, 405) diff --git a/document-parser/tests/test_repos.py b/document-parser/tests/test_repos.py index 60e01c9..69a606c 100644 --- a/document-parser/tests/test_repos.py +++ b/document-parser/tests/test_repos.py @@ -153,6 +153,49 @@ class TestAnalysisRepo: deleted = await analysis_repo.delete("nonexistent") assert deleted is False + async def test_find_latest_completed_by_document(self, document_repo, analysis_repo): + """Reasoning tunnel helper: latest COMPLETED analysis with document_json.""" + await self._insert_doc(document_repo) + + # Each job must be insert()'d before update_status can touch it. + # Scenarios: pending (excluded — not COMPLETED), old completed without + # document_json (excluded — NULL json), recent completed with + # document_json (the one we want), running (excluded). + pending = AnalysisJob(id="job-pending", document_id="doc-1") + await analysis_repo.insert(pending) + + old_completed = AnalysisJob(id="job-old", document_id="doc-1") + await analysis_repo.insert(old_completed) + old_completed.mark_running() + old_completed.mark_completed(markdown="", html="", pages_json="[]") + await analysis_repo.update_status(old_completed) + + latest = AnalysisJob(id="job-latest", document_id="doc-1") + await analysis_repo.insert(latest) + latest.mark_running() + latest.mark_completed( + markdown="md", + html="

", + pages_json="[]", + document_json='{"body":{"children":[]},"texts":[]}', + ) + await analysis_repo.update_status(latest) + + running = AnalysisJob(id="job-running", document_id="doc-1") + await analysis_repo.insert(running) + running.mark_running() + await analysis_repo.update_status(running) + + found = await analysis_repo.find_latest_completed_by_document("doc-1") + assert found is not None + assert found.id == "job-latest" + assert found.document_json is not None + + async def test_find_latest_completed_by_document_none(self, document_repo, analysis_repo): + await self._insert_doc(document_repo) + found = await analysis_repo.find_latest_completed_by_document("doc-1") + assert found is None + async def test_delete_by_document(self, document_repo, analysis_repo): await self._insert_doc(document_repo) for i in range(3): diff --git a/experiments/reasoning-trace/.gitignore b/experiments/reasoning-trace/.gitignore new file mode 100644 index 0000000..ea1472e --- /dev/null +++ b/experiments/reasoning-trace/.gitignore @@ -0,0 +1 @@ +output/ diff --git a/experiments/reasoning-trace/README.md b/experiments/reasoning-trace/README.md new file mode 100644 index 0000000..65ae766 --- /dev/null +++ b/experiments/reasoning-trace/README.md @@ -0,0 +1,139 @@ +# Reasoning Trace — R&D sandbox + +Goal: run `docling-agent`'s RAG loop against a document already ingested in +Docling-Studio, capture the `RAGResult` (per-iteration reasoning trace), and +inspect what the agent does. + +Fully **isolated** from the Studio backend: no deps added to +`document-parser/`, no services modified. Just a script + uv inline deps. + +--- + +## What it does + +1. Reads the pre-parsed `DoclingDocument` directly from Studio's SQLite + (`analysis_jobs.document_json`) — no PDF re-conversion. +2. Instantiates `DoclingRAGAgent` against a local Ollama model. +3. Calls `agent._rag_loop()` directly (the public `.run()` method discards the + `RAGResult`; we need the iterations to see the reasoning trace). +4. Dumps the full `RAGResult` as JSON to `output/`. + +--- + +## Prerequisites + +### 1. Ollama running +```sh +# If not already running as a service: +ollama serve # in another terminal +``` + +### 2. A model pulled +Recommended (Peter Staar's default, ~3B params, good JSON adherence): +```sh +ollama pull granite4:micro-h +``` + +Alternative already on your machine (2 GB, may struggle with strict JSON +rejection sampling): +``` +llama3.2:3b +``` + +Bigger/more reliable but slower (20B): +```sh +ollama pull gpt-oss:20b +``` + +### 3. Pick an analysis job id +Any `COMPLETED` row from `analysis_jobs` with a non-null `document_json`: + +```sh +sqlite3 document-parser/data/docling_studio.db \ + "SELECT aj.id, d.filename, length(aj.document_json) + FROM analysis_jobs aj JOIN documents d ON d.id=aj.document_id + WHERE aj.document_json IS NOT NULL AND aj.status='COMPLETED' + ORDER BY length(aj.document_json) DESC LIMIT 5;" +``` + +On this machine, the biggest one right now is: +``` +722d5631-0089-44a3-a64a-7ce5b99579d3 — CCI - Conférence IA - Offre Commerciale v1.0 +``` + +--- + +## Run + +```sh +uv run experiments/reasoning-trace/inspect_doc.py \ + --job-id 722d5631-0089-44a3-a64a-7ce5b99579d3 \ + --query "Quels sont les livrables principaux proposés ?" \ + --model granite4:micro-h +``` + +Flags: +- `--job-id` — required, analysis_jobs.id +- `--query` — required, the question +- `--model` — either a mellea catalog constant (`IBM_GRANITE_4_HYBRID_MICRO`) + or a raw Ollama tag (`granite4:micro-h`, `llama3.2:3b`). Default: + `granite4:micro-h`. +- `--max-iters` — default 5 (agent's own default) +- `--quiet` — disable the rich panels during the loop + +First run will take ~1–2 min: `uv` solves the `docling-agent` env (pulls +docling-core, mellea, pydantic, rich, …) into a cached virtualenv. Subsequent +runs are instant. + +--- + +## Output + +`experiments/reasoning-trace/output/_.json` + +Schema: +```json +{ + "job_id": "…", + "filename": "…", + "query": "…", + "model": { "ollama_name": "…", "hf_model_name": "…" }, + "max_iterations": 5, + "result": { + "answer": "…", + "converged": true, + "iterations": [ + { "iteration": 1, "section_ref": "#/texts/3", + "reason": "…", "section_text_length": 412, + "can_answer": false, "response": "…" }, + … + ] + } +} +``` + +This is the artifact the v1 Studio endpoint (`POST /api/rag/inspect`) will +import — so anything that works here should work there. + +--- + +## Things to check on first run + +- **Do we actually get a trace?** `iterations` list should have ≥ 1 entries + (empty means "no section headers found" fallback — bad sign for the viz idea). +- **Are `section_ref` values `#/texts/N` paths or `#/groups/N`?** Determines + how the resolver walks the tree. +- **Reasoning quality**: does `reason` actually explain the pick, or is it + LLM filler? That affects whether the trace is worth surfacing visually. +- **Convergence rate**: with `max_iters=5`, does a small model converge at all, + or hit the cap and return a partial answer? +- **Latency**: per-iteration wall-clock on your M-series machine with granite4. + +--- + +## Next step (if the above looks promising) + +Resolve each `iteration.section_ref` → `(page_no, bbox)` using the same +`DoclingDocument` that was loaded here. That's the `reasoning_service.py` +resolver described in `docs/design/reasoning-trace.md` §3.2 — implement it in +a second script here (`resolve_trace.py`) before touching Studio. diff --git a/experiments/reasoning-trace/inspect_doc.py b/experiments/reasoning-trace/inspect_doc.py new file mode 100644 index 0000000..0f1846e --- /dev/null +++ b/experiments/reasoning-trace/inspect_doc.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "docling-agent", +# "rich", +# ] +# /// +""" +Run a docling-agent RAG inspection on a Docling-Studio analysis job and dump the RAGResult. + +Bypasses `DoclingRAGAgent.run()` (which discards the RAGResult) and calls the private +`_rag_loop()` directly so we can capture the per-iteration trace. + +Loads the DoclingDocument from Studio's SQLite (`analysis_jobs.document_json`), so no +re-parsing of the PDF is needed — same doc the UI is showing. + +Usage: + uv run experiments/reasoning-trace/inspect_doc.py \\ + --job-id 722d5631-0089-44a3-a64a-7ce5b99579d3 \\ + --query "Quels sont les points clés de l'offre ?" \\ + --model granite4:micro-h + +Output: + experiments/reasoning-trace/output/_.json +""" +from __future__ import annotations + +import argparse +import json +import sqlite3 +import sys +from datetime import datetime, timezone +from pathlib import Path + +from docling_agent.agent.rag import DoclingRAGAgent +from docling_core.types.doc.document import DoclingDocument +from mellea.backends import model_ids as M +from mellea.backends.model_ids import ModelIdentifier + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[1] +DB_PATH = REPO / "document-parser" / "data" / "docling_studio.db" +OUT_DIR = HERE / "output" + + +def load_doc(job_id: str) -> tuple[DoclingDocument, str]: + if not DB_PATH.exists(): + sys.exit(f"SQLite DB not found at {DB_PATH}") + con = sqlite3.connect(DB_PATH) + con.row_factory = sqlite3.Row + row = con.execute( + """ + SELECT aj.document_json, d.filename + FROM analysis_jobs aj + JOIN documents d ON d.id = aj.document_id + WHERE aj.id = ? + """, + (job_id,), + ).fetchone() + con.close() + if row is None: + sys.exit(f"No analysis job with id {job_id}") + if not row["document_json"]: + sys.exit(f"Analysis job {job_id} has no document_json (not completed?)") + return DoclingDocument.model_validate_json(row["document_json"]), row["filename"] + + +def resolve_model(name: str) -> ModelIdentifier: + """Accept either a mellea catalog constant name (e.g. 'IBM_GRANITE_4_HYBRID_MICRO') + or a raw Ollama tag (e.g. 'granite4:micro-h', 'llama3.2:3b').""" + const = getattr(M, name.upper(), None) + if isinstance(const, ModelIdentifier): + return const + return ModelIdentifier(ollama_name=name) + + +def summarize_structure(doc: DoclingDocument) -> str: + from docling_core.types.doc.document import SectionHeaderItem, TitleItem + + headers = [ + item for item, _ in doc.iterate_items() + if isinstance(item, (TitleItem, SectionHeaderItem)) + ] + return ( + f"texts={len(doc.texts)} " + f"tables={len(doc.tables)} " + f"pictures={len(doc.pictures)} " + f"groups={len(doc.groups)} " + f"section_headers={len(headers)}" + ) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--job-id", required=True, help="analysis_jobs.id from Studio SQLite") + p.add_argument("--query", required=True, help="Question to ask the document") + p.add_argument( + "--model", + default="granite4:micro-h", + help="Ollama tag or mellea catalog constant (default: granite4:micro-h)", + ) + p.add_argument("--max-iters", type=int, default=5) + p.add_argument("--quiet", action="store_true", help="disable rich progress panels") + args = p.parse_args() + + print(f"→ Loading DoclingDocument from analysis {args.job_id[:8]}…") + doc, filename = load_doc(args.job_id) + print(f" {filename}") + print(f" {summarize_structure(doc)}") + + model_id = resolve_model(args.model) + print(f"→ Model: ollama={model_id.ollama_name!r} hf={model_id.hf_model_name!r}") + + agent = DoclingRAGAgent( + model_id=model_id, + tools=[], + max_iterations=args.max_iters, + verbose=not args.quiet, + ) + + print(f"→ Running RAG loop (query: {args.query!r})\n") + # Intentional: agent.run() discards the RAGResult. _rag_loop gives us the trace. + result = agent._rag_loop(query=args.query, doc=doc) + + OUT_DIR.mkdir(exist_ok=True) + ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + out_path = OUT_DIR / f"{args.job_id[:8]}_{ts}.json" + payload = { + "job_id": args.job_id, + "filename": filename, + "query": args.query, + "model": { + "ollama_name": model_id.ollama_name, + "hf_model_name": model_id.hf_model_name, + }, + "max_iterations": args.max_iters, + "result": json.loads(result.model_dump_json()), + } + out_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False)) + + print() + print(f"✓ Wrote {out_path.relative_to(REPO)}") + print( + f" converged={result.converged} " + f"iterations={len(result.iterations)} " + f"answer_chars={len(result.answer)}" + ) + if result.iterations: + print(" section_refs visited:", [it.section_ref for it in result.iterations]) + + +if __name__ == "__main__": + main() diff --git a/experiments/reasoning-trace/inspect_graph.py b/experiments/reasoning-trace/inspect_graph.py new file mode 100644 index 0000000..5bc2858 --- /dev/null +++ b/experiments/reasoning-trace/inspect_graph.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "neo4j>=5.20,<6", +# ] +# /// +""" +Compare the Neo4j graph for a given doc_id against the source DoclingDocument +stored in SQLite. Reports whether the graph is "well-formed" — all elements +present, hierarchy intact, no orphans, reading order faithful. + +Use when you want to sanity-check the Neo4j writer before building anything +on top of the graph (e.g. the live RAG overlay). + +Usage: + uv run experiments/reasoning-trace/inspect_graph.py \\ + --doc-id 307ad2ba-93d8-4dfd-8e38-c1ea06d23f0d + +Env: + NEO4J_URI default bolt://localhost:7687 + NEO4J_USER default neo4j + NEO4J_PASSWORD default changeme +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sqlite3 +import sys +from collections import Counter +from pathlib import Path + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[1] +DB_PATH = REPO / "document-parser" / "data" / "docling_studio.db" + + +def _load_doc_json(doc_id: str) -> tuple[str, dict]: + """Return (filename, parsed DoclingDocument dict) for the latest completed + analysis of this doc.""" + if not DB_PATH.exists(): + sys.exit(f"SQLite DB not found at {DB_PATH}") + con = sqlite3.connect(DB_PATH) + con.row_factory = sqlite3.Row + row = con.execute( + """ + SELECT d.filename, aj.document_json + FROM analysis_jobs aj + JOIN documents d ON d.id = aj.document_id + WHERE aj.document_id = ? + AND aj.status = 'COMPLETED' + AND aj.document_json IS NOT NULL + ORDER BY aj.completed_at DESC LIMIT 1 + """, + (doc_id,), + ).fetchone() + con.close() + if not row: + sys.exit(f"No completed analysis with document_json for doc_id={doc_id}") + return row["filename"], json.loads(row["document_json"]) + + +# --------------------------------------------------------------------------- +# Source-side summaries (DoclingDocument JSON) +# --------------------------------------------------------------------------- +_ITEM_LISTS = ("texts", "tables", "pictures", "groups") + + +def _iter_items(doc: dict): + for key in _ITEM_LISTS: + for it in doc.get(key) or []: + yield key, it + + +def _parent_ref(item: dict) -> str | None: + parent = item.get("parent") + if isinstance(parent, dict): + return parent.get("$ref") or parent.get("cref") + return None + + +def _dfs_reading_order(doc: dict) -> list[str]: + by_ref: dict[str, dict] = {} + for _, it in _iter_items(doc): + ref = it.get("self_ref") + if ref: + by_ref[ref] = it + body = doc.get("body") or {} + order: list[str] = [] + + def walk(children): + for ch in children or []: + ref = ch.get("$ref") or ch.get("cref") + if not ref: + continue + order.append(ref) + walk((by_ref.get(ref) or {}).get("children")) + + walk(body.get("children")) + return order + + +def _summarize_source(doc: dict) -> dict: + items_by_kind = Counter() + labels = Counter() + roots_from_body = 0 + with_prov = 0 + no_text_no_caption = 0 + section_headers: list[str] = [] + + for kind, it in _iter_items(doc): + items_by_kind[kind] += 1 + label = (it.get("label") or "").lower() + labels[label] += 1 + if _parent_ref(it) == "#/body": + roots_from_body += 1 + if it.get("prov"): + with_prov += 1 + if not (it.get("text") or it.get("caption")): + no_text_no_caption += 1 + if label in ("section_header", "title"): + section_headers.append(it.get("self_ref") or "") + + pages = doc.get("pages") or {} + reading_order = _dfs_reading_order(doc) + + return { + "filename_in_doc": doc.get("name"), + "items_by_kind": dict(items_by_kind), + "labels": dict(labels), + "total_items": sum(items_by_kind.values()), + "section_headers_count": len(section_headers), + "section_headers_sample": section_headers[:5], + "roots_from_body": roots_from_body, + "items_with_prov": with_prov, + "items_without_prov": sum(items_by_kind.values()) - with_prov, + "items_without_text_nor_caption": no_text_no_caption, + "pages_count": len(pages), + "reading_order_length": len(reading_order), + "reading_order_sample": reading_order[:5], + } + + +# --------------------------------------------------------------------------- +# Graph-side summaries (Neo4j) +# --------------------------------------------------------------------------- +async def _summarize_graph(neo, doc_id: str) -> dict: + async with neo.driver.session(database=neo.database) as s: + + async def q(cypher: str, **params) -> list[dict]: + res = await s.run(cypher, doc_id=doc_id, **params) + return [dict(r) async for r in res] + + doc = await q("MATCH (d:Document {id: $doc_id}) RETURN d.title AS title") + if not doc: + return {"exists": False} + + elements = await q( + "MATCH (e:Element {doc_id: $doc_id}) " + "RETURN labels(e) AS labels, e.docling_label AS docling_label, " + " e.self_ref AS self_ref, e.prov_page AS page" + ) + pages = await q( + "MATCH (p:Page {doc_id: $doc_id}) RETURN p.page_no AS page_no" + ) + + edges_by_type = await q( + """ + MATCH (n {doc_id: $doc_id})-[r]->(m) + WHERE (m:Document OR m:Element OR m:Page OR m:Chunk) + RETURN type(r) AS type, count(r) AS n + """ + ) + # HAS_ROOT starts at :Document, so n.doc_id filter above miss the edge + # (doc has id= not doc_id=); fetch separately. + has_root = await q( + "MATCH (d:Document {id: $doc_id})-[:HAS_ROOT]->(c:Element) " + "RETURN count(c) AS n" + ) + + # Orphan detection: elements with no incoming PARENT_OF and not a root. + orphans = await q( + """ + MATCH (e:Element {doc_id: $doc_id}) + WHERE NOT (()-[:PARENT_OF]->(e)) + AND NOT ((:Document {id: $doc_id})-[:HAS_ROOT]->(e)) + RETURN e.self_ref AS self_ref, e.docling_label AS label + """ + ) + # Elements with prov_page but no ON_PAGE edge. + on_page_missing = await q( + """ + MATCH (e:Element {doc_id: $doc_id}) + WHERE e.prov_page IS NOT NULL AND NOT (e)-[:ON_PAGE]->(:Page) + RETURN count(e) AS n + """ + ) + # Pages with no element attached. + empty_pages = await q( + """ + MATCH (p:Page {doc_id: $doc_id}) + WHERE NOT (:Element {doc_id: $doc_id})-[:ON_PAGE]->(p) + RETURN p.page_no AS page_no + """ + ) + + specific_label_counter: Counter = Counter() + docling_label_counter: Counter = Counter() + for e in elements: + specifics = [lbl for lbl in e["labels"] if lbl != "Element"] + specific_label_counter[specifics[0] if specifics else ""] += 1 + docling_label_counter[e["docling_label"] or ""] += 1 + + edges = {row["type"]: row["n"] for row in edges_by_type} + edges["HAS_ROOT"] = has_root[0]["n"] if has_root else 0 + + return { + "exists": True, + "title": doc[0]["title"], + "element_count": len(elements), + "page_count": len(pages), + "specific_labels": dict(specific_label_counter), + "docling_labels": dict(docling_label_counter), + "edges": edges, + "orphan_elements": orphans, + "elements_with_prov_but_no_on_page": on_page_missing[0]["n"] if on_page_missing else 0, + "empty_pages": [r["page_no"] for r in empty_pages], + } + + +# --------------------------------------------------------------------------- +# Coherence checks +# --------------------------------------------------------------------------- +async def _coherence(neo, doc_id: str, doc: dict) -> dict: + source_refs = {it.get("self_ref") for _, it in _iter_items(doc) if it.get("self_ref")} + async with neo.driver.session(database=neo.database) as s: + res = await s.run( + "MATCH (e:Element {doc_id: $doc_id}) RETURN e.self_ref AS r", + doc_id=doc_id, + ) + graph_refs = {row["r"] async for row in res} + + return { + "source_refs_count": len(source_refs), + "graph_refs_count": len(graph_refs), + "missing_in_graph": sorted(source_refs - graph_refs)[:20], + "missing_in_graph_count": len(source_refs - graph_refs), + "extra_in_graph": sorted(graph_refs - source_refs)[:20], + "extra_in_graph_count": len(graph_refs - source_refs), + "reading_order_source_len": len(_dfs_reading_order(doc)), + } + + +# --------------------------------------------------------------------------- +# Pretty report +# --------------------------------------------------------------------------- +def _section(title: str, payload: dict) -> str: + lines = [f"\n── {title} " + "─" * max(2, 70 - len(title))] + for k, v in payload.items(): + if isinstance(v, (dict, list)): + lines.append(f" {k}: {json.dumps(v, ensure_ascii=False, default=str)[:220]}") + else: + lines.append(f" {k}: {v}") + return "\n".join(lines) + + +def _verdict(source: dict, graph: dict, coherence: dict) -> str: + issues: list[str] = [] + + if not graph.get("exists"): + return "❌ No graph in Neo4j for this doc." + + if coherence["missing_in_graph_count"]: + issues.append( + f"MISSING — {coherence['missing_in_graph_count']} self_ref(s) in " + "the source DoclingDocument are not in Neo4j" + ) + if coherence["extra_in_graph_count"]: + issues.append( + f"EXTRA — {coherence['extra_in_graph_count']} self_ref(s) in the " + "graph are not in the source doc" + ) + + # Edge integrity: + edges = graph.get("edges", {}) + expected_has_root = source["roots_from_body"] + actual_has_root = edges.get("HAS_ROOT", 0) + if actual_has_root != expected_has_root: + issues.append( + f"HAS_ROOT mismatch — source has {expected_has_root} top-level " + f"items, graph has {actual_has_root}" + ) + + # Reading order: NEXT chain should be reading_order_length - 1. + expected_next = max(source["reading_order_length"] - 1, 0) + actual_next = edges.get("NEXT", 0) + if actual_next != expected_next: + issues.append( + f"NEXT chain mismatch — expected {expected_next}, got {actual_next}" + ) + + # ON_PAGE: every item with prov should have one ON_PAGE edge. + expected_on_page = source["items_with_prov"] + actual_on_page = edges.get("ON_PAGE", 0) + if actual_on_page != expected_on_page: + issues.append( + f"ON_PAGE mismatch — {expected_on_page} items have prov, " + f"{actual_on_page} ON_PAGE edges exist" + ) + + if graph.get("elements_with_prov_but_no_on_page"): + issues.append( + f"ON_PAGE broken — {graph['elements_with_prov_but_no_on_page']} " + "elements have prov_page but no ON_PAGE edge" + ) + if graph.get("orphan_elements"): + orphans = graph["orphan_elements"] + issues.append(f"ORPHANS — {len(orphans)} disconnected Element node(s)") + if graph.get("empty_pages"): + issues.append( + f"EMPTY PAGES — pages with no element attached: {graph['empty_pages']}" + ) + + if not issues: + return "✅ Graph is well-formed — matches source doc 1:1." + return "⚠️ Findings:\n " + "\n ".join(f"• {x}" for x in issues) + + +async def main_async() -> None: + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--doc-id", required=True, help="documents.id (not analysis_jobs.id)") + args = p.parse_args() + + filename, doc = _load_doc_json(args.doc_id) + + # Imports deferred so we can patch sys.path from within the script runtime. + sys.path.insert(0, str(REPO / "document-parser")) + from infra.neo4j import close_driver, get_driver + + neo = await get_driver( + os.environ.get("NEO4J_URI", "bolt://localhost:7687"), + os.environ.get("NEO4J_USER", "neo4j"), + os.environ.get("NEO4J_PASSWORD", "changeme"), + ) + try: + source = _summarize_source(doc) + graph = await _summarize_graph(neo, args.doc_id) + coh = await _coherence(neo, args.doc_id, doc) + + print(f"Document: {filename}") + print(f"doc_id: {args.doc_id}") + print(_section("SOURCE (DoclingDocument from SQLite)", source)) + print(_section("GRAPH (Neo4j)", graph)) + print(_section("COHERENCE", coh)) + print() + print(_verdict(source, graph, coh)) + finally: + await close_driver() + + +if __name__ == "__main__": + asyncio.run(main_async()) diff --git a/experiments/reasoning-trace/inspect_semantics.py b/experiments/reasoning-trace/inspect_semantics.py new file mode 100644 index 0000000..1145022 --- /dev/null +++ b/experiments/reasoning-trace/inspect_semantics.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = ["neo4j>=5.20,<6"] +# /// +""" +Second-pass inspector: how well does our Neo4j graph model the *semantics* +of a DoclingDocument — the parts Peter Staar would care about? + +Goes beyond structural 1:1 coverage (inspect_graph.py) to answer: + - Are SectionHeader levels preserved (outline hierarchy intact)? + - Can we reconstruct "section scope" (section → its content) from the graph? + - Are lists vs list-items distinguishable? + - Are figure captions linked to their figures? + - What's the depth of the actual hierarchy vs the flat body? +""" +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[1] + + +async def main_async(doc_id: str) -> None: + sys.path.insert(0, str(REPO / "document-parser")) + from infra.neo4j import close_driver, get_driver + + neo = await get_driver( + os.environ.get("NEO4J_URI", "bolt://localhost:7687"), + os.environ.get("NEO4J_USER", "neo4j"), + os.environ.get("NEO4J_PASSWORD", "changeme"), + ) + try: + async with neo.driver.session(database=neo.database) as s: + + async def q(cypher: str, **params): + res = await s.run(cypher, doc_id=doc_id, **params) + return [dict(r) async for r in res] + + print("=== 1. OUTLINE HIERARCHY (section header levels) ===") + rows = await q( + """ + MATCH (e:SectionHeader {doc_id: $doc_id}) + RETURN e.self_ref AS ref, e.level AS level, substring(e.text, 0, 60) AS text + ORDER BY e.self_ref + """ + ) + for r in rows: + print(f" level={r['level']} {r['ref']:<14} {r['text']}") + levels = [r["level"] for r in rows] + print(f" → distinct levels: {sorted(set(levels))}") + + print("\n=== 2. DIRECT TREE DEPTH via PARENT_OF ===") + rows = await q( + """ + MATCH (root:Element {doc_id: $doc_id}) + WHERE NOT (()-[:PARENT_OF]->(root)) + OPTIONAL MATCH path = (root)-[:PARENT_OF*]->(leaf) + WITH root, max(length(path)) AS depth + RETURN + labels(root) AS labels, + root.docling_label AS docling_label, + coalesce(depth, 0) AS depth + ORDER BY depth DESC + LIMIT 10 + """ + ) + for r in rows: + specific = [l for l in r["labels"] if l != "Element"][0] + print(f" depth={r['depth']} {specific:<15} ({r['docling_label']})") + + print("\n=== 3. SECTION SCOPE (can we infer section content from NEXT?) ===") + # For each section header, walk NEXT until the next section header — + # that's the section's content span as per docling-agent's logic. + rows = await q( + """ + MATCH (sh:SectionHeader {doc_id: $doc_id}) + OPTIONAL MATCH p = (sh)-[:NEXT*]->(next:SectionHeader {doc_id: $doc_id}) + WITH sh, min(length(p)) AS span + RETURN + sh.self_ref AS ref, + sh.level AS level, + coalesce(span - 1, -1) AS items_in_scope_if_span_works, + substring(sh.text, 0, 50) AS title + ORDER BY sh.self_ref + """ + ) + for r in rows: + span = r["items_in_scope_if_span_works"] + label = f"~{span} items" if span >= 0 else "last (unknown span)" + print(f" level={r['level']} {r['ref']:<14} {label:<18} {r['title']}") + + print("\n=== 4. LIST CONTAINER vs LIST ITEM distinction ===") + rows = await q( + """ + MATCH (e:ListItem {doc_id: $doc_id}) + RETURN e.docling_label AS docling_label, count(*) AS n + ORDER BY docling_label + """ + ) + for r in rows: + print(f" docling_label={r['docling_label']:<12} neo4j_label=:ListItem count={r['n']}") + print(" ⚠️ Both 'list' (container) and 'list_item' get :ListItem in Neo4j.") + + print("\n=== 5. FIGURE ↔ CAPTION linkage ===") + captions = await q( + "MATCH (c:Caption {doc_id: $doc_id}) RETURN count(c) AS n" + ) + linked = await q( + """ + MATCH (fig:Figure {doc_id: $doc_id})-[:PARENT_OF]-(c:Caption {doc_id: $doc_id}) + RETURN count(DISTINCT fig) AS figs_with_caption, count(DISTINCT c) AS captions_linked + """ + ) + print( + f" captions={captions[0]['n']} " + f"figures_with_caption={linked[0]['figs_with_caption']} " + f"captions_linked={linked[0]['captions_linked']}" + ) + + print("\n=== 6. TABLE CELL CONTENT — graph-addressable or opaque? ===") + rows = await q( + """ + MATCH (t:Table {doc_id: $doc_id}) + RETURN t.self_ref AS ref, + CASE WHEN t.cells_json IS NOT NULL THEN 'JSON-blob' ELSE 'missing' END + AS cells_mode, + size(coalesce(t.cells_json, '')) AS cells_bytes + """ + ) + for r in rows: + print(f" {r['ref']:<14} cells={r['cells_mode']} ({r['cells_bytes']} bytes)") + if rows: + print(" ⚠️ Cells are a JSON string on the Table node — not queryable as graph nodes.") + + print("\n=== 7. WHAT AN AGENT VISIT LOOKS LIKE (section subgraph) ===") + # Pick the first section header and show what would be highlighted + # if we traversed its scope (via NEXT until next same-or-higher-level section). + rows = await q( + """ + MATCH (sh:SectionHeader {doc_id: $doc_id}) + WITH sh ORDER BY sh.self_ref LIMIT 1 + OPTIONAL MATCH chain = (sh)-[:NEXT*0..50]->(e:Element {doc_id: $doc_id}) + WITH sh, e, length(chain) AS pos + OPTIONAL MATCH (e)<-[:NEXT*0..]-(_stop:SectionHeader) + WHERE _stop <> sh AND _stop.level <= sh.level + WITH sh, e, pos + ORDER BY pos + LIMIT 8 + RETURN pos, labels(e) AS labels, e.docling_label AS kind, + substring(e.text, 0, 60) AS text + """ + ) + for r in rows: + specific = [l for l in r["labels"] if l != "Element"][0] + print(f" pos={r['pos']:<3} {specific:<15} {r['text']}") + print( + " → Visiting a section on the graph shows ONE node. Its 'scope' " + "(the content) must be inferred via NEXT-walk — not materialized as edges." + ) + + finally: + await close_driver() + + +if __name__ == "__main__": + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--doc-id", required=True) + args = p.parse_args() + asyncio.run(main_async(args.doc_id)) diff --git a/experiments/reasoning-trace/prime_neo4j.py b/experiments/reasoning-trace/prime_neo4j.py new file mode 100644 index 0000000..0a1c83f --- /dev/null +++ b/experiments/reasoning-trace/prime_neo4j.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "neo4j>=5.20,<6", +# "python-dotenv>=1.0", +# ] +# /// +""" +Populate Neo4j with the DoclingDocument tree + pages for one or more analysis +jobs, **without re-parsing the PDFs**. Reuses Studio's own TreeWriter so the +graph is byte-identical to what an in-UI analysis would produce. + +Use when: + - Neo4j was brought up after some docs were already analyzed (orphan graphs). + - You want to prime a demo environment from existing SQLite state. + +Usage: + # Single job + uv run experiments/reasoning-trace/prime_neo4j.py \\ + --job-id 722d5631-0089-44a3-a64a-7ce5b99579d3 + + # All completed analyses that have a document_json but no graph yet + uv run experiments/reasoning-trace/prime_neo4j.py --all-missing + +Env (defaults match docker-compose.dev.yml): + NEO4J_URI default bolt://localhost:7687 + NEO4J_USER default neo4j + NEO4J_PASSWORD default changeme +""" +from __future__ import annotations + +import argparse +import asyncio +import os +import sqlite3 +import sys +from datetime import UTC, datetime +from pathlib import Path + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[1] +DB_PATH = REPO / "document-parser" / "data" / "docling_studio.db" + +# Studio's own TreeWriter lives in document-parser/infra/neo4j. Import it by +# adding document-parser to sys.path — this keeps us byte-identical with what +# the live backend writes, instead of re-implementing the walk. +sys.path.insert(0, str(REPO / "document-parser")) + + +def _fetch_row(job_id: str) -> tuple[str, str, str] | None: + con = sqlite3.connect(DB_PATH) + con.row_factory = sqlite3.Row + row = con.execute( + """ + SELECT aj.document_id, d.filename, aj.document_json + FROM analysis_jobs aj + JOIN documents d ON d.id = aj.document_id + WHERE aj.id = ? AND aj.document_json IS NOT NULL + """, + (job_id,), + ).fetchone() + con.close() + return (row["document_id"], row["filename"], row["document_json"]) if row else None + + +def _fetch_all_completed() -> list[tuple[str, str, str, str]]: + """Latest completed analysis per document that has a document_json.""" + con = sqlite3.connect(DB_PATH) + con.row_factory = sqlite3.Row + rows = con.execute( + """ + SELECT aj.id, aj.document_id, d.filename, aj.document_json + FROM analysis_jobs aj + JOIN documents d ON d.id = aj.document_id + WHERE aj.document_json IS NOT NULL + AND aj.status = 'COMPLETED' + GROUP BY aj.document_id + HAVING MAX(aj.completed_at) + """, + ).fetchall() + con.close() + return [(r["id"], r["document_id"], r["filename"], r["document_json"]) for r in rows] + + +async def prime(job_id: str, doc_id: str, filename: str, document_json: str) -> None: + # Imports deferred until after sys.path is patched. + from infra.neo4j import bootstrap_schema, close_driver, get_driver, write_document + + uri = os.environ.get("NEO4J_URI", "bolt://localhost:7687") + user = os.environ.get("NEO4J_USER", "neo4j") + pwd = os.environ.get("NEO4J_PASSWORD", "changeme") + + neo = await get_driver(uri, user, pwd) + try: + # Schema is idempotent; safe to run every time. + await bootstrap_schema(neo) + result = await write_document( + neo, + doc_id=doc_id, + filename=filename, + document_json=document_json, + ) + print( + f" ✓ {doc_id[:8]} {filename[:40]:<40} " + f"elements={result.element_count} pages={result.page_count}" + ) + finally: + await close_driver() + + +async def main_async() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + g = p.add_mutually_exclusive_group(required=True) + g.add_argument("--job-id", help="analysis_jobs.id to prime") + g.add_argument( + "--all-missing", + action="store_true", + help="prime every completed analysis with a document_json (latest per doc)", + ) + args = p.parse_args() + + if not DB_PATH.exists(): + sys.exit(f"SQLite DB not found at {DB_PATH}") + + started = datetime.now(tz=UTC) + if args.job_id: + row = _fetch_row(args.job_id) + if row is None: + sys.exit(f"No analysis with id {args.job_id} or no document_json") + doc_id, filename, document_json = row + print(f"→ Priming Neo4j for job {args.job_id[:8]} (doc {doc_id[:8]})") + await prime(args.job_id, doc_id, filename, document_json) + else: + rows = _fetch_all_completed() + print(f"→ Priming Neo4j for {len(rows)} document(s)") + for job_id, doc_id, filename, document_json in rows: + try: + await prime(job_id, doc_id, filename, document_json) + except Exception as e: + print(f" ✗ {doc_id[:8]} {filename[:40]:<40} FAILED: {e}") + + elapsed = (datetime.now(tz=UTC) - started).total_seconds() + print(f"Done in {elapsed:.1f}s") + + +if __name__ == "__main__": + asyncio.run(main_async()) diff --git a/frontend/env.d.ts b/frontend/env.d.ts index f1a72ef..bdf131b 100644 --- a/frontend/env.d.ts +++ b/frontend/env.d.ts @@ -7,3 +7,10 @@ declare module '*.vue' { const component: DefineComponent, Record, unknown> export default component } + +// Cytoscape plugins we use in GraphView — they ship no types. Treated as +// opaque plugin objects; the runtime APIs they add to `cy` are called via +// `(cy as any)` at the call site and typed loosely. +declare module 'cytoscape-expand-collapse' +declare module 'cytoscape-navigator' +declare module 'cytoscape-navigator/cytoscape.js-navigator.css' diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 21307ad..6be331f 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -10,6 +10,8 @@ "dependencies": { "cytoscape": "^3.30.0", "cytoscape-dagre": "^2.5.0", + "cytoscape-expand-collapse": "^4.1.1", + "cytoscape-navigator": "^2.0.2", "dompurify": "^3.3.3", "marked": "^17.0.4", "pinia": "^2.3.0", @@ -1866,6 +1868,22 @@ "cytoscape": "^3.2.22" } }, + "node_modules/cytoscape-expand-collapse": { + "version": "4.1.1", + "resolved": "https://registry.npmjs.org/cytoscape-expand-collapse/-/cytoscape-expand-collapse-4.1.1.tgz", + "integrity": "sha512-MI4/GsA6Rf6RRzNR1aCitBLSnxiIKLxvZyCzF+oti/zn/ui1jmf769VcEFAEbjjsAtwteGsTmczI+niCMWJNvA==", + "peerDependencies": { + "cytoscape": "^3.3.0" + } + }, + "node_modules/cytoscape-navigator": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/cytoscape-navigator/-/cytoscape-navigator-2.0.2.tgz", + "integrity": "sha512-TZFBLFWEMW858UOt4rzusOjtDj7YT5vNx2uCwpUuicUYbaWCHHcUROBZWO+hiuSPWpVhvGLFlOq3NBcAVYOAgw==", + "peerDependencies": { + "cytoscape": "^2.6.0 || ^3.0.0" + } + }, "node_modules/dagre": { "version": "0.8.5", "resolved": "https://registry.npmjs.org/dagre/-/dagre-0.8.5.tgz", diff --git a/frontend/package.json b/frontend/package.json index 9c10205..f33564b 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -18,6 +18,7 @@ "dependencies": { "cytoscape": "^3.30.0", "cytoscape-dagre": "^2.5.0", + "cytoscape-expand-collapse": "^4.1.1", "dompurify": "^3.3.3", "marked": "^17.0.4", "pinia": "^2.3.0", @@ -28,9 +29,9 @@ "@eslint/js": "^9.0.0", "@types/cytoscape": "^3.21.4", "@types/cytoscape-dagre": "^2.3.3", - "@vitest/mocker": "^4.1.2", "@types/dompurify": "^3.2.0", "@vitejs/plugin-vue": "^6.0.5", + "@vitest/mocker": "^4.1.2", "eslint": "^9.0.0", "eslint-plugin-vue": "^9.32.0", "prettier": "^3.4.0", diff --git a/frontend/src/app/router/index.ts b/frontend/src/app/router/index.ts index 19f8141..adb9397 100644 --- a/frontend/src/app/router/index.ts +++ b/frontend/src/app/router/index.ts @@ -27,6 +27,22 @@ const routes: RouteRecordRaw[] = [ name: 'search', component: () => import('../../pages/SearchPage.vue'), }, + { + // Reasoning-trace tunnel. Route is always registered; the page shows + // an empty state when the `reasoning` feature flag is off (same pattern + // as /search does for ingestion). + path: '/reasoning', + name: 'reasoning', + component: () => import('../../pages/ReasoningPage.vue'), + }, + { + // Deep-link into a specific document's reasoning workspace, e.g. shared + // by Peter to a teammate. + path: '/reasoning/:docId', + name: 'reasoning-doc', + component: () => import('../../pages/ReasoningPage.vue'), + props: true, + }, { path: '/settings', name: 'settings', diff --git a/frontend/src/features/analysis/graphApi.ts b/frontend/src/features/analysis/graphApi.ts index e8e09fb..a4dece7 100644 --- a/frontend/src/features/analysis/graphApi.ts +++ b/frontend/src/features/analysis/graphApi.ts @@ -1,5 +1,22 @@ import { apiFetch } from '../../shared/api/http' +/** + * A single provenance entry for an element — matches Docling's + * `ProvenanceItem`. One element may have multiple (e.g. a paragraph that + * spans a page break has two entries, one per page). + */ +export interface GraphProvenance { + order: number + page_no: number | null + bbox_l: number + bbox_t: number + bbox_r: number + bbox_b: number + coord_origin: string + charspan_start: number | null + charspan_end: number | null +} + export interface GraphNode { id: string group: 'document' | 'element' | 'page' | 'chunk' @@ -8,6 +25,8 @@ export interface GraphNode { self_ref?: string text?: string prov_page?: number | null + /** Full list of provenances (page + bbox) — used by the bbox-highlight UI. */ + provs?: GraphProvenance[] level?: number | null page_no?: number chunk_index?: number diff --git a/frontend/src/features/analysis/legendFilters.test.ts b/frontend/src/features/analysis/legendFilters.test.ts new file mode 100644 index 0000000..b408887 --- /dev/null +++ b/frontend/src/features/analysis/legendFilters.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, it } from 'vitest' + +import type { GraphNode } from './graphApi' +import { LEGEND_CHIPS, findChip, partitionByHiddenChips } from './legendFilters' + +function node(overrides: Partial): GraphNode { + return { + id: overrides.id ?? 'n1', + group: overrides.group ?? 'element', + ...overrides, + } +} + +describe('LEGEND_CHIPS matching', () => { + it('section matches only SectionHeader elements', () => { + const chip = findChip('section')! + expect(chip.match(node({ group: 'element', label: 'SectionHeader' }))).toBe(true) + expect(chip.match(node({ group: 'element', label: 'Paragraph' }))).toBe(false) + // An element with no label should NOT match section. + expect(chip.match(node({ group: 'element' }))).toBe(false) + }) + + it('paragraph matches Paragraph and TextElement (fallback)', () => { + const chip = findChip('paragraph')! + expect(chip.match(node({ group: 'element', label: 'Paragraph' }))).toBe(true) + // TextElement is the generic fallback Docling-labels that don't have a + // specific mapping end up with. It's visually rendered in the paragraph + // color, so the chip covers both. + expect(chip.match(node({ group: 'element', label: 'TextElement' }))).toBe(true) + expect(chip.match(node({ group: 'element', label: 'Table' }))).toBe(false) + }) + + it('group-based chips use `group` not `label`', () => { + expect(findChip('document')!.match(node({ group: 'document' }))).toBe(true) + expect(findChip('page')!.match(node({ group: 'page', page_no: 1 }))).toBe(true) + expect(findChip('chunk')!.match(node({ group: 'chunk' }))).toBe(true) + // document chip must NOT match elements that happen to have label===undefined + expect(findChip('document')!.match(node({ group: 'element' }))).toBe(false) + }) + + it('exposes a stable list of chips', () => { + // The sidebar CSS targets `.legend-${key}` — renaming a key here would + // silently break the styles. Pin the expected order + keys. + expect(LEGEND_CHIPS.map((c) => c.key)).toEqual([ + 'document', + 'section', + 'paragraph', + 'table', + 'figure', + 'page', + 'chunk', + ]) + }) +}) + +describe('partitionByHiddenChips', () => { + const nodes: GraphNode[] = [ + node({ id: 'd', group: 'document' }), + node({ id: 's1', group: 'element', label: 'SectionHeader' }), + node({ id: 'p1', group: 'element', label: 'Paragraph' }), + node({ id: 'f1', group: 'element', label: 'Figure' }), + node({ id: 'pg1', group: 'page' }), + ] + + it('returns everything in `show` when the hidden set is empty', () => { + const { hide, show } = partitionByHiddenChips(nodes, new Set()) + expect(hide).toHaveLength(0) + expect(show).toHaveLength(nodes.length) + }) + + it('hides the nodes matching the selected chips', () => { + const { hide, show } = partitionByHiddenChips(nodes, new Set(['figure', 'document'])) + expect(hide.map((n) => n.id).sort()).toEqual(['d', 'f1']) + expect(show.map((n) => n.id).sort()).toEqual(['p1', 'pg1', 's1']) + }) + + it('ignores unknown chip keys (no silent breakage)', () => { + const { hide, show } = partitionByHiddenChips(nodes, new Set(['some-new-kind'])) + expect(hide).toHaveLength(0) + expect(show).toHaveLength(nodes.length) + }) + + it('keeps un-classifiable nodes visible even if all chips are off', () => { + // A hypothetical element with an unmapped label (e.g. :Formula is in the + // writer but has no chip yet) must not be hidden when toggling unrelated + // chips — only explicit chip matches cause hiding. + const weird = node({ id: 'w', group: 'element', label: 'Formula' }) + const { hide, show } = partitionByHiddenChips( + [weird, ...nodes], + new Set(['section', 'paragraph']), + ) + expect(hide.map((n) => n.id)).toEqual(['s1', 'p1']) + expect(show.some((n) => n.id === 'w')).toBe(true) + }) +}) diff --git a/frontend/src/features/analysis/legendFilters.ts b/frontend/src/features/analysis/legendFilters.ts new file mode 100644 index 0000000..0d60aea --- /dev/null +++ b/frontend/src/features/analysis/legendFilters.ts @@ -0,0 +1,108 @@ +/** + * Declarative definition of GraphView's legend chips. + * + * Each chip maps to (a) a swatch color shown in the toolbar, and (b) a + * predicate that decides whether a given node is "of that kind". Clicking a + * chip toggles the matching nodes' visibility via the `.hidden` Cytoscape + * class — pure function kept here so it's unit-testable without mounting + * the component or spinning up Cytoscape. + */ + +import type { GraphNode } from './graphApi' + +export interface LegendChip { + /** Stable ID used as the key in the `hiddenChips` Set + as a CSS slug. */ + key: string + /** Display label shown in the chip. Kept English for now — matches + * existing `legend-*` class names that the design stylesheet targets. */ + label: string + /** Swatch color (also used for the CSS class `legend-${key}`). */ + color: string + /** Returns true when a node belongs to this chip's category. */ + match: (node: GraphNode) => boolean +} + +/** + * Legend order matches the chips currently rendered in GraphView's toolbar, + * so swapping to this source of truth is a drop-in replacement. + * + * `kindLabel` is the specific Neo4j label returned by `fetch_graph` + * (e.g. `SectionHeader`, `Paragraph`, `Table`, …). Group-based chips target + * the `group` attribute (document / page / chunk) — only `element` nodes + * have a useful `kindLabel`. + */ +export const LEGEND_CHIPS: LegendChip[] = [ + { + key: 'document', + label: 'Document', + color: '#1E293B', + match: (n) => n.group === 'document', + }, + { + key: 'section', + label: 'Section', + color: '#F97316', + match: (n) => n.group === 'element' && n.label === 'SectionHeader', + }, + { + key: 'paragraph', + label: 'Paragraph', + color: '#3B82F6', + match: (n) => n.group === 'element' && (n.label === 'Paragraph' || n.label === 'TextElement'), + }, + { + key: 'table', + label: 'Table', + color: '#8B5CF6', + match: (n) => n.group === 'element' && n.label === 'Table', + }, + { + key: 'figure', + label: 'Figure', + color: '#22C55E', + match: (n) => n.group === 'element' && n.label === 'Figure', + }, + { + key: 'page', + label: 'Page', + color: '#94A3B8', + match: (n) => n.group === 'page', + }, + { + key: 'chunk', + label: 'Chunk', + color: '#DC2626', + match: (n) => n.group === 'chunk', + }, +] + +/** Lookup by key — tiny helper to keep GraphView concise. */ +export function findChip(key: string): LegendChip | undefined { + return LEGEND_CHIPS.find((c) => c.key === key) +} + +/** + * Partition an iterable of nodes into `{ hide, show }` based on a set of + * chip keys to hide. A node is hidden if it matches any chip in the set. + * Nodes matching no chip at all (e.g. a future new kind we haven't added + * to the legend) are always shown — no silent disappearance. + */ +export function partitionByHiddenChips( + nodes: Iterable, + hiddenChipKeys: ReadonlySet, +): { hide: GraphNode[]; show: GraphNode[] } { + const hide: GraphNode[] = [] + const show: GraphNode[] = [] + for (const n of nodes) { + let matchedHidden = false + for (const key of hiddenChipKeys) { + const chip = findChip(key) + if (chip && chip.match(n)) { + matchedHidden = true + break + } + } + ;(matchedHidden ? hide : show).push(n) + } + return { hide, show } +} diff --git a/frontend/src/features/analysis/sectionParenting.test.ts b/frontend/src/features/analysis/sectionParenting.test.ts new file mode 100644 index 0000000..0fed70c --- /dev/null +++ b/frontend/src/features/analysis/sectionParenting.test.ts @@ -0,0 +1,136 @@ +import { describe, expect, it } from 'vitest' + +import type { GraphEdge, GraphNode } from './graphApi' +import { computeSectionParents, explicitParentMap, mergeParentMaps } from './sectionParenting' + +function el(selfRef: string, label: string): GraphNode { + return { + id: `elem::${selfRef}`, + group: 'element', + label, + self_ref: selfRef, + } +} + +function nextEdge(from: string, to: string): GraphEdge { + return { + id: `NEXT::elem::${from}::elem::${to}`, + source: `elem::${from}`, + target: `elem::${to}`, + type: 'NEXT', + } +} + +function parentEdge(parent: string, child: string): GraphEdge { + return { + id: `PARENT_OF::elem::${parent}::elem::${child}`, + source: `elem::${parent}`, + target: `elem::${child}`, + type: 'PARENT_OF', + } +} + +describe('computeSectionParents', () => { + it('returns empty when no SectionHeader exists', () => { + const nodes = [el('#/texts/0', 'Paragraph'), el('#/texts/1', 'Paragraph')] + const edges = [nextEdge('#/texts/0', '#/texts/1')] + expect(computeSectionParents(nodes, edges)).toEqual(new Map()) + }) + + it('parents every element to the preceding SectionHeader', () => { + // SH ─▶ P ─▶ P ─▶ SH ─▶ P + const nodes = [ + el('#/texts/0', 'SectionHeader'), + el('#/texts/1', 'Paragraph'), + el('#/texts/2', 'Paragraph'), + el('#/texts/3', 'SectionHeader'), + el('#/texts/4', 'Paragraph'), + ] + const edges = [ + nextEdge('#/texts/0', '#/texts/1'), + nextEdge('#/texts/1', '#/texts/2'), + nextEdge('#/texts/2', '#/texts/3'), + nextEdge('#/texts/3', '#/texts/4'), + ] + const parents = computeSectionParents(nodes, edges) + expect(parents.get('elem::#/texts/1')).toBe('elem::#/texts/0') + expect(parents.get('elem::#/texts/2')).toBe('elem::#/texts/0') + expect(parents.get('elem::#/texts/4')).toBe('elem::#/texts/3') + // Headers are never themselves synthetically parented. + expect(parents.has('elem::#/texts/0')).toBe(false) + expect(parents.has('elem::#/texts/3')).toBe(false) + }) + + it('does not override an explicit PARENT_OF', () => { + // SH ─▶ List ─▶ ListItem, with PARENT_OF(List, ListItem) + const nodes = [ + el('#/texts/0', 'SectionHeader'), + el('#/groups/0', 'List'), + el('#/texts/1', 'ListItem'), + ] + const edges: GraphEdge[] = [ + nextEdge('#/texts/0', '#/groups/0'), + nextEdge('#/groups/0', '#/texts/1'), + parentEdge('#/groups/0', '#/texts/1'), + ] + const parents = computeSectionParents(nodes, edges) + // The List falls under the section… + expect(parents.get('elem::#/groups/0')).toBe('elem::#/texts/0') + // …but the ListItem keeps its explicit List parent, not the section. + expect(parents.has('elem::#/texts/1')).toBe(false) + }) + + it('leaves elements preceding the first SectionHeader unparented', () => { + // Page-header stuff can come before any SectionHeader — don't force-parent it. + const nodes = [ + el('#/texts/0', 'PageHeader'), + el('#/texts/1', 'SectionHeader'), + el('#/texts/2', 'Paragraph'), + ] + const edges = [nextEdge('#/texts/0', '#/texts/1'), nextEdge('#/texts/1', '#/texts/2')] + const parents = computeSectionParents(nodes, edges) + expect(parents.has('elem::#/texts/0')).toBe(false) + expect(parents.get('elem::#/texts/2')).toBe('elem::#/texts/1') + }) + + it('handles multiple disconnected NEXT chains deterministically', () => { + // Two chains, each with its own section. Heads are sorted by id so the + // walk order is stable run-to-run. + const nodes = [ + el('#/texts/10', 'SectionHeader'), + el('#/texts/11', 'Paragraph'), + el('#/texts/20', 'SectionHeader'), + el('#/texts/21', 'Paragraph'), + ] + const edges = [nextEdge('#/texts/10', '#/texts/11'), nextEdge('#/texts/20', '#/texts/21')] + const parents = computeSectionParents(nodes, edges) + expect(parents.get('elem::#/texts/11')).toBe('elem::#/texts/10') + expect(parents.get('elem::#/texts/21')).toBe('elem::#/texts/20') + }) +}) + +describe('explicitParentMap + mergeParentMaps', () => { + it('extracts explicit PARENT_OF into a child→parent map', () => { + const edges = [parentEdge('#/groups/0', '#/texts/1')] + const m = explicitParentMap(edges) + expect(m.get('elem::#/texts/1')).toBe('elem::#/groups/0') + }) + + it('explicit wins over synthetic when both map the same child', () => { + const synthetic = new Map([['c', 'synthetic-parent']]) + const explicit = new Map([['c', 'real-parent']]) + const merged = mergeParentMaps(explicit, synthetic) + expect(merged.get('c')).toBe('real-parent') + }) + + it('preserves synthetic entries that have no explicit counterpart', () => { + const synthetic = new Map([ + ['c1', 'section-1'], + ['c2', 'section-1'], + ]) + const explicit = new Map([['c2', 'list-a']]) + const merged = mergeParentMaps(explicit, synthetic) + expect(merged.get('c1')).toBe('section-1') + expect(merged.get('c2')).toBe('list-a') + }) +}) diff --git a/frontend/src/features/analysis/sectionParenting.ts b/frontend/src/features/analysis/sectionParenting.ts new file mode 100644 index 0000000..fe1adab --- /dev/null +++ b/frontend/src/features/analysis/sectionParenting.ts @@ -0,0 +1,111 @@ +/** + * Synthesize "section scope" compound-node parents for GraphView. + * + * The Docling graph is structurally flat — most elements are direct children + * of `#/body` (HAS_ROOT) with no explicit parent relationship to the section + * header that "owns" them in the reader's mental model. To enable the + * expand-collapse UX per section, we derive that ownership here, at render + * time, from the NEXT reading-order chain: + * + * - Walk NEXT from the head of the chain. + * - Every time we hit a SectionHeader, it becomes the "current section". + * - Every non-SectionHeader node thereafter — unless it has an explicit + * PARENT_OF edge (e.g. list_item → list) — gets that section as its + * compound parent. + * + * The function returns a `Map` where both ids are full + * Cytoscape node ids (`elem::`). The caller merges this with any + * existing PARENT_OF to produce the final `data.parent` on each node. + * + * Edge cases handled: + * - No SectionHeader in the doc → returns an empty map (nothing to collapse). + * - Multiple disconnected NEXT chains → each is walked independently. + * - Explicit PARENT_OF on a child → never overridden. + * - SectionHeader that is itself a list_item or nested → its descendants + * still use it as their section anchor; the SectionHeader keeps its own + * explicit parent. + */ + +import type { GraphEdge, GraphNode } from './graphApi' + +export function computeSectionParents( + nodes: readonly GraphNode[], + edges: readonly GraphEdge[], +): Map { + const elementIds = new Set() + const kindById = new Map() + for (const n of nodes) { + if (n.group !== 'element') continue + elementIds.add(n.id) + kindById.set(n.id, n.label) + } + if (elementIds.size === 0) return new Map() + + const explicitParentOf = new Set() + const nextMap = new Map() + const hasIncomingNext = new Set() + + for (const e of edges) { + if (!elementIds.has(e.source) || !elementIds.has(e.target)) continue + if (e.type === 'PARENT_OF') { + explicitParentOf.add(e.target) + } else if (e.type === 'NEXT') { + // Preserve first-seen NEXT edge per source (should be unique anyway). + if (!nextMap.has(e.source)) nextMap.set(e.source, e.target) + hasIncomingNext.add(e.target) + } + } + + // Walk heads deterministically (sorted) so the same graph always produces + // the same parenting map — useful for tests and reproducible debugging. + const heads = [...elementIds].filter((id) => !hasIncomingNext.has(id)).sort() + + const parents = new Map() + const visited = new Set() + + for (const head of heads) { + let currentSection: string | null = null + let node: string | undefined = head + + while (node && !visited.has(node)) { + visited.add(node) + const kind = kindById.get(node) + + if (kind === 'SectionHeader') { + currentSection = node + // The SectionHeader itself does NOT get a synthetic parent — it is + // the anchor. If it has an explicit PARENT_OF, that one remains + // authoritative downstream (the merger respects it). + } else if (currentSection && !explicitParentOf.has(node)) { + parents.set(node, currentSection) + } + + node = nextMap.get(node) + } + } + + return parents +} + +/** + * Merge an explicit PARENT_OF map with synthetic section parents. Explicit + * wins. Produces the final `childId → parentId` map that callers set as + * `data.parent` on Cytoscape nodes. + */ +export function mergeParentMaps( + explicit: ReadonlyMap, + synthetic: ReadonlyMap, +): Map { + const out = new Map(synthetic) + for (const [child, parent] of explicit) out.set(child, parent) + return out +} + +/** Extract the explicit PARENT_OF map from the edge list — convenience. */ +export function explicitParentMap(edges: readonly GraphEdge[]): Map { + const out = new Map() + for (const e of edges) { + if (e.type === 'PARENT_OF') out.set(e.target, e.source) + } + return out +} diff --git a/frontend/src/features/analysis/ui/GraphView.vue b/frontend/src/features/analysis/ui/GraphView.vue index 7589fd0..7a31f9e 100644 --- a/frontend/src/features/analysis/ui/GraphView.vue +++ b/frontend/src/features/analysis/ui/GraphView.vue @@ -18,26 +18,65 @@ {{ payload?.page_count }} pages - Document - Section - Paragraph - Table - Figure - Page - Chunk + -

+
+
+
+ +
+ {{ tooltip.text }} +
+
+ +
diff --git a/frontend/src/features/feature-flags/store.ts b/frontend/src/features/feature-flags/store.ts index 72057d9..d2cb865 100644 --- a/frontend/src/features/feature-flags/store.ts +++ b/frontend/src/features/feature-flags/store.ts @@ -16,7 +16,7 @@ interface HealthResponse { ingestionAvailable?: boolean } -export type FeatureFlag = 'chunking' | 'disclaimer' | 'ingestion' +export type FeatureFlag = 'chunking' | 'disclaimer' | 'ingestion' | 'reasoning' interface FeatureFlagDef { description: string @@ -42,6 +42,13 @@ const featureRegistry: Record = { description: 'OpenSearch ingestion pipeline (embedding + vector indexing)', isEnabled: (ctx) => ctx.ingestionAvailable, }, + reasoning: { + // SQLite-backed (builds the graph from `document_json` on the fly), so no + // server-side gating needed. Kept as a flag so a future deployment can + // still kill-switch the UI if it wants to. + description: 'Reasoning trace tunnel (docling-agent RAGResult viewer)', + isEnabled: () => true, + }, } export const useFeatureFlagStore = defineStore('feature-flags', () => { diff --git a/frontend/src/features/reasoning/api.ts b/frontend/src/features/reasoning/api.ts new file mode 100644 index 0000000..59c5c8a --- /dev/null +++ b/frontend/src/features/reasoning/api.ts @@ -0,0 +1,15 @@ +import { apiFetch } from '../../shared/api/http' +import type { GraphPayload } from '../analysis/graphApi' + +/** + * Fetch the reasoning-trace graph for a document — built on the backend from + * the SQLite `document_json` blob, not Neo4j. This is intentionally decoupled + * from Maintain's richer Neo4j graph: reasoning only needs the structural + * view (sections, parent/child, reading order, pages) to overlay iterations + * onto, and should work even if Neo4j isn't configured. + * + * 404 if no completed analysis with `document_json` exists for the doc. + */ +export function fetchReasoningGraph(docId: string): Promise { + return apiFetch(`/api/documents/${encodeURIComponent(docId)}/reasoning-graph`) +} diff --git a/frontend/src/features/reasoning/graphReasoningOverlay.test.ts b/frontend/src/features/reasoning/graphReasoningOverlay.test.ts new file mode 100644 index 0000000..6b8fa0a --- /dev/null +++ b/frontend/src/features/reasoning/graphReasoningOverlay.test.ts @@ -0,0 +1,245 @@ +import cytoscape from 'cytoscape' +import type { Core } from 'cytoscape' +import { beforeEach, describe, expect, it } from 'vitest' + +import { + REASONING_EDGE_TYPE, + applyReasoningOverlay, + buildDegradedOverlay, + clearReasoningOverlay, + focusIteration, + nodeIdForSectionRef, +} from './graphReasoningOverlay' +import type { RAGResult } from './types' + +function seed(): Core { + // Headless mode — no DOM container needed. + return cytoscape({ + headless: true, + elements: [ + { data: { id: 'elem::#/texts/0', group: 'element' } }, + { data: { id: 'elem::#/texts/3', group: 'element' } }, + { data: { id: 'elem::#/texts/7', group: 'element' } }, + { data: { id: 'elem::#/groups/1', group: 'element' } }, + ], + }) +} + +function result(iterations: Array>): RAGResult { + return { + answer: 'x', + converged: true, + iterations: iterations.map((it, i) => ({ + iteration: i + 1, + section_ref: '', + reason: '', + section_text_length: 0, + can_answer: false, + response: '', + ...it, + })), + } +} + +describe('nodeIdForSectionRef', () => { + it('matches the backend _element_node id format', () => { + // Kept in sync with document-parser/infra/neo4j/queries.py::_element_node. + expect(nodeIdForSectionRef('#/texts/3')).toBe('elem::#/texts/3') + }) +}) + +describe('applyReasoningOverlay', () => { + let cy: Core + beforeEach(() => { + cy = seed() + }) + + it('marks every resolvable section as visited with its iteration order', () => { + const res = applyReasoningOverlay( + cy, + result([{ section_ref: '#/texts/0' }, { section_ref: '#/texts/3' }]), + ) + + expect(res.presentCount).toBe(2) + expect(res.missingCount).toBe(0) + expect(cy.getElementById('elem::#/texts/0').hasClass('visited')).toBe(true) + expect(cy.getElementById('elem::#/texts/0').data('visitOrder')).toBe(1) + expect(cy.getElementById('elem::#/texts/3').data('visitOrder')).toBe(2) + }) + + it('adds a REASONING_NEXT edge between each consecutive visited pair', () => { + applyReasoningOverlay( + cy, + result([ + { section_ref: '#/texts/0' }, + { section_ref: '#/texts/3' }, + { section_ref: '#/texts/7' }, + ]), + ) + + const edges = cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`) + expect(edges.length).toBe(2) + const ids = edges.map((e) => e.id()).sort() + expect(ids).toEqual([ + 'REASONING_NEXT::elem::#/texts/0::elem::#/texts/3', + 'REASONING_NEXT::elem::#/texts/3::elem::#/texts/7', + ]) + }) + + it('reports missing section refs and does not crash on them', () => { + const res = applyReasoningOverlay( + cy, + result([ + { section_ref: '#/texts/0' }, + { section_ref: '#/texts/999' }, // not in graph + { section_ref: '#/texts/3' }, + ]), + ) + + expect(res.presentCount).toBe(2) + expect(res.missingCount).toBe(1) + expect(res.resolved[1].present).toBe(false) + expect(cy.getElementById('elem::#/texts/999').nonempty()).toBe(false) + }) + + it('breaks the arrow chain across a missing iteration (no ghost edges)', () => { + applyReasoningOverlay( + cy, + result([ + { section_ref: '#/texts/0' }, + { section_ref: '#/texts/999' }, // missing + { section_ref: '#/texts/3' }, + ]), + ) + + // Only the chain between present-to-present pairs gets edges. With one + // missing in the middle, we still draw 0→3 because the filter keeps the + // present ones adjacent in the sequence used for edge drawing. Assert it: + // one edge between 0 and 3. + const edges = cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`) + expect(edges.length).toBe(1) + expect(edges[0].data('source')).toBe('elem::#/texts/0') + expect(edges[0].data('target')).toBe('elem::#/texts/3') + }) + + it('is idempotent — re-applying replaces the previous overlay', () => { + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/0' }])) + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/3' }, { section_ref: '#/texts/7' }])) + + expect(cy.getElementById('elem::#/texts/0').hasClass('visited')).toBe(false) + expect(cy.getElementById('elem::#/texts/3').hasClass('visited')).toBe(true) + expect(cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`).length).toBe(1) + }) + + it('dims every non-visited node so the trace pops visually', () => { + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/0' }, { section_ref: '#/texts/3' }])) + + // Visited ones keep their full opacity (no dimmed class). + expect(cy.getElementById('elem::#/texts/0').hasClass('dimmed')).toBe(false) + expect(cy.getElementById('elem::#/texts/3').hasClass('dimmed')).toBe(false) + // Everything else gets dimmed. + expect(cy.getElementById('elem::#/texts/7').hasClass('dimmed')).toBe(true) + expect(cy.getElementById('elem::#/groups/1').hasClass('dimmed')).toBe(true) + }) + + it('does not dim anything when the trace has no resolvable iterations', () => { + // All missing → we don't want to wash out the whole graph for nothing. + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/999' }])) + expect(cy.$('.dimmed').length).toBe(0) + }) + + it('skips dimming entirely when focusMode is off', () => { + applyReasoningOverlay( + cy, + result([{ section_ref: '#/texts/0' }, { section_ref: '#/texts/3' }]), + { focusMode: false }, + ) + // Visited still highlighted... + expect(cy.getElementById('elem::#/texts/0').hasClass('visited')).toBe(true) + // ...but nothing else gets dimmed. + expect(cy.$('.dimmed').length).toBe(0) + }) + + it('does not dim the synthetic REASONING_NEXT edges', () => { + applyReasoningOverlay( + cy, + result([ + { section_ref: '#/texts/0' }, + { section_ref: '#/texts/3' }, + { section_ref: '#/texts/7' }, + ]), + ) + const reasoningEdges = cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`) + expect(reasoningEdges.length).toBe(2) + reasoningEdges.forEach((e) => { + expect(e.hasClass('dimmed')).toBe(false) + }) + }) + + it('preserves the original graph elements (no destructive mutations)', () => { + const beforeIds = cy + .elements() + .map((e) => e.id()) + .sort() + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/0' }])) + clearReasoningOverlay(cy) + const afterIds = cy + .elements() + .map((e) => e.id()) + .sort() + expect(afterIds).toEqual(beforeIds) + }) +}) + +describe('clearReasoningOverlay', () => { + it('removes class, data, dimming, and synthetic edges', () => { + const cy = seed() + applyReasoningOverlay(cy, result([{ section_ref: '#/texts/0' }, { section_ref: '#/texts/3' }])) + clearReasoningOverlay(cy) + + expect(cy.$('.visited').length).toBe(0) + expect(cy.$('.dimmed').length).toBe(0) + expect(cy.getElementById('elem::#/texts/0').data('visitOrder')).toBeUndefined() + expect(cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`).length).toBe(0) + }) + + it('is a no-op when nothing is overlaid', () => { + const cy = seed() + expect(() => clearReasoningOverlay(cy)).not.toThrow() + }) +}) + +describe('buildDegradedOverlay', () => { + // Used when the Neo4j graph isn't loaded — we still want to show the + // reasoning trace cards, just without graph positioning. + it('returns every iteration with present=false', () => { + const out = buildDegradedOverlay( + result([ + { section_ref: '#/texts/0', reason: 'r0', can_answer: false }, + { section_ref: '#/texts/3', reason: 'r1', can_answer: true, response: 'done' }, + ]), + ) + expect(out.resolved).toHaveLength(2) + expect(out.presentCount).toBe(0) + expect(out.missingCount).toBe(2) + expect(out.resolved.every((r) => !r.present)).toBe(true) + expect(out.resolved[0].reason).toBe('r0') + expect(out.resolved[1].canAnswer).toBe(true) + expect(out.resolved[1].response).toBe('done') + expect(out.resolved[0].nodeId).toBe('elem::#/texts/0') + }) + + it('handles an empty iterations list', () => { + const out = buildDegradedOverlay(result([])) + expect(out.resolved).toEqual([]) + expect(out.presentCount).toBe(0) + expect(out.missingCount).toBe(0) + }) +}) + +describe('focusIteration', () => { + it('is a no-op for a missing node', () => { + const cy = seed() + expect(() => focusIteration(cy, 'elem::#/texts/does-not-exist')).not.toThrow() + }) +}) diff --git a/frontend/src/features/reasoning/graphReasoningOverlay.ts b/frontend/src/features/reasoning/graphReasoningOverlay.ts new file mode 100644 index 0000000..17e1a6b --- /dev/null +++ b/frontend/src/features/reasoning/graphReasoningOverlay.ts @@ -0,0 +1,278 @@ +/** + * Pure Cytoscape-manipulation helpers for the reasoning-trace overlay. + * + * Keeping this as a plain TS module (not a component) makes it trivially + * testable against `cytoscape({ headless: true })` and lets the same code + * drive both the v1 static-import UX and the v2 streaming-runner UX. + * + * The overlay is purely visual: it adds a class + a transient `visitOrder` + * data attribute to existing element nodes and injects synthetic + * `REASONING_NEXT` edges between successive visited nodes. None of this is + * persisted to Neo4j. + */ +import type { Core } from 'cytoscape' + +import type { OverlayResult, RAGResult, ResolvedIteration } from './types' + +export const REASONING_EDGE_TYPE = 'REASONING_NEXT' +const VISITED_CLASS = 'visited' +const DIMMED_CLASS = 'dimmed' + +/** + * Build the Cytoscape node id for a given `section_ref`. + * + * Must stay in sync with `document-parser/infra/neo4j/queries.py::_element_node` + * which emits `f"elem::{self_ref}"`. The `#` that lives inside a Docling + * self_ref (e.g. `#/texts/3`) is fine inside an id string — we just always + * look up via `cy.getElementById()` rather than the `#id` selector syntax, + * which would conflict with the leading `#`. + */ +export function nodeIdForSectionRef(sectionRef: string): string { + return `elem::${sectionRef}` +} + +export interface OverlayOptions { + /** + * When true (default), dim every non-visited element and hide the Document + * root — the user sees the trace pop against a muted background. + * When false, the graph keeps its full colors; only visited nodes get the + * orange ring + numbered badge and the REASONING_NEXT arrows are drawn. + * Toggled from the ReasoningPanel header. + */ + focusMode?: boolean +} + +/** + * Apply the overlay for a freshly imported `RAGResult`. + * + * Idempotent: any previous overlay is cleared first. Returns a summary that + * the caller (store / panel) uses to drive the UI (list of iterations, + * missing count, etc). + */ +export function applyReasoningOverlay( + cy: Core, + result: RAGResult, + options: OverlayOptions = {}, +): OverlayResult { + const focusMode = options.focusMode ?? true + clearReasoningOverlay(cy) + + const resolved: ResolvedIteration[] = result.iterations.map((it) => { + const nodeId = nodeIdForSectionRef(it.section_ref) + const node = cy.getElementById(nodeId) + const present = node.nonempty() + if (present) { + node.addClass(VISITED_CLASS) + node.data('visitOrder', it.iteration) + } + return { + iteration: it.iteration, + sectionRef: it.section_ref, + nodeId, + present, + reason: it.reason, + canAnswer: it.can_answer, + response: it.response, + sectionTextLength: it.section_text_length, + } + }) + + // Dim everything that isn't part of the trace. We do this BEFORE injecting + // the synthetic REASONING_NEXT edges so those edges never receive `.dimmed` + // — they're the foreground of the viz. Takes the inspiration from the + // BboxOverlay approach (dim non-highlighted bboxes, keep the active ones + // fully opaque) and applies it to the graph. + // + // The Document root node is hidden outright (via a `display: none` rule on + // `node.dimmed[group = "document"]`) — it sits at the center of the layout + // and adds zero signal to a reasoning trace, only visual noise. Its + // connected `HAS_ROOT` edge is hidden automatically by Cytoscape. + const present = resolved.filter((r) => r.present) + if (focusMode && present.length > 0) { + cy.elements().not(`.${VISITED_CLASS}`).addClass(DIMMED_CLASS) + } + + // Draw trace arrows only between successively-present iterations. Gaps + // (missing nodes) break the chain — we don't draw edges to ghosts. + for (let i = 0; i < present.length - 1; i++) { + const src = present[i] + const tgt = present[i + 1] + cy.add({ + group: 'edges', + data: { + id: `${REASONING_EDGE_TYPE}::${src.nodeId}::${tgt.nodeId}`, + source: src.nodeId, + target: tgt.nodeId, + type: REASONING_EDGE_TYPE, + order: i, + }, + }) + } + + const visitedEles = cy.$(`.${VISITED_CLASS}`) + if (visitedEles.nonempty()) { + // Padding keeps the arrows readable when the trace is one or two nodes. + cy.fit(visitedEles, 80) + } + + return { + resolved, + presentCount: present.length, + missingCount: resolved.length - present.length, + } +} + +/** + * Build a "degraded" overlay result when no Cytoscape instance is available + * (graph failed to load, or graph is empty for this document). Every iteration + * is marked `present: false`; the panel can still render the iteration cards + * so the user sees the reasoning — they just don't get highlights on the graph. + */ +export function buildDegradedOverlay(result: RAGResult): OverlayResult { + const resolved: ResolvedIteration[] = result.iterations.map((it) => ({ + iteration: it.iteration, + sectionRef: it.section_ref, + nodeId: nodeIdForSectionRef(it.section_ref), + present: false, + reason: it.reason, + canAnswer: it.can_answer, + response: it.response, + sectionTextLength: it.section_text_length, + })) + return { + resolved, + presentCount: 0, + missingCount: resolved.length, + } +} + +/** + * Remove every overlay artifact from the graph. Safe to call when nothing + * is overlaid — it becomes a no-op. + */ +export function clearReasoningOverlay(cy: Core): void { + cy.$(`.${VISITED_CLASS}`).forEach((n) => { + n.removeClass(VISITED_CLASS) + n.removeData('visitOrder') + }) + cy.$(`.${DIMMED_CLASS}`).removeClass(DIMMED_CLASS) + cy.$(`edge[type = "${REASONING_EDGE_TYPE}"]`).remove() +} + +/** + * Center the viewport on a single iteration's node and give it a quick + * visual pulse. No-op if the node isn't present (missing iteration). + */ +export function focusIteration(cy: Core, nodeId: string): void { + const node = cy.getElementById(nodeId) + if (!node.nonempty()) return + cy.animate( + { + center: { eles: node }, + zoom: Math.max(cy.zoom(), 1.0), + }, + { duration: 300 }, + ) + // flashClass is a cytoscape-builtin: add the class then remove it after N ms. + node.flashClass('pulse', 800) +} + +/** + * Shape of a single Cytoscape stylesheet block used for the overlay. We keep + * it intentionally loose (`Record` for `style`) because the + * upstream cytoscape types fork selectors by node-vs-edge prefix and flag + * cross-type properties, while the runtime itself accepts everything we emit + * here. + */ +interface StyleBlock { + selector: string + style: Record +} + +/** + * Style rules appended to the existing GraphView stylesheet. Exported so the + * component can spread them into its Cytoscape config. + */ +export const reasoningOverlayStyles: StyleBlock[] = [ + // Non-trace elements get dimmed — BboxOverlay-inspired: keep the colors, + // drop the opacity hard so the trace pops. Node opacity cascades to label + // + border; edges get a stronger fade because they add visual noise. + { + selector: `node.${DIMMED_CLASS}`, + style: { + opacity: 0.18, + 'text-opacity': 0.25, + }, + }, + { + selector: `edge.${DIMMED_CLASS}`, + style: { + opacity: 0.08, + }, + }, + // Hide the Document root node entirely when a trace is active — it sits + // at the center of the dagre layout and is pure noise for reasoning-path + // inspection. Its `HAS_ROOT` edge is hidden automatically by Cytoscape + // because `display: none` cascades to connected edges. + { + selector: `node.${DIMMED_CLASS}[group = "document"]`, + style: { + display: 'none', + }, + }, + // Visited nodes: orange ring + label-less numbered badge above them. + // Explicit opacity: 1 prevents inheritance quirks when the user re-applies + // an overlay on top of an existing one. + { + selector: `node.${VISITED_CLASS}`, + style: { + 'border-color': '#EA580C', + 'border-width': 4, + 'overlay-color': '#EA580C', + 'overlay-opacity': 0.08, + 'overlay-padding': 4, + opacity: 1, + 'z-index': 50, + }, + }, + { + selector: `node.${VISITED_CLASS}[visitOrder]`, + style: { + label: 'data(visitOrder)', + 'text-valign': 'top', + 'text-margin-y': -6, + 'text-background-color': '#EA580C', + 'text-background-opacity': 1, + 'text-background-padding': '3px', + 'text-background-shape': 'roundrectangle', + 'text-border-color': '#FFFFFF', + 'text-border-width': 1.5, + 'text-border-opacity': 1, + color: '#FFFFFF', + 'font-weight': 700, + 'font-size': 12, + 'text-opacity': 1, + }, + }, + { + selector: `edge[type = "${REASONING_EDGE_TYPE}"]`, + style: { + 'line-color': '#EA580C', + 'target-arrow-color': '#EA580C', + 'target-arrow-shape': 'triangle', + 'arrow-scale': 1.4, + 'curve-style': 'bezier', + 'control-point-step-size': 40, + width: 3, + opacity: 0.95, + 'z-index': 99, + }, + }, + { + selector: 'node.pulse', + style: { + 'border-width': 7, + 'border-color': '#F59E0B', + }, + }, +] diff --git a/frontend/src/features/reasoning/store.test.ts b/frontend/src/features/reasoning/store.test.ts new file mode 100644 index 0000000..5d0e25a --- /dev/null +++ b/frontend/src/features/reasoning/store.test.ts @@ -0,0 +1,84 @@ +import { createPinia, setActivePinia } from 'pinia' +import { beforeEach, describe, expect, it } from 'vitest' + +import { parseImportedTrace, useReasoningStore } from './store' + +describe('parseImportedTrace', () => { + const bare = { + answer: 'ok', + converged: true, + iterations: [], + } + + it('accepts a bare RAGResult', () => { + const parsed = parseImportedTrace(bare) + expect(parsed?.result.answer).toBe('ok') + expect(parsed?.envelope).toBeNull() + }) + + it('accepts a sidecar envelope and extracts the result', () => { + const parsed = parseImportedTrace({ + job_id: 'abc', + filename: 'x.pdf', + result: bare, + }) + expect(parsed?.result.answer).toBe('ok') + expect(parsed?.envelope?.job_id).toBe('abc') + }) + + it('rejects shapes that are neither envelope nor bare RAGResult', () => { + expect(parseImportedTrace(null)).toBeNull() + expect(parseImportedTrace('string')).toBeNull() + expect(parseImportedTrace({ foo: 'bar' })).toBeNull() + expect(parseImportedTrace({ answer: 'x' })).toBeNull() // missing converged + iterations + }) +}) + +describe('useReasoningStore', () => { + beforeEach(() => { + setActivePinia(createPinia()) + }) + + it('starts empty', () => { + const s = useReasoningStore() + expect(s.hasTrace).toBe(false) + expect(s.iterations).toEqual([]) + expect(s.presentCount).toBe(0) + }) + + it('setResult populates and resets active iteration', () => { + const s = useReasoningStore() + s.setActiveIteration(3) + s.setResult({ answer: 'a', converged: true, iterations: [] }, null) + expect(s.hasTrace).toBe(true) + expect(s.activeIteration).toBeNull() + }) + + it('toggleFocusMode flips the flag', () => { + const s = useReasoningStore() + expect(s.focusMode).toBe(true) + s.toggleFocusMode() + expect(s.focusMode).toBe(false) + s.toggleFocusMode() + expect(s.focusMode).toBe(true) + }) + + it('reset restores focusMode to the default on', () => { + const s = useReasoningStore() + s.toggleFocusMode() + expect(s.focusMode).toBe(false) + s.reset() + expect(s.focusMode).toBe(true) + }) + + it('reset clears everything including the dialog flag', () => { + const s = useReasoningStore() + s.openImportDialog() + s.setResult({ answer: 'a', converged: true, iterations: [] }, null) + s.setError('boom') + s.reset() + expect(s.hasTrace).toBe(false) + expect(s.error).toBeNull() + expect(s.importDialogOpen).toBe(false) + }) +}) diff --git a/frontend/src/features/reasoning/store.ts b/frontend/src/features/reasoning/store.ts new file mode 100644 index 0000000..b82eaaf --- /dev/null +++ b/frontend/src/features/reasoning/store.ts @@ -0,0 +1,128 @@ +import { defineStore } from 'pinia' +import { computed, ref } from 'vue' + +import type { OverlayResult, RAGResult, ResolvedIteration, SidecarEnvelope } from './types' + +/** + * Parse an arbitrary JSON payload as either: + * - a bare `RAGResult` (what `docling-agent` emits directly), or + * - the sidecar envelope (`{ job_id, filename, query, model, result }`) + * + * Returns `null` if the shape doesn't match either. + */ +export function parseImportedTrace( + raw: unknown, +): { result: RAGResult; envelope: SidecarEnvelope | null } | null { + if (!raw || typeof raw !== 'object') return null + const obj = raw as Record + + // Envelope shape + if (obj.result && typeof obj.result === 'object') { + const result = obj.result as RAGResult + if (isRAGResult(result)) { + return { result, envelope: obj as unknown as SidecarEnvelope } + } + } + // Bare RAGResult + if (isRAGResult(obj as unknown as RAGResult)) { + return { result: obj as unknown as RAGResult, envelope: null } + } + return null +} + +function isRAGResult(x: RAGResult | undefined): boolean { + if (!x || typeof x !== 'object') return false + return ( + typeof x.answer === 'string' && typeof x.converged === 'boolean' && Array.isArray(x.iterations) + ) +} + +export const useReasoningStore = defineStore('reasoning', () => { + const importDialogOpen = ref(false) + const rawResult = ref(null) + const envelope = ref(null) + const overlay = ref(null) + const activeIteration = ref(null) + const error = ref(null) + // Focus mode: when true, non-visited elements are dimmed so the trace pops. + // Default ON because that's the primary value of the feature; user can + // switch it off from the panel to see the trace in the full graph context. + const focusMode = ref(true) + + const hasTrace = computed(() => rawResult.value !== null) + const iterations = computed(() => overlay.value?.resolved ?? []) + const presentCount = computed(() => overlay.value?.presentCount ?? 0) + const missingCount = computed(() => overlay.value?.missingCount ?? 0) + + function openImportDialog(): void { + importDialogOpen.value = true + } + + function closeImportDialog(): void { + importDialogOpen.value = false + } + + /** + * Called by `ImportTraceDialog` once the user has supplied a JSON file. + * Does NOT touch Cytoscape — the `ReasoningPanel` watches `rawResult` and + * reapplies the overlay via `graphReasoningOverlay.applyReasoningOverlay`. + */ + function setResult(result: RAGResult, env: SidecarEnvelope | null): void { + rawResult.value = result + envelope.value = env + error.value = null + activeIteration.value = null + } + + function setOverlay(o: OverlayResult | null): void { + overlay.value = o + } + + function setActiveIteration(n: number | null): void { + activeIteration.value = n + } + + function setError(msg: string | null): void { + error.value = msg + } + + function toggleFocusMode(): void { + focusMode.value = !focusMode.value + } + + /** Full reset — e.g. when the user switches document. */ + function reset(): void { + rawResult.value = null + envelope.value = null + overlay.value = null + activeIteration.value = null + error.value = null + importDialogOpen.value = false + focusMode.value = true + } + + return { + // state + importDialogOpen, + rawResult, + envelope, + overlay, + activeIteration, + error, + focusMode, + // computed + hasTrace, + iterations, + presentCount, + missingCount, + // actions + openImportDialog, + closeImportDialog, + setResult, + setOverlay, + setActiveIteration, + setError, + toggleFocusMode, + reset, + } +}) diff --git a/frontend/src/features/reasoning/types.ts b/frontend/src/features/reasoning/types.ts new file mode 100644 index 0000000..eef6de1 --- /dev/null +++ b/frontend/src/features/reasoning/types.ts @@ -0,0 +1,65 @@ +/** + * Types mirroring the `docling-agent` RAG output. + * + * The JSON imported by the user is produced either by: + * - the R&D sidecar (`experiments/reasoning-trace/inspect_doc.py`), or + * - any external `docling-agent` run that was serialized to JSON. + * + * Since `docling-agent` uses plain pydantic (no alias generator), field names + * are **snake_case** here. This is one of the rare spots in the frontend where + * we don't normalize to camelCase — keeping the shape 1:1 with upstream means + * a schema drift upstream gives us a clean type error rather than silent + * re-mapping. + * + * Source of truth: docling-project/docling-agent @ docling_agent/agent/rag_models.py + */ + +export interface RAGIteration { + iteration: number + section_ref: string + reason: string + section_text_length: number + can_answer: boolean + response: string +} + +export interface RAGResult { + answer: string + iterations: RAGIteration[] + converged: boolean +} + +/** + * Envelope written by the R&D sidecar. The viewer also accepts a bare + * `RAGResult` (see `parseImportedTrace` in the store). + */ +export interface SidecarEnvelope { + job_id?: string + filename?: string + query?: string + model?: { ollama_name?: string | null; hf_model_name?: string | null } + max_iterations?: number + result: RAGResult +} + +/** + * One iteration after matching its `section_ref` against the currently-loaded + * Cytoscape graph. `present=false` means the section ref has no corresponding + * node (doc not through Maintain, or a different version of the doc). + */ +export interface ResolvedIteration { + iteration: number + sectionRef: string + nodeId: string + present: boolean + reason: string + canAnswer: boolean + response: string + sectionTextLength: number +} + +export interface OverlayResult { + resolved: ResolvedIteration[] + presentCount: number + missingCount: number +} diff --git a/frontend/src/features/reasoning/ui/ImportTraceDialog.vue b/frontend/src/features/reasoning/ui/ImportTraceDialog.vue new file mode 100644 index 0000000..db5ce55 --- /dev/null +++ b/frontend/src/features/reasoning/ui/ImportTraceDialog.vue @@ -0,0 +1,340 @@ +