c2550867b7
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
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
"""Local Docling chunker — runs chunking in-process using docling-core.
|
|
|
|
This adapter implements the DocumentChunker port. It deserializes a
|
|
DoclingDocument from JSON, applies the requested chunker, and returns
|
|
domain ChunkResult objects.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
|
|
from docling_core.transforms.chunker import HierarchicalChunker
|
|
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
from domain.value_objects import ChunkBbox, ChunkDocItem, ChunkingOptions, ChunkResult
|
|
from infra.bbox import EMPTY_BBOX, to_topleft_list
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _chunk_sync(document_json: str, options: ChunkingOptions) -> list[ChunkResult]:
|
|
if not document_json or not document_json.strip():
|
|
raise ValueError("Empty document JSON — nothing to chunk")
|
|
|
|
try:
|
|
doc_data = json.loads(document_json)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Malformed document JSON: {e}") from e
|
|
|
|
doc = DoclingDocument.model_validate(doc_data)
|
|
|
|
chunker = _build_chunker(options)
|
|
results: list[ChunkResult] = []
|
|
|
|
for chunk in chunker.chunk(doc):
|
|
source_page = None
|
|
token_count = 0
|
|
bboxes: list[ChunkBbox] = []
|
|
doc_items: list[ChunkDocItem] = []
|
|
|
|
if hasattr(chunk, "meta") and chunk.meta and chunk.meta.doc_items:
|
|
for doc_item in chunk.meta.doc_items:
|
|
ref = getattr(doc_item, "self_ref", None)
|
|
if ref:
|
|
doc_items.append(
|
|
ChunkDocItem(
|
|
self_ref=ref,
|
|
label=str(getattr(doc_item, "label", "") or ""),
|
|
)
|
|
)
|
|
if not hasattr(doc_item, "prov") or not doc_item.prov:
|
|
continue
|
|
for prov in doc_item.prov:
|
|
page_no = prov.page_no
|
|
if source_page is None:
|
|
source_page = page_no
|
|
if prov.bbox:
|
|
page_obj = doc.pages.get(page_no)
|
|
if page_obj:
|
|
bbox = to_topleft_list(prov.bbox, page_obj.size.height)
|
|
if bbox != EMPTY_BBOX:
|
|
bboxes.append(ChunkBbox(page=page_no, bbox=bbox))
|
|
|
|
if hasattr(chunker, "tokenizer") and chunker.tokenizer:
|
|
token_count = chunker.tokenizer.count_tokens(chunk.text)
|
|
|
|
headings = list(chunk.meta.headings) if chunk.meta and chunk.meta.headings else []
|
|
|
|
results.append(
|
|
ChunkResult(
|
|
text=chunk.text,
|
|
headings=headings,
|
|
source_page=source_page,
|
|
token_count=token_count,
|
|
bboxes=bboxes,
|
|
doc_items=doc_items,
|
|
)
|
|
)
|
|
|
|
logger.info("Chunked document into %d chunks (chunker=%s)", len(results), options.chunker_type)
|
|
return results
|
|
|
|
|
|
def _build_chunker(options: ChunkingOptions) -> HierarchicalChunker | HybridChunker:
|
|
if options.chunker_type == "hierarchical":
|
|
return HierarchicalChunker()
|
|
|
|
return HybridChunker(
|
|
max_tokens=options.max_tokens,
|
|
merge_peers=options.merge_peers,
|
|
repeat_table_header=options.repeat_table_header,
|
|
)
|
|
|
|
|
|
class LocalChunker:
|
|
"""Adapter that runs docling-core chunking locally."""
|
|
|
|
async def chunk(
|
|
self,
|
|
document_json: str,
|
|
options: ChunkingOptions,
|
|
) -> list[ChunkResult]:
|
|
return await asyncio.to_thread(_chunk_sync, document_json, options)
|