c9359f60e1
ChunkWriter mirrors chunks into Neo4j after OpenSearch indexing, creating
HAS_CHUNK edges and DERIVED_FROM back-references to the source Elements
(via doc_items propagated from the local chunker).
Graph API: GET /api/documents/{id}/graph returns a cytoscape-shaped
payload with nodes + edges for Document / Element / Page / Chunk.
Hard cap at 200 pages returns HTTP 413 per design §8.4.
Frontend: new Graph tab in Studio results, rendered with Cytoscape.js +
dagre layout (lazy-loaded, ~175 KB gz). Legend, node styling per element
label, directional edges styled per edge type.
README gains a Neo4j section with the schema, three demo Cypher
queries, and env vars. Backend tests skip cleanly when the neo4j python
package is not installed locally.
Refs #186
113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
"""ChunkWriter creates Chunk nodes + DERIVED_FROM links.
|
|
|
|
Builds on the tree_writer fixture — writes the tree first so that DERIVED_FROM
|
|
has Elements to link against.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from infra.neo4j import fetch_graph, write_chunks, write_document
|
|
from infra.neo4j.schema import bootstrap_schema
|
|
from tests.neo4j.test_tree_writer import FIXTURE
|
|
|
|
|
|
CHUNKS = [
|
|
{
|
|
"text": "Introduction. First paragraph on page 1.",
|
|
"sourcePage": 1,
|
|
"tokenCount": 8,
|
|
"docItems": [
|
|
{"selfRef": "#/texts/0", "label": "section_header"},
|
|
{"selfRef": "#/texts/1", "label": "paragraph"},
|
|
],
|
|
},
|
|
{
|
|
"text": "Continued on page 2.",
|
|
"sourcePage": 2,
|
|
"tokenCount": 4,
|
|
"docItems": [{"selfRef": "#/texts/2", "label": "paragraph"}],
|
|
"deleted": False,
|
|
},
|
|
# soft-deleted chunk: must be ignored
|
|
{"text": "gone", "deleted": True, "docItems": []},
|
|
]
|
|
|
|
|
|
async def test_write_chunks_and_derived_from(neo4j_driver):
|
|
await bootstrap_schema(neo4j_driver)
|
|
await write_document(
|
|
neo4j_driver,
|
|
doc_id="doc-fixture",
|
|
filename="fixture.pdf",
|
|
document_json=json.dumps(FIXTURE),
|
|
)
|
|
|
|
result = await write_chunks(
|
|
neo4j_driver,
|
|
doc_id="doc-fixture",
|
|
chunks_json=json.dumps(CHUNKS),
|
|
)
|
|
|
|
assert result.chunks_written == 2
|
|
assert result.derived_from_edges == 3
|
|
|
|
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
|
|
count = await (
|
|
await s.run(
|
|
"MATCH (:Document {id: $id})-[:HAS_CHUNK]->(c:Chunk) RETURN count(c) AS n",
|
|
id="doc-fixture",
|
|
)
|
|
).single()
|
|
assert count["n"] == 2
|
|
|
|
# First chunk derives from 2 elements, second from 1.
|
|
for idx, expected in [(0, 2), (1, 1)]:
|
|
cnt = await (
|
|
await s.run(
|
|
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) "
|
|
"RETURN count(e) AS n",
|
|
cid=f"doc-fixture::chunk::{idx}",
|
|
)
|
|
).single()
|
|
assert cnt["n"] == expected
|
|
|
|
stages = await (
|
|
await s.run(
|
|
"MATCH (d:Document {id: $id}) RETURN d.stages_applied AS s", id="doc-fixture"
|
|
)
|
|
).single()
|
|
assert "chunks" in stages["s"]
|
|
|
|
|
|
async def test_fetch_graph_returns_full_payload(neo4j_driver):
|
|
await bootstrap_schema(neo4j_driver)
|
|
await write_document(
|
|
neo4j_driver,
|
|
doc_id="doc-fixture",
|
|
filename="fixture.pdf",
|
|
document_json=json.dumps(FIXTURE),
|
|
)
|
|
await write_chunks(
|
|
neo4j_driver,
|
|
doc_id="doc-fixture",
|
|
chunks_json=json.dumps(CHUNKS),
|
|
)
|
|
|
|
payload = await fetch_graph(neo4j_driver, "doc-fixture")
|
|
assert payload is not None
|
|
assert payload.truncated is False
|
|
assert payload.page_count == 2
|
|
|
|
groups = {n["group"] for n in payload.nodes}
|
|
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
|
|
|
|
|
|
async def test_fetch_graph_missing_doc_returns_none(neo4j_driver):
|
|
await bootstrap_schema(neo4j_driver)
|
|
assert await fetch_graph(neo4j_driver, "no-such-doc") is None
|