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
This commit is contained in:
Pier-Jean Malandrino
parent
75a2db087b
commit
ef80e22342
9 changed files with 1658 additions and 0 deletions
212
document-parser/api/document_chunks.py
Normal file
212
document-parser/api/document_chunks.py
Normal file
|
|
@ -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"],
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
700
document-parser/services/chunk_service.py
Normal file
700
document-parser/services/chunk_service.py
Normal file
|
|
@ -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
|
||||
]
|
||||
378
document-parser/tests/test_chunk_service.py
Normal file
378
document-parser/tests/test_chunk_service.py
Normal file
|
|
@ -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
|
||||
225
document-parser/tests/test_document_chunks_api.py
Normal file
225
document-parser/tests/test_document_chunks_api.py
Normal file
|
|
@ -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
|
||||
Loading…
Reference in a new issue