danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1
docling-studio/document-parser/services/chunk_service.py
Pier-Jean Malandrino 70972e807a feat(#256): add /api/documents/{id}/chunks routes + ChunkService 
The frontend was wired against /api/documents/{id}/chunks/* (canonical
doc-centric chunkset) but the backend never exposed those routes — the
chunk tab in the doc workspace 404'd. The domain entities (Chunk,
ChunkEdit, ChunkPush) and persistence repos already existed since #205;
what was missing was the service + API layer that connects them.

ChunkService owns all canonical chunkset invariants (sequence ordering,
soft-delete + audit log atomicity) and shares the chunker port with
AnalysisService so chunking strategy stays a single implementation.

AnalysisService grew a duck-typed promoter hook that copies the chunks
of the first successful analysis into the canonical chunkset. The hook
is idempotent so subsequent ad-hoc analyses (Studio / OCR Debug) never
overwrite hand-edited state.

Routes added (all additive, /api/documents prefix):
  GET    /{id}/chunks
  POST   /{id}/chunks
  PATCH  /{id}/chunks/{chunkId}
  DELETE /{id}/chunks/{chunkId}
  POST   /{id}/chunks/{chunkId}/split
  POST   /{id}/chunks/merge
  POST   /{id}/rechunk
  GET    /{id}/tree
  GET    /{id}/diff?store=...
  POST   /{id}/chunks/push
2026-05-07 11:09:04 +02:00

700 lines
25 KiB
Python
Raw Blame History

"""Chunk service — canonical chunk lifecycle for a document (#256).
Sits between the API layer and the chunk / chunk_edit / chunk_push
repositories. Owns the invariants of the canonical chunkset:
- ordering by `sequence` (dense ascending, gaps allowed after split)
- soft-delete (audit log keeps before/after pointers valid)
- atomic mutation + audit row (one ChunkEdit per mutation)
- promotion from the first completed analysis (idempotent)
Re-uses `DocumentChunker` for rechunk (same port that
`AnalysisService.rechunk` uses), so chunking strategy logic is not
duplicated.
"""
from __future__ import annotations
import json
import logging
import uuid
from dataclasses import asdict
from datetime import UTC, datetime
from typing import TYPE_CHECKING
from domain.models import Chunk, ChunkEdit, ChunkPush
from domain.value_objects import (
ChunkBbox,
ChunkDocItem,
ChunkEditAction,
ChunkingOptions,
)
if TYPE_CHECKING:
from domain.ports import (
AnalysisRepository,
ChunkEditRepository,
ChunkPushRepository,
ChunkRepository,
DocumentChunker,
DocumentRepository,
)
from services.ingestion_service import IngestionService
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Errors — carry an http_status hint, mirrors store_service.py convention.
# ---------------------------------------------------------------------------
class ChunkServiceError(Exception):
http_status: int = 400
def __init__(self, message: str, *, http_status: int | None = None):
super().__init__(message)
if http_status is not None:
self.http_status = http_status
class ChunkNotFoundError(ChunkServiceError):
http_status = 404
class DocumentNotFoundError(ChunkServiceError):
http_status = 404
class ChunkConflictError(ChunkServiceError):
http_status = 409
class ChunkValidationError(ChunkServiceError):
http_status = 400
# ---------------------------------------------------------------------------
# Helpers — chunk ↔ dict conversions for audit log + analysis chunks_json.
# ---------------------------------------------------------------------------
def _utcnow() -> datetime:
return datetime.now(UTC)
def _new_id() -> str:
return str(uuid.uuid4())
def _chunk_to_audit_dict(c: Chunk) -> dict:
"""Serializable snapshot for ChunkEdit.before / .after."""
return {
"id": c.id,
"sequence": c.sequence,
"text": c.text,
"headings": list(c.headings),
"sourcePage": c.source_page,
"tokenCount": c.token_count,
"bboxes": [asdict(b) for b in c.bboxes],
"docItems": [asdict(d) for d in c.doc_items],
}
def _bbox_from_dict(d: dict) -> ChunkBbox:
return ChunkBbox(page=d["page"], bbox=list(d["bbox"]))
def _doc_item_from_dict(d: dict) -> ChunkDocItem:
return ChunkDocItem(
self_ref=d.get("selfRef") or d.get("self_ref", ""), label=d.get("label", "")
)
def _analysis_chunk_to_canonical(
document_id: str,
sequence: int,
raw: dict,
) -> Chunk:
"""Convert an entry from `AnalysisJob.chunks_json` (camelCase) into a
canonical `Chunk`. Used by `_promote_from_analysis`."""
return Chunk(
document_id=document_id,
sequence=sequence,
text=raw.get("text", ""),
headings=list(raw.get("headings", [])),
source_page=raw.get("sourcePage"),
bboxes=[_bbox_from_dict(b) for b in raw.get("bboxes", [])],
doc_items=[_doc_item_from_dict(d) for d in raw.get("docItems", [])],
token_count=raw.get("tokenCount"),
)
# ---------------------------------------------------------------------------
# Service
# ---------------------------------------------------------------------------
class ChunkService:
"""Orchestrates canonical chunk operations for a document."""
def __init__(
self,
chunk_repo: ChunkRepository,
chunk_edit_repo: ChunkEditRepository,
chunk_push_repo: ChunkPushRepository,
document_repo: DocumentRepository,
analysis_repo: AnalysisRepository,
chunker: DocumentChunker | None = None,
ingestion_service: IngestionService | None = None,
actor: str = "user",
) -> None:
self._chunks = chunk_repo
self._edits = chunk_edit_repo
self._pushes = chunk_push_repo
self._documents = document_repo
self._analyses = analysis_repo
self._chunker = chunker
self._ingestion = ingestion_service
self._actor = actor
# -- promotion (called by AnalysisService after first successful analysis)
async def promote_from_analysis_if_empty(self, document_id: str, chunks_json: str) -> int:
"""Populate the canonical chunkset from an analysis result, ONLY if
the document has no canonical chunks yet. Idempotent.
Returns the number of chunks promoted (0 if skipped).
"""
if await self._chunks.count_for_document(document_id) > 0:
return 0
try:
raw_chunks = json.loads(chunks_json)
except json.JSONDecodeError:
logger.exception("Invalid chunks_json for doc %s — skipping promotion", document_id)
return 0
if not isinstance(raw_chunks, list) or not raw_chunks:
return 0
canonical = [
_analysis_chunk_to_canonical(document_id, seq, raw)
for seq, raw in enumerate(raw_chunks)
if not raw.get("deleted")
]
if not canonical:
return 0
await self._chunks.insert_many(canonical)
for c in canonical:
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=c.id,
action=ChunkEditAction.INSERT,
actor="system:initial-analysis",
at=_utcnow(),
after=_chunk_to_audit_dict(c),
)
)
logger.info(
"chunk.promote docId=%s count=%d (initial-analysis)", document_id, len(canonical)
)
return len(canonical)
# -- read
async def list_chunks(self, document_id: str) -> list[Chunk]:
await self._require_doc(document_id)
return await self._chunks.find_for_document(document_id)
# -- mutations
async def add_chunk(self, document_id: str, *, text: str, after_id: str | None = None) -> Chunk:
await self._require_doc(document_id)
existing = await self._chunks.find_for_document(document_id)
sequence = self._sequence_after(existing, after_id)
await self._shift_sequences(existing, from_sequence=sequence)
new_chunk = Chunk(
document_id=document_id,
sequence=sequence,
text=text,
created_at=_utcnow(),
updated_at=_utcnow(),
)
await self._chunks.insert(new_chunk)
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=new_chunk.id,
action=ChunkEditAction.INSERT,
actor=self._actor,
at=_utcnow(),
after=_chunk_to_audit_dict(new_chunk),
)
)
logger.info(
"chunk.add docId=%s chunkId=%s sequence=%d", document_id, new_chunk.id, sequence
)
return new_chunk
async def update_chunk(
self,
document_id: str,
chunk_id: str,
*,
text: str | None = None,
headings: list[str] | None = None,
) -> Chunk:
chunk = await self._require_chunk(document_id, chunk_id)
before = _chunk_to_audit_dict(chunk)
if text is not None:
chunk.text = text
if headings is not None:
chunk.headings = list(headings)
chunk.updated_at = _utcnow()
await self._chunks.update(chunk)
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=chunk.id,
action=ChunkEditAction.UPDATE,
actor=self._actor,
at=_utcnow(),
before=before,
after=_chunk_to_audit_dict(chunk),
)
)
logger.info("chunk.update docId=%s chunkId=%s", document_id, chunk.id)
return chunk
async def delete_chunk(self, document_id: str, chunk_id: str) -> None:
chunk = await self._require_chunk(document_id, chunk_id)
before = _chunk_to_audit_dict(chunk)
deleted = await self._chunks.soft_delete(chunk_id, at=_utcnow())
if not deleted:
raise ChunkNotFoundError(f"Chunk not found: {chunk_id}")
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=chunk_id,
action=ChunkEditAction.DELETE,
actor=self._actor,
at=_utcnow(),
before=before,
)
)
logger.info("chunk.delete docId=%s chunkId=%s", document_id, chunk.id)
async def split_chunk(self, document_id: str, chunk_id: str, cursor_offset: int) -> list[Chunk]:
source = await self._require_chunk(document_id, chunk_id)
if cursor_offset <= 0 or cursor_offset >= len(source.text):
raise ChunkValidationError(
f"cursorOffset {cursor_offset} out of range for chunk of length {len(source.text)}"
)
existing = await self._chunks.find_for_document(document_id)
before = _chunk_to_audit_dict(source)
# Both halves inherit headings, source_page, bboxes, doc_items.
# Token counts are unknown post-split; leave as None.
head_text = source.text[:cursor_offset]
tail_text = source.text[cursor_offset:]
head = Chunk(
document_id=document_id,
sequence=source.sequence,
text=head_text,
headings=list(source.headings),
source_page=source.source_page,
bboxes=list(source.bboxes),
doc_items=list(source.doc_items),
)
tail = Chunk(
document_id=document_id,
sequence=source.sequence + 1,
text=tail_text,
headings=list(source.headings),
source_page=source.source_page,
bboxes=list(source.bboxes),
doc_items=list(source.doc_items),
)
# Push subsequent sequences by 1 to make room for `tail`.
await self._shift_sequences(existing, from_sequence=source.sequence + 1)
await self._chunks.soft_delete(source.id, at=_utcnow())
await self._chunks.insert_many([head, tail])
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=source.id,
action=ChunkEditAction.SPLIT,
actor=self._actor,
at=_utcnow(),
before=before,
children=[head.id, tail.id],
)
)
logger.info(
"chunk.split docId=%s sourceId=%s newIds=[%s,%s]",
document_id,
source.id,
head.id,
tail.id,
)
return [head, tail]
async def merge_chunks(self, document_id: str, ids: list[str]) -> Chunk:
if len(ids) < 2:
raise ChunkValidationError("merge requires at least 2 chunk ids")
existing = await self._chunks.find_for_document(document_id)
by_id = {c.id: c for c in existing}
targets = [by_id.get(i) for i in ids]
if any(t is None for t in targets):
missing = [i for i, t in zip(ids, targets, strict=True) if t is None]
raise ChunkNotFoundError(f"Chunks not found: {missing}")
ordered = sorted(targets, key=lambda c: c.sequence)
sequences = [c.sequence for c in ordered]
if sequences != list(range(sequences[0], sequences[0] + len(sequences))):
raise ChunkConflictError("merge requires contiguous chunks (by sequence)")
merged_text = "\n".join(c.text for c in ordered)
bboxes: list[ChunkBbox] = []
doc_items: list[ChunkDocItem] = []
for c in ordered:
bboxes.extend(c.bboxes)
doc_items.extend(c.doc_items)
token_total = sum((c.token_count or 0) for c in ordered) or None
merged = Chunk(
document_id=document_id,
sequence=ordered[0].sequence,
text=merged_text,
headings=list(ordered[0].headings),
source_page=ordered[0].source_page,
bboxes=bboxes,
doc_items=doc_items,
token_count=token_total,
)
for c in ordered:
await self._chunks.soft_delete(c.id, at=_utcnow())
await self._chunks.insert(merged)
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=merged.id,
action=ChunkEditAction.MERGE,
actor=self._actor,
at=_utcnow(),
parents=[c.id for c in ordered],
after=_chunk_to_audit_dict(merged),
)
)
logger.info(
"chunk.merge docId=%s sourceIds=%s newId=%s",
document_id,
[c.id for c in ordered],
merged.id,
)
return merged
async def rechunk_document(self, document_id: str, options: dict | None = None) -> list[Chunk]:
"""Re-run the chunker on the latest completed analysis JSON and
replace the canonical chunkset.
Emits one INSERT edit per new chunk and one DELETE per old chunk
— keeps the audit log within the existing `ChunkEditAction` enum.
"""
await self._require_doc(document_id)
if not self._chunker:
raise ChunkServiceError("Chunker not configured", http_status=503)
job = await self._analyses.find_latest_completed_by_document(document_id)
if not job or not job.document_json:
raise ChunkServiceError(
"No completed analysis with document_json available for rechunk",
http_status=409,
)
chunk_opts = ChunkingOptions(**options) if options else ChunkingOptions()
new_results = await self._chunker.chunk(job.document_json, chunk_opts)
existing = await self._chunks.find_for_document(document_id)
now = _utcnow()
for c in existing:
await self._chunks.soft_delete(c.id, at=now)
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=c.id,
action=ChunkEditAction.DELETE,
actor="system:rechunk",
at=now,
before=_chunk_to_audit_dict(c),
)
)
new_chunks = [
Chunk(
document_id=document_id,
sequence=seq,
text=r.text,
headings=list(r.headings),
source_page=r.source_page,
bboxes=list(r.bboxes),
doc_items=[], # ChunkResult has no doc_items currently
token_count=r.token_count or None,
)
for seq, r in enumerate(new_results)
]
if new_chunks:
await self._chunks.insert_many(new_chunks)
for c in new_chunks:
await self._edits.insert(
ChunkEdit(
id=_new_id(),
document_id=document_id,
chunk_id=c.id,
action=ChunkEditAction.INSERT,
actor="system:rechunk",
at=now,
after=_chunk_to_audit_dict(c),
)
)
logger.info(
"chunk.rechunk docId=%s oldCount=%d newCount=%d",
document_id,
len(existing),
len(new_chunks),
)
return new_chunks
# -- diff (against last push to a store)
async def diff_against_store(self, document_id: str, store_id: str) -> list[dict]:
"""Compare the canonical chunkset to the last push for `store_id`.
Returns a list of `ChunkDiff`-shaped dicts (camelCase) covering:
- canonical chunks not in the last push → status "added"
- canonical chunks updated since last push → status "modified"
- canonical chunks unchanged since push → status "unchanged"
- chunk ids in last push absent from canonical → status "removed"
Coarse-grained — does not produce a textDiff today (follow-up).
"""
await self._require_doc(document_id)
canonical = await self._chunks.find_for_document(document_id)
last_push = await self._pushes.find_latest(document_id, store_id)
if last_push is None:
return [{"chunkId": c.id, "status": "added", "textDiff": None} for c in canonical]
pushed_ids = set(last_push.chunk_ids)
diffs: list[dict] = []
for c in canonical:
if c.id not in pushed_ids:
diffs.append({"chunkId": c.id, "status": "added", "textDiff": None})
continue
if c.updated_at > last_push.pushed_at:
diffs.append({"chunkId": c.id, "status": "modified", "textDiff": None})
else:
diffs.append({"chunkId": c.id, "status": "unchanged", "textDiff": None})
canonical_ids = {c.id for c in canonical}
for cid in pushed_ids - canonical_ids:
diffs.append({"chunkId": cid, "status": "removed", "textDiff": None})
return diffs
# -- push (delegates to IngestionService; per-store dispatch is a follow-up)
async def push_to_store(self, document_id: str, store_id: str) -> dict:
"""Push the canonical chunkset to a store and record a `ChunkPush`.
Today this delegates to the globally-configured `IngestionService`
(single index). Per-store dispatch by slug is a follow-up issue —
the `store_id` argument is recorded on the `ChunkPush` row so the
diff endpoint can answer "what was last pushed to store X" even
once the dispatch lands.
"""
await self._require_doc(document_id)
if self._ingestion is None:
raise ChunkServiceError(
"Ingestion not available (EMBEDDING_URL and OPENSEARCH_URL required)",
http_status=503,
)
doc = await self._documents.find_by_id(document_id)
canonical = await self._chunks.find_for_document(document_id)
if not canonical:
raise ChunkServiceError(
"No canonical chunks to push — run analysis or rechunk first",
http_status=409,
)
chunks_payload = [_chunk_to_ingestion_dict(c) for c in canonical]
chunks_json_payload = json.dumps(chunks_payload)
ingestion_result = await self._ingestion.ingest(
doc_id=document_id,
filename=(doc.filename if doc else document_id),
chunks_json=chunks_json_payload,
)
chunk_ids = [c.id for c in canonical]
push = ChunkPush(
id=_new_id(),
document_id=document_id,
store_id=store_id,
chunkset_hash=_compute_chunkset_hash(canonical),
chunk_ids=chunk_ids,
pushed_at=_utcnow(),
)
await self._pushes.insert(push)
token_total = sum((c.token_count or 0) for c in canonical)
logger.info(
"chunk.push docId=%s store=%s count=%d tokens=%d",
document_id,
store_id,
ingestion_result.chunks_indexed,
token_total,
)
return {
"jobId": push.id,
"summary": {
"embeds": ingestion_result.chunks_indexed,
"tokens": token_total,
},
}
# -- tree (read from latest analysis document_json)
async def get_tree(self, document_id: str) -> list[dict]:
"""Build a doc tree from the latest completed analysis.
Returns a list of `DocTreeNode`-shaped dicts (camelCase). Empty
list if no analysis is available yet — caller decides if that is
an error or just "not parsed yet".
"""
await self._require_doc(document_id)
job = await self._analyses.find_latest_completed_by_document(document_id)
if not job or not job.document_json:
return []
try:
doc_data = json.loads(job.document_json)
except json.JSONDecodeError:
logger.exception("Invalid document_json for analysis %s", job.id)
return []
return _build_tree_nodes(doc_data)
# -- guards
async def _require_doc(self, document_id: str) -> None:
doc = await self._documents.find_by_id(document_id)
if not doc:
raise DocumentNotFoundError(f"Document not found: {document_id}")
async def _require_chunk(self, document_id: str, chunk_id: str) -> Chunk:
chunk = await self._chunks.find_by_id(chunk_id)
if not chunk or chunk.document_id != document_id or chunk.deleted_at is not None:
raise ChunkNotFoundError(f"Chunk not found: {chunk_id}")
return chunk
# -- sequence helpers
@staticmethod
def _sequence_after(existing: list[Chunk], after_id: str | None) -> int:
if after_id is None:
return (max((c.sequence for c in existing), default=-1)) + 1
anchor = next((c for c in existing if c.id == after_id), None)
if anchor is None:
raise ChunkNotFoundError(f"Anchor chunk not found: {after_id}")
return anchor.sequence + 1
async def _shift_sequences(self, existing: list[Chunk], *, from_sequence: int) -> None:
"""Push chunks at >= from_sequence one slot up to make room."""
affected = [c for c in existing if c.sequence >= from_sequence]
for c in affected:
c.sequence += 1
c.updated_at = _utcnow()
await self._chunks.update(c)
# ---------------------------------------------------------------------------
# Tree projection — extract a hierarchical outline from a DoclingDocument.
# Kept module-level so it stays cheap to test in isolation.
# ---------------------------------------------------------------------------
def _chunk_to_ingestion_dict(c: Chunk) -> dict:
"""Convert a canonical `Chunk` into the legacy chunks_json shape that
`IngestionService.ingest` consumes (camelCase, modeled after
`analysis_service._chunk_to_dict`)."""
return {
"text": c.text,
"headings": list(c.headings),
"sourcePage": c.source_page,
"tokenCount": c.token_count or 0,
"bboxes": [{"page": b.page, "bbox": list(b.bbox)} for b in c.bboxes],
"docItems": [{"selfRef": d.self_ref, "label": d.label} for d in c.doc_items],
}
def _compute_chunkset_hash(chunks: list[Chunk]) -> str:
"""Stable hash of the canonical chunkset content.
Used by `ChunkPush` snapshots so we can answer 'is the store in sync
with the current canonical state' without listing chunks from the
vector store.
"""
import hashlib
h = hashlib.sha256()
for c in chunks:
h.update(c.id.encode("utf-8"))
h.update(b"\x00")
h.update(c.text.encode("utf-8"))
h.update(b"\x00")
h.update(str(c.updated_at).encode("utf-8"))
h.update(b"\x01")
return h.hexdigest()
def _build_tree_nodes(doc_data: dict) -> list[dict]:
"""Project a Docling document JSON into a `[DocTreeNode]` outline.
The Docling document layout has top-level lists like `texts`, `tables`,
`pictures`, `groups`, etc. Each entry carries a `self_ref` and a
`label`. We surface a flat outline grouped by label families, which
is enough for the Inspect tab — full hierarchy reconstruction is a
follow-up (#216).
"""
sections = (
("section_header", "Sections", []),
("title", "Titles", []),
("text", "Paragraphs", []),
("table", "Tables", []),
("picture", "Pictures", []),
)
section_map = {label: bucket for label, _, bucket in sections}
for collection in ("texts", "tables", "pictures", "groups"):
for item in doc_data.get(collection, []) or []:
label = item.get("label") or collection.rstrip("s")
ref = item.get("self_ref") or item.get("selfRef") or ""
display = item.get("text") or item.get("name") or label
bucket = section_map.get(label)
if bucket is None:
continue
bucket.append({"ref": ref, "type": label, "label": display, "children": []})
return [
{"ref": f"#group/{key}", "type": "group", "label": title, "children": children}
for key, title, children in sections
if children
]
Reference in a new issue View git blame Copy permalink