diff --git a/document-parser/api/document_chunks.py b/document-parser/api/document_chunks.py new file mode 100644 index 0000000..65fee49 --- /dev/null +++ b/document-parser/api/document_chunks.py @@ -0,0 +1,212 @@ +"""Document chunks API router (#256). + +Exposes the canonical doc-centric chunkset (CRUD + tree + rechunk + diff ++ push). Distinct from `/api/analyses/{id}/...` which operates on +ephemeral analysis runs (StudioPage / OCR Debug). + +All routes mount under `/api/documents/{doc_id}/...`. The router +delegates to `ChunkService` — no direct repo / persistence access. +""" + +from __future__ import annotations + +import logging +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException, Query, Request + +from api.schemas import ( + AddChunkRequest, + ChunkDiffResponse, + DocChunkResponse, + DocRechunkRequest, + DocTreeNodeResponse, + MergeChunksRequest, + PushChunksRequest, + PushChunksResponse, + SplitChunkRequest, + UpdateChunkRequest, +) +from services.chunk_service import ( + ChunkService, + ChunkServiceError, +) + +logger = logging.getLogger(__name__) +router = APIRouter(prefix="/api/documents", tags=["documents"]) + + +def _get_service(request: Request) -> ChunkService: + svc = getattr(request.app.state, "chunk_service", None) + if svc is None: + raise HTTPException(status_code=503, detail="Chunk service not available") + return svc + + +ServiceDep = Annotated[ChunkService, Depends(_get_service)] + + +def _to_response(chunk) -> DocChunkResponse: + return DocChunkResponse( + id=chunk.id, + doc_id=chunk.document_id, + sequence=chunk.sequence, + text=chunk.text, + headings=list(chunk.headings), + source_page=chunk.source_page, + token_count=chunk.token_count, + created_at=str(chunk.created_at), + updated_at=str(chunk.updated_at), + ) + + +def _raise_for(error: ChunkServiceError) -> None: + raise HTTPException(status_code=error.http_status, detail=str(error)) + + +# --------------------------------------------------------------------------- +# CRUD +# --------------------------------------------------------------------------- + + +@router.get("/{doc_id}/chunks", response_model=list[DocChunkResponse]) +async def list_chunks(doc_id: str, service: ServiceDep) -> list[DocChunkResponse]: + """List the canonical chunkset for a document, ordered by sequence.""" + try: + chunks = await service.list_chunks(doc_id) + except ChunkServiceError as e: + _raise_for(e) + return [_to_response(c) for c in chunks] + + +@router.post("/{doc_id}/chunks", response_model=DocChunkResponse, status_code=201) +async def add_chunk(doc_id: str, body: AddChunkRequest, service: ServiceDep) -> DocChunkResponse: + """Insert a new chunk (optionally after `afterId`).""" + if not body.text: + raise HTTPException(status_code=400, detail="text is required") + try: + chunk = await service.add_chunk(doc_id, text=body.text, after_id=body.after_id) + except ChunkServiceError as e: + _raise_for(e) + return _to_response(chunk) + + +@router.patch("/{doc_id}/chunks/{chunk_id}", response_model=DocChunkResponse) +async def update_chunk( + doc_id: str, + chunk_id: str, + body: UpdateChunkRequest, + service: ServiceDep, +) -> DocChunkResponse: + """Update a chunk's text or title (mapped to first heading).""" + if body.text is None and body.title is None: + raise HTTPException(status_code=400, detail="text or title is required") + headings: list[str] | None = None + if body.title is not None: + headings = [body.title] if body.title else [] + try: + chunk = await service.update_chunk(doc_id, chunk_id, text=body.text, headings=headings) + except ChunkServiceError as e: + _raise_for(e) + return _to_response(chunk) + + +@router.delete("/{doc_id}/chunks/{chunk_id}", status_code=204, response_model=None) +async def delete_chunk(doc_id: str, chunk_id: str, service: ServiceDep) -> None: + """Soft-delete a chunk.""" + try: + await service.delete_chunk(doc_id, chunk_id) + except ChunkServiceError as e: + _raise_for(e) + + +@router.post("/{doc_id}/chunks/{chunk_id}/split", response_model=list[DocChunkResponse]) +async def split_chunk( + doc_id: str, + chunk_id: str, + body: SplitChunkRequest, + service: ServiceDep, +) -> list[DocChunkResponse]: + """Split a chunk in two at `cursorOffset` characters.""" + try: + chunks = await service.split_chunk(doc_id, chunk_id, body.cursor_offset) + except ChunkServiceError as e: + _raise_for(e) + return [_to_response(c) for c in chunks] + + +@router.post("/{doc_id}/chunks/merge", response_model=DocChunkResponse) +async def merge_chunks( + doc_id: str, body: MergeChunksRequest, service: ServiceDep +) -> DocChunkResponse: + """Merge contiguous chunks into one. Order in `ids` is irrelevant — the + service sorts by sequence.""" + try: + merged = await service.merge_chunks(doc_id, body.ids) + except ChunkServiceError as e: + _raise_for(e) + return _to_response(merged) + + +# --------------------------------------------------------------------------- +# Rechunk + tree + diff + push +# --------------------------------------------------------------------------- + + +@router.post("/{doc_id}/rechunk", response_model=list[DocChunkResponse]) +async def rechunk_document( + doc_id: str, + body: DocRechunkRequest | None, + service: ServiceDep, +) -> list[DocChunkResponse]: + """Re-run the chunker on the latest analysis JSON; replace canonical.""" + options = body.chunking_options.model_dump() if body and body.chunking_options else None + try: + chunks = await service.rechunk_document(doc_id, options) + except ChunkServiceError as e: + _raise_for(e) + return [_to_response(c) for c in chunks] + + +@router.get("/{doc_id}/tree", response_model=list[DocTreeNodeResponse]) +async def get_tree(doc_id: str, service: ServiceDep) -> list[DocTreeNodeResponse]: + """Outline of the document built from the latest completed analysis. + + Returns `[]` if no analysis is available yet. + """ + try: + tree = await service.get_tree(doc_id) + except ChunkServiceError as e: + _raise_for(e) + return [DocTreeNodeResponse(**node) for node in tree] + + +@router.get("/{doc_id}/diff", response_model=list[ChunkDiffResponse]) +async def diff_against_store( + doc_id: str, + service: ServiceDep, + store: str = Query(..., min_length=1), +) -> list[ChunkDiffResponse]: + """Diff the canonical chunkset against the last push to `store`.""" + try: + diffs = await service.diff_against_store(doc_id, store) + except ChunkServiceError as e: + _raise_for(e) + return [ChunkDiffResponse(**d) for d in diffs] + + +@router.post("/{doc_id}/chunks/push", response_model=PushChunksResponse) +async def push_chunks( + doc_id: str, + body: PushChunksRequest, + service: ServiceDep, +) -> PushChunksResponse: + """Push the canonical chunkset to a store and record a `ChunkPush`.""" + try: + result = await service.push_to_store(doc_id, body.store) + except ChunkServiceError as e: + _raise_for(e) + return PushChunksResponse( + job_id=result["jobId"], + summary=result["summary"], + ) diff --git a/document-parser/api/schemas.py b/document-parser/api/schemas.py index ad034d9..7d3d329 100644 --- a/document-parser/api/schemas.py +++ b/document-parser/api/schemas.py @@ -293,3 +293,79 @@ class StoreDocEntryResponse(_CamelModel): state: str chunk_count: int pushed_at: str | None = None + + +# --------------------------------------------------------------------------- +# Doc-centric chunks (#256) — canonical chunkset, distinct from analysis chunks +# --------------------------------------------------------------------------- + + +class DocChunkResponse(_CamelModel): + """Canonical doc chunk — wire shape consumed by `features/chunks` on the front.""" + + id: str + doc_id: str + sequence: int + text: str + headings: list[str] = [] + source_page: int | None = None + token_count: int | None = None + created_at: str | datetime + updated_at: str | datetime + + +class AddChunkRequest(_CamelModel): + text: str + after_id: str | None = None + + +class UpdateChunkRequest(_CamelModel): + """Either or both fields. Empty body is a 400 — handled in the router.""" + + text: str | None = None + title: str | None = None # surfaced as first heading; future: dedicated field + + +class SplitChunkRequest(_CamelModel): + cursor_offset: int + + +class MergeChunksRequest(_CamelModel): + ids: list[str] + + +class DocRechunkRequest(_CamelModel): + """Optional chunking options. Empty body uses defaults.""" + + chunking_options: ChunkingOptionsRequest | None = None + + +class DocTreeNodeResponse(_CamelModel): + ref: str + type: str + label: str + children: list[DocTreeNodeResponse] = [] + + +# Forward-ref resolution (children references the same class). +DocTreeNodeResponse.model_rebuild() + + +class ChunkDiffResponse(_CamelModel): + chunk_id: str + status: str # 'added' | 'modified' | 'removed' | 'unchanged' + text_diff: str | None = None + + +class PushSummaryResponse(_CamelModel): + embeds: int + tokens: int + + +class PushChunksResponse(_CamelModel): + job_id: str + summary: PushSummaryResponse + + +class PushChunksRequest(_CamelModel): + store: str diff --git a/document-parser/domain/ports.py b/document-parser/domain/ports.py index f46c8b7..51ff6e9 100644 --- a/document-parser/domain/ports.py +++ b/document-parser/domain/ports.py @@ -156,6 +156,8 @@ class ChunkRepository(Protocol): async def find_by_id(self, chunk_id: str) -> Chunk | None: ... + async def count_for_document(self, document_id: str) -> int: ... + class ChunkEditRepository(Protocol): """Port for the immutable chunk_edits audit log (introduced by #205).""" @@ -192,6 +194,8 @@ class AnalysisRepository(Protocol): async def find_by_id(self, job_id: str) -> AnalysisJob | None: ... + async def find_latest_completed_by_document(self, document_id: str) -> AnalysisJob | None: ... + async def update_status(self, job: AnalysisJob) -> None: ... async def update_progress(self, job_id: str, current: int, total: int) -> None: ... diff --git a/document-parser/main.py b/document-parser/main.py index 58d2778..ce00bdd 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -19,6 +19,7 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from api.analyses import router as analyses_router +from api.document_chunks import router as document_chunks_router from api.documents import router as documents_router from api.ingestion import router as ingestion_router from api.schemas import HealthResponse @@ -26,11 +27,14 @@ from api.stores import router as stores_router from infra.rate_limiter import RateLimiterMiddleware from infra.settings import settings from persistence.analysis_repo import SqliteAnalysisRepository +from persistence.chunk_edit_repo import SqliteChunkEditRepository, SqliteChunkPushRepository +from persistence.chunk_repo import SqliteChunkRepository from persistence.database import get_connection, init_db from persistence.document_repo import SqliteDocumentRepository from persistence.document_store_link_repo import SqliteDocumentStoreLinkRepository from persistence.store_repo import SqliteStoreRepository from services.analysis_service import AnalysisConfig, AnalysisService +from services.chunk_service import ChunkService from services.document_service import DocumentConfig, DocumentService from services.ingestion_service import IngestionConfig, IngestionService from services.store_service import StoreService @@ -205,6 +209,26 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: if ingestion_service is not None: app.include_router(ingestion_router) logger.info("Ingestion router mounted") + + # Doc-centric chunks (#256). Wires the canonical chunkset CRUD on top + # of the chunk / chunk_edit / chunk_push repos introduced by #205. + chunk_repo = SqliteChunkRepository() + chunk_edit_repo = SqliteChunkEditRepository() + chunk_push_repo = SqliteChunkPushRepository() + app.state.chunk_repo = chunk_repo + app.state.chunk_service = ChunkService( + chunk_repo=chunk_repo, + chunk_edit_repo=chunk_edit_repo, + chunk_push_repo=chunk_push_repo, + document_repo=document_repo, + analysis_repo=analysis_repo, + chunker=_build_chunker(), + ingestion_service=ingestion_service, + ) + # The analysis service promotes the first analysis's chunks into the + # canonical chunkset (idempotent), so the doc workspace lights up the + # moment a doc is parsed for the first time. + app.state.analysis_service.set_chunk_promoter(app.state.chunk_service) logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine) try: yield @@ -236,6 +260,7 @@ if settings.rate_limit_rpm > 0: ) app.include_router(documents_router) +app.include_router(document_chunks_router) app.include_router(analyses_router) app.include_router(stores_router) diff --git a/document-parser/persistence/chunk_repo.py b/document-parser/persistence/chunk_repo.py index c0aa24c..104d928 100644 --- a/document-parser/persistence/chunk_repo.py +++ b/document-parser/persistence/chunk_repo.py @@ -132,3 +132,12 @@ class SqliteChunkRepository: cursor = await db.execute("SELECT * FROM chunks WHERE id = ?", (chunk_id,)) row = await cursor.fetchone() return _row_to_chunk(row) if row else None + + async def count_for_document(self, document_id: str) -> int: + async with get_connection() as db: + cursor = await db.execute( + "SELECT COUNT(*) AS n FROM chunks WHERE document_id = ? AND deleted_at IS NULL", + (document_id,), + ) + row = await cursor.fetchone() + return int(row["n"]) if row else 0 diff --git a/document-parser/services/analysis_service.py b/document-parser/services/analysis_service.py index 04ad7dd..ac90ed7 100644 --- a/document-parser/services/analysis_service.py +++ b/document-parser/services/analysis_service.py @@ -99,6 +99,21 @@ class AnalysisService: self._background_tasks: set[asyncio.Task] = set() self._config = config or AnalysisConfig() self._neo4j = neo4j_driver + # Duck-typed callable injected at startup. Wired in main.py to + # `ChunkService.promote_from_analysis_if_empty` so the canonical + # chunkset (#205) is populated on the first successful analysis, + # making the Doc workspace tab functional immediately (#256). + # Optional: when None, analysis behaviour is unchanged. + self._chunk_promoter = None + + def set_chunk_promoter(self, chunk_service) -> None: + """Inject the canonical-chunk promoter (post-construction wiring). + + Kept loosely typed to avoid an import cycle between + `analysis_service` and `chunk_service`. The contract is a + coroutine `(document_id: str, chunks_json: str) -> int`. + """ + self._chunk_promoter = chunk_service async def create( self, @@ -421,6 +436,20 @@ class AnalysisService: if result.page_count: await self._document_repo.update_page_count(job.document_id, result.page_count) + # Promote chunks into the canonical doc-centric chunkset on the + # first analysis (#256). Idempotent: a doc that already has + # canonical chunks is left untouched, so subsequent ad-hoc + # analyses (Studio / OCR Debug) never overwrite hand-edited state. + if chunks_json is not None and self._chunk_promoter is not None: + try: + await self._chunk_promoter.promote_from_analysis_if_empty( + job.document_id, chunks_json + ) + except Exception: + # Promotion is a best-effort hook — never fail an analysis + # because the canonical promotion hit a snag. + logger.exception("Canonical chunk promotion failed for doc %s", job.document_id) + # Drive the document lifecycle (#202): chunks present → Chunked, # otherwise → Parsed. target_state = ( diff --git a/document-parser/services/chunk_service.py b/document-parser/services/chunk_service.py new file mode 100644 index 0000000..4481ee3 --- /dev/null +++ b/document-parser/services/chunk_service.py @@ -0,0 +1,700 @@ +"""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 + ] diff --git a/document-parser/tests/test_chunk_service.py b/document-parser/tests/test_chunk_service.py new file mode 100644 index 0000000..5f1f104 --- /dev/null +++ b/document-parser/tests/test_chunk_service.py @@ -0,0 +1,378 @@ +"""Tests for ChunkService — canonical chunk lifecycle (#256).""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from domain.models import AnalysisJob, AnalysisStatus, Chunk, Document +from domain.value_objects import ChunkEditAction, ChunkResult +from persistence.analysis_repo import SqliteAnalysisRepository +from persistence.chunk_edit_repo import ( + SqliteChunkEditRepository, + SqliteChunkPushRepository, +) +from persistence.chunk_repo import SqliteChunkRepository +from persistence.database import init_db +from persistence.document_repo import SqliteDocumentRepository +from services.chunk_service import ( + ChunkConflictError, + ChunkNotFoundError, + ChunkService, + ChunkServiceError, + ChunkValidationError, + DocumentNotFoundError, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +async def setup_db(monkeypatch, tmp_path): + db_path = str(tmp_path / "test.db") + monkeypatch.setattr("persistence.database.DB_PATH", db_path) + await init_db() + yield + + +@pytest.fixture +async def doc(): + repo = SqliteDocumentRepository() + document = Document(id="doc-1", filename="t.pdf", storage_path="/tmp/t.pdf") + await repo.insert(document) + return document + + +@pytest.fixture +def repos(): + return { + "chunks": SqliteChunkRepository(), + "edits": SqliteChunkEditRepository(), + "pushes": SqliteChunkPushRepository(), + "documents": SqliteDocumentRepository(), + "analyses": SqliteAnalysisRepository(), + } + + +@pytest.fixture +def service(repos): + return ChunkService( + chunk_repo=repos["chunks"], + chunk_edit_repo=repos["edits"], + chunk_push_repo=repos["pushes"], + document_repo=repos["documents"], + analysis_repo=repos["analyses"], + chunker=None, + ingestion_service=None, + ) + + +# --------------------------------------------------------------------------- +# list_chunks +# --------------------------------------------------------------------------- + + +class TestListChunks: + async def test_list_empty(self, service, doc): + assert await service.list_chunks(doc.id) == [] + + async def test_list_filters_deleted(self, service, repos, doc): + a = Chunk(document_id=doc.id, sequence=0, text="alpha") + b = Chunk(document_id=doc.id, sequence=1, text="beta", deleted_at=datetime.now(UTC)) + await repos["chunks"].insert_many([a, b]) + chunks = await service.list_chunks(doc.id) + assert [c.id for c in chunks] == [a.id] + + async def test_404_on_missing_doc(self, service): + with pytest.raises(DocumentNotFoundError): + await service.list_chunks("no-such") + + +# --------------------------------------------------------------------------- +# add_chunk +# --------------------------------------------------------------------------- + + +class TestAddChunk: + async def test_append_to_empty(self, service, repos, doc): + new = await service.add_chunk(doc.id, text="hello") + assert new.sequence == 0 + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["hello"] + edits = await repos["edits"].find_for_document(doc.id) + assert len(edits) == 1 + assert edits[0].action == ChunkEditAction.INSERT + assert edits[0].after is not None + + async def test_append_after_anchor_shifts_sequences(self, service, repos, doc): + first = await service.add_chunk(doc.id, text="a") + last = await service.add_chunk(doc.id, text="c") + middle = await service.add_chunk(doc.id, text="b", after_id=first.id) + + chunks = await service.list_chunks(doc.id) + # Order should now be a, b, c + assert [c.text for c in chunks] == ["a", "b", "c"] + # Sequences must be dense ascending + assert [c.sequence for c in chunks] == [0, 1, 2] + assert chunks[1].id == middle.id + # `last` was shifted from sequence=1 to sequence=2 + refreshed_last = next(c for c in chunks if c.id == last.id) + assert refreshed_last.sequence == 2 + + async def test_anchor_not_found(self, service, doc): + with pytest.raises(ChunkNotFoundError): + await service.add_chunk(doc.id, text="x", after_id="no-such") + + +# --------------------------------------------------------------------------- +# update_chunk +# --------------------------------------------------------------------------- + + +class TestUpdateChunk: + async def test_update_text_records_audit(self, service, repos, doc): + new = await service.add_chunk(doc.id, text="hi") + updated = await service.update_chunk(doc.id, new.id, text="hi there") + assert updated.text == "hi there" + edits = await repos["edits"].find_for_document(doc.id) + assert {e.action for e in edits} == {ChunkEditAction.INSERT, ChunkEditAction.UPDATE} + update_edit = next(e for e in edits if e.action == ChunkEditAction.UPDATE) + assert update_edit.before["text"] == "hi" + assert update_edit.after["text"] == "hi there" + + async def test_404_on_missing_chunk(self, service, doc): + with pytest.raises(ChunkNotFoundError): + await service.update_chunk(doc.id, "no-such", text="x") + + async def test_404_on_chunk_from_other_doc(self, service, repos, doc): + other = Document(id="doc-2", filename="o.pdf", storage_path="/tmp/o.pdf") + await repos["documents"].insert(other) + new = await service.add_chunk(other.id, text="x") + with pytest.raises(ChunkNotFoundError): + await service.update_chunk(doc.id, new.id, text="y") + + +# --------------------------------------------------------------------------- +# delete_chunk +# --------------------------------------------------------------------------- + + +class TestDeleteChunk: + async def test_soft_deletes_and_records_audit(self, service, repos, doc): + new = await service.add_chunk(doc.id, text="x") + await service.delete_chunk(doc.id, new.id) + # Soft-deleted: still in repo with deleted_at set, list filters it. + assert await service.list_chunks(doc.id) == [] + edits = await repos["edits"].find_for_document(doc.id) + assert any(e.action == ChunkEditAction.DELETE for e in edits) + + +# --------------------------------------------------------------------------- +# split_chunk +# --------------------------------------------------------------------------- + + +class TestSplitChunk: + async def test_split_produces_two_chunks(self, service, repos, doc): + src = await service.add_chunk(doc.id, text="abcdef") + head, tail = await service.split_chunk(doc.id, src.id, cursor_offset=3) + assert head.text == "abc" + assert tail.text == "def" + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["abc", "def"] + edits = await repos["edits"].find_for_document(doc.id) + split_edit = next(e for e in edits if e.action == ChunkEditAction.SPLIT) + assert split_edit.children == [head.id, tail.id] + assert split_edit.before is not None + + async def test_split_400_on_offset_out_of_range(self, service, doc): + src = await service.add_chunk(doc.id, text="abc") + with pytest.raises(ChunkValidationError): + await service.split_chunk(doc.id, src.id, cursor_offset=0) + with pytest.raises(ChunkValidationError): + await service.split_chunk(doc.id, src.id, cursor_offset=99) + + async def test_split_shifts_subsequent_chunks(self, service, doc): + a = await service.add_chunk(doc.id, text="abcdef") + await service.add_chunk(doc.id, text="next") + await service.split_chunk(doc.id, a.id, cursor_offset=3) + chunks = await service.list_chunks(doc.id) + # New head, new tail, then `next` at sequence 2 + assert [c.text for c in chunks] == ["abc", "def", "next"] + assert [c.sequence for c in chunks] == [0, 1, 2] + + +# --------------------------------------------------------------------------- +# merge_chunks +# --------------------------------------------------------------------------- + + +class TestMergeChunks: + async def test_merge_contiguous(self, service, repos, doc): + a = await service.add_chunk(doc.id, text="a") + b = await service.add_chunk(doc.id, text="b") + await service.add_chunk(doc.id, text="c") + merged = await service.merge_chunks(doc.id, [b.id, a.id]) # order irrelevant + # Merged contains a,b + assert merged.text == "a\nb" + chunks = await service.list_chunks(doc.id) + # New chunk at sequence 0, then `c` still at 2 (gap is allowed) + assert {c2.text for c2 in chunks} == {"a\nb", "c"} + edit = next( + e + for e in await repos["edits"].find_for_document(doc.id) + if e.action == ChunkEditAction.MERGE + ) + assert set(edit.parents) == {a.id, b.id} + + async def test_merge_409_on_non_contiguous(self, service, doc): + a = await service.add_chunk(doc.id, text="a") + await service.add_chunk(doc.id, text="b") + c = await service.add_chunk(doc.id, text="c") + with pytest.raises(ChunkConflictError): + await service.merge_chunks(doc.id, [a.id, c.id]) + + async def test_merge_validates_min_two(self, service, doc): + a = await service.add_chunk(doc.id, text="a") + with pytest.raises(ChunkValidationError): + await service.merge_chunks(doc.id, [a.id]) + + +# --------------------------------------------------------------------------- +# rechunk_document +# --------------------------------------------------------------------------- + + +class TestRechunkDocument: + async def test_rechunk_replaces_canonical(self, service, repos, doc): + # Seed an existing canonical chunk + await service.add_chunk(doc.id, text="old") + # Seed a completed analysis with document_json + job = AnalysisJob(document_id=doc.id, status=AnalysisStatus.COMPLETED) + await repos["analyses"].insert(job) + job.document_json = json.dumps({"texts": []}) + job.completed_at = datetime.now(UTC) + await repos["analyses"].update_status(job) + + chunker = MagicMock() + chunker.chunk = AsyncMock( + return_value=[ + ChunkResult(text="new1", source_page=1, token_count=4), + ChunkResult(text="new2", source_page=2, token_count=4), + ] + ) + service._chunker = chunker + + result = await service.rechunk_document(doc.id) + assert [c.text for c in result] == ["new1", "new2"] + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["new1", "new2"] + + async def test_rechunk_409_when_no_completed_analysis(self, service, doc): + service._chunker = MagicMock() + with pytest.raises(ChunkServiceError): + await service.rechunk_document(doc.id) + + async def test_rechunk_503_when_no_chunker(self, service, doc): + with pytest.raises(ChunkServiceError) as exc: + await service.rechunk_document(doc.id) + assert exc.value.http_status == 503 + + +# --------------------------------------------------------------------------- +# promote_from_analysis_if_empty +# --------------------------------------------------------------------------- + + +class TestPromote: + async def test_promotes_when_canonical_empty(self, service, repos, doc): + chunks_json = json.dumps( + [ + {"text": "first", "headings": ["H"], "sourcePage": 1, "tokenCount": 2}, + {"text": "second", "sourcePage": 2, "tokenCount": 3}, + ] + ) + promoted = await service.promote_from_analysis_if_empty(doc.id, chunks_json) + assert promoted == 2 + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["first", "second"] + # Audit should record both INSERTs + edits = await repos["edits"].find_for_document(doc.id) + assert sum(1 for e in edits if e.action == ChunkEditAction.INSERT) == 2 + + async def test_idempotent_when_canonical_not_empty(self, service, doc): + await service.add_chunk(doc.id, text="manual") + promoted = await service.promote_from_analysis_if_empty( + doc.id, json.dumps([{"text": "auto"}]) + ) + assert promoted == 0 + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["manual"] + + async def test_skips_invalid_json(self, service, doc): + promoted = await service.promote_from_analysis_if_empty(doc.id, "not-json") + assert promoted == 0 + + async def test_skips_deleted_chunks_from_analysis(self, service, doc): + chunks_json = json.dumps( + [ + {"text": "keep"}, + {"text": "drop", "deleted": True}, + ] + ) + promoted = await service.promote_from_analysis_if_empty(doc.id, chunks_json) + assert promoted == 1 + chunks = await service.list_chunks(doc.id) + assert [c.text for c in chunks] == ["keep"] + + +# --------------------------------------------------------------------------- +# diff_against_store +# --------------------------------------------------------------------------- + + +class TestDiff: + async def test_no_push_history_returns_all_added(self, service, doc): + c1 = await service.add_chunk(doc.id, text="a") + c2 = await service.add_chunk(doc.id, text="b") + diffs = await service.diff_against_store(doc.id, "store-1") + statuses = {d["chunkId"]: d["status"] for d in diffs} + assert statuses == {c1.id: "added", c2.id: "added"} + + +# --------------------------------------------------------------------------- +# get_tree +# --------------------------------------------------------------------------- + + +class TestGetTree: + async def test_tree_empty_when_no_analysis(self, service, doc): + assert await service.get_tree(doc.id) == [] + + async def test_tree_groups_by_label(self, service, repos, doc): + job = AnalysisJob(document_id=doc.id, status=AnalysisStatus.COMPLETED) + await repos["analyses"].insert(job) + job.document_json = json.dumps( + { + "texts": [ + {"self_ref": "#/texts/0", "label": "title", "text": "Hi"}, + {"self_ref": "#/texts/1", "label": "text", "text": "Body"}, + ], + "tables": [], + "pictures": [], + } + ) + job.completed_at = datetime.now(UTC) + await repos["analyses"].update_status(job) + tree = await service.get_tree(doc.id) + labels = {n["type"] for n in tree} + assert labels == {"group"} + # Both Titles and Paragraphs groups should be present + group_titles = {n["label"] for n in tree} + assert "Titles" in group_titles + assert "Paragraphs" in group_titles diff --git a/document-parser/tests/test_document_chunks_api.py b/document-parser/tests/test_document_chunks_api.py new file mode 100644 index 0000000..ba0ff21 --- /dev/null +++ b/document-parser/tests/test_document_chunks_api.py @@ -0,0 +1,225 @@ +"""Integration tests for the /api/documents/{id}/chunks router (#256).""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +from domain.models import Chunk +from main import app +from services.chunk_service import ( + ChunkConflictError, + ChunkNotFoundError, + ChunkValidationError, + DocumentNotFoundError, +) + + +@pytest.fixture +def client(): + return TestClient(app, raise_server_exceptions=False) + + +@pytest.fixture +def mock_service(client): + svc = MagicMock() + original = getattr(app.state, "chunk_service", None) + app.state.chunk_service = svc + yield svc + app.state.chunk_service = original + + +def _chunk(*, id="c-1", doc_id="d-1", sequence=0, text="hi") -> Chunk: + return Chunk(id=id, document_id=doc_id, sequence=sequence, text=text) + + +class TestListChunks: + def test_200(self, client, mock_service): + mock_service.list_chunks = AsyncMock(return_value=[_chunk(id="c1"), _chunk(id="c2")]) + resp = client.get("/api/documents/d-1/chunks") + assert resp.status_code == 200 + data = resp.json() + assert [c["id"] for c in data] == ["c1", "c2"] + # camelCase serialization + assert "docId" in data[0] + assert "sourcePage" in data[0] + + def test_404_when_doc_missing(self, client, mock_service): + mock_service.list_chunks = AsyncMock(side_effect=DocumentNotFoundError("nope")) + resp = client.get("/api/documents/no-such/chunks") + assert resp.status_code == 404 + + +class TestAddChunk: + def test_201(self, client, mock_service): + mock_service.add_chunk = AsyncMock(return_value=_chunk(id="new")) + resp = client.post("/api/documents/d-1/chunks", json={"text": "hello"}) + assert resp.status_code == 201 + assert resp.json()["id"] == "new" + mock_service.add_chunk.assert_awaited_once_with("d-1", text="hello", after_id=None) + + def test_400_on_empty_body(self, client, mock_service): + resp = client.post("/api/documents/d-1/chunks", json={"text": ""}) + assert resp.status_code == 400 + + def test_passes_after_id(self, client, mock_service): + mock_service.add_chunk = AsyncMock(return_value=_chunk(id="new")) + client.post( + "/api/documents/d-1/chunks", + json={"text": "x", "afterId": "c-prev"}, + ) + mock_service.add_chunk.assert_awaited_once_with("d-1", text="x", after_id="c-prev") + + +class TestUpdateChunk: + def test_200(self, client, mock_service): + mock_service.update_chunk = AsyncMock(return_value=_chunk(text="new")) + resp = client.patch( + "/api/documents/d-1/chunks/c-1", + json={"text": "new"}, + ) + assert resp.status_code == 200 + assert resp.json()["text"] == "new" + + def test_400_when_no_field_set(self, client, mock_service): + resp = client.patch("/api/documents/d-1/chunks/c-1", json={}) + assert resp.status_code == 400 + + def test_404_when_chunk_missing(self, client, mock_service): + mock_service.update_chunk = AsyncMock(side_effect=ChunkNotFoundError("nope")) + resp = client.patch("/api/documents/d-1/chunks/c-1", json={"text": "x"}) + assert resp.status_code == 404 + + def test_title_maps_to_first_heading(self, client, mock_service): + mock_service.update_chunk = AsyncMock(return_value=_chunk()) + client.patch("/api/documents/d-1/chunks/c-1", json={"title": "Section A"}) + mock_service.update_chunk.assert_awaited_once_with( + "d-1", "c-1", text=None, headings=["Section A"] + ) + + +class TestDeleteChunk: + def test_204(self, client, mock_service): + mock_service.delete_chunk = AsyncMock(return_value=None) + resp = client.delete("/api/documents/d-1/chunks/c-1") + assert resp.status_code == 204 + + def test_404(self, client, mock_service): + mock_service.delete_chunk = AsyncMock(side_effect=ChunkNotFoundError("nope")) + resp = client.delete("/api/documents/d-1/chunks/c-1") + assert resp.status_code == 404 + + +class TestSplit: + def test_200(self, client, mock_service): + mock_service.split_chunk = AsyncMock( + return_value=[_chunk(id="head"), _chunk(id="tail", sequence=1)] + ) + resp = client.post( + "/api/documents/d-1/chunks/c-1/split", + json={"cursorOffset": 3}, + ) + assert resp.status_code == 200 + assert [c["id"] for c in resp.json()] == ["head", "tail"] + + def test_400_on_invalid_offset(self, client, mock_service): + mock_service.split_chunk = AsyncMock(side_effect=ChunkValidationError("oor")) + resp = client.post( + "/api/documents/d-1/chunks/c-1/split", + json={"cursorOffset": 999}, + ) + assert resp.status_code == 400 + + +class TestMerge: + def test_200(self, client, mock_service): + mock_service.merge_chunks = AsyncMock(return_value=_chunk(id="merged")) + resp = client.post( + "/api/documents/d-1/chunks/merge", + json={"ids": ["a", "b"]}, + ) + assert resp.status_code == 200 + assert resp.json()["id"] == "merged" + + def test_409_on_non_contiguous(self, client, mock_service): + mock_service.merge_chunks = AsyncMock(side_effect=ChunkConflictError("nope")) + resp = client.post( + "/api/documents/d-1/chunks/merge", + json={"ids": ["a", "z"]}, + ) + assert resp.status_code == 409 + + +class TestRechunk: + def test_200_no_options(self, client, mock_service): + mock_service.rechunk_document = AsyncMock(return_value=[_chunk(id="r1")]) + resp = client.post("/api/documents/d-1/rechunk", json={}) + assert resp.status_code == 200 + mock_service.rechunk_document.assert_awaited_once_with("d-1", None) + + def test_200_with_options(self, client, mock_service): + mock_service.rechunk_document = AsyncMock(return_value=[]) + resp = client.post( + "/api/documents/d-1/rechunk", + json={"chunkingOptions": {"chunkerType": "hybrid", "maxTokens": 256}}, + ) + assert resp.status_code == 200 + call_args = mock_service.rechunk_document.await_args + assert call_args.args[0] == "d-1" + passed_opts = call_args.args[1] + # snake_case dump from Pydantic ChunkingOptionsRequest + assert passed_opts["chunker_type"] == "hybrid" + assert passed_opts["max_tokens"] == 256 + + +class TestTree: + def test_200_empty(self, client, mock_service): + mock_service.get_tree = AsyncMock(return_value=[]) + resp = client.get("/api/documents/d-1/tree") + assert resp.status_code == 200 + assert resp.json() == [] + + def test_200_groups(self, client, mock_service): + mock_service.get_tree = AsyncMock( + return_value=[ + {"ref": "#group/title", "type": "group", "label": "Titles", "children": []} + ] + ) + resp = client.get("/api/documents/d-1/tree") + assert resp.status_code == 200 + data = resp.json() + assert data[0]["label"] == "Titles" + + +class TestDiff: + def test_200(self, client, mock_service): + mock_service.diff_against_store = AsyncMock( + return_value=[ + {"chunkId": "c1", "status": "added", "textDiff": None}, + ] + ) + resp = client.get("/api/documents/d-1/diff?store=mystore") + assert resp.status_code == 200 + assert resp.json()[0]["chunkId"] == "c1" + + def test_400_when_store_missing(self, client, mock_service): + resp = client.get("/api/documents/d-1/diff") + # FastAPI auto-422 on missing required query + assert resp.status_code in (400, 422) + + +class TestPush: + def test_200(self, client, mock_service): + mock_service.push_to_store = AsyncMock( + return_value={"jobId": "p1", "summary": {"embeds": 3, "tokens": 30}} + ) + resp = client.post( + "/api/documents/d-1/chunks/push", + json={"store": "mystore"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["jobId"] == "p1" + assert body["summary"]["embeds"] == 3