diff --git a/document-parser/api/analyses.py b/document-parser/api/analyses.py index 4fbf988..123f12f 100644 --- a/document-parser/api/analyses.py +++ b/document-parser/api/analyses.py @@ -7,7 +7,7 @@ from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, Request -from api.schemas import AnalysisResponse, CreateAnalysisRequest +from api.schemas import AnalysisResponse, ChunkResponse, CreateAnalysisRequest, RechunkRequest from services.analysis_service import AnalysisService logger = logging.getLogger(__name__) @@ -30,6 +30,8 @@ def _to_response(job) -> AnalysisResponse: content_markdown=job.content_markdown, content_html=job.content_html, pages_json=job.pages_json, + chunks_json=job.chunks_json, + has_document_json=job.document_json is not None, error_message=job.error_message, started_at=str(job.started_at) if job.started_at else None, completed_at=str(job.completed_at) if job.completed_at else None, @@ -47,8 +49,16 @@ async def create_analysis(body: CreateAnalysisRequest, service: ServiceDep): if body.pipelineOptions: pipeline_opts = body.pipelineOptions.model_dump() + chunking_opts = None + if body.chunkingOptions: + chunking_opts = body.chunkingOptions.model_dump() + try: - job = await service.create(body.documentId, pipeline_options=pipeline_opts) + job = await service.create( + body.documentId, + pipeline_options=pipeline_opts, + chunking_options=chunking_opts, + ) except ValueError as e: raise HTTPException(status_code=404, detail=str(e)) from e @@ -71,6 +81,24 @@ async def get_analysis(job_id: str, service: ServiceDep): return _to_response(job) +@router.post("/{job_id}/rechunk", response_model=list[ChunkResponse]) +async def rechunk_analysis(job_id: str, body: RechunkRequest, service: ServiceDep): + """Re-chunk a completed analysis with new chunking options.""" + try: + chunks = await service.rechunk(job_id, body.chunkingOptions.model_dump()) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) from e + return [ + ChunkResponse( + text=c.text, + headings=c.headings, + source_page=c.source_page, + token_count=c.token_count, + ) + for c in chunks + ] + + @router.delete("/{job_id}", status_code=204) async def delete_analysis(job_id: str, service: ServiceDep): """Delete an analysis job.""" diff --git a/document-parser/api/schemas.py b/document-parser/api/schemas.py index 4e30153..bf01fd5 100644 --- a/document-parser/api/schemas.py +++ b/document-parser/api/schemas.py @@ -18,6 +18,7 @@ def _to_camel(name: str) -> str: class _CamelModel(BaseModel): """Base model that serializes field names to camelCase.""" + model_config = ConfigDict( alias_generator=_to_camel, populate_by_name=True, @@ -42,6 +43,8 @@ class AnalysisResponse(_CamelModel): content_markdown: str | None = None content_html: str | None = None pages_json: str | None = None + chunks_json: str | None = None + has_document_json: bool = False error_message: str | None = None started_at: str | datetime | None = None completed_at: str | datetime | None = None @@ -50,6 +53,7 @@ class AnalysisResponse(_CamelModel): class PipelineOptionsRequest(BaseModel): """Docling pipeline configuration options.""" + do_ocr: bool = True do_table_structure: bool = True table_mode: str = "accurate" # "accurate" or "fast" @@ -76,6 +80,41 @@ class PipelineOptionsRequest(BaseModel): return v +class ChunkingOptionsRequest(BaseModel): + """Docling chunking configuration options.""" + + chunker_type: str = "hybrid" # "hybrid", "hierarchical" + max_tokens: int = 512 + merge_peers: bool = True + repeat_table_header: bool = True + + @field_validator("chunker_type") + @classmethod + def validate_chunker_type(cls, v: str) -> str: + if v not in ("hybrid", "hierarchical"): + raise ValueError('chunker_type must be "hybrid" or "hierarchical"') + return v + + @field_validator("max_tokens") + @classmethod + def validate_max_tokens(cls, v: int) -> int: + if v < 64 or v > 8192: + raise ValueError("max_tokens must be between 64 and 8192") + return v + + +class ChunkResponse(_CamelModel): + text: str + headings: list[str] = [] + source_page: int | None = None + token_count: int = 0 + + class CreateAnalysisRequest(BaseModel): documentId: str # camelCase to match existing frontend contract pipelineOptions: PipelineOptionsRequest | None = None + chunkingOptions: ChunkingOptionsRequest | None = None + + +class RechunkRequest(BaseModel): + chunkingOptions: ChunkingOptionsRequest diff --git a/document-parser/conftest.py b/document-parser/conftest.py index 60a2633..ddf789c 100644 --- a/document-parser/conftest.py +++ b/document-parser/conftest.py @@ -1,3 +1 @@ - - pytest_plugins = ["pytest_asyncio"] diff --git a/document-parser/domain/bbox.py b/document-parser/domain/bbox.py index 41530a7..4804df9 100644 --- a/document-parser/domain/bbox.py +++ b/document-parser/domain/bbox.py @@ -43,7 +43,11 @@ def to_topleft_list(bbox: BoundingBox, page_height: float) -> list[float]: if right <= left or bottom <= top: logger.debug( "Degenerate bbox skipped: [%.1f, %.1f, %.1f, %.1f] (page_height=%.1f)", - left, top, right, bottom, page_height, + left, + top, + right, + bottom, + page_height, ) return list(EMPTY_BBOX) diff --git a/document-parser/domain/models.py b/document-parser/domain/models.py index 3df9eb3..66fc3fe 100644 --- a/document-parser/domain/models.py +++ b/document-parser/domain/models.py @@ -42,6 +42,8 @@ class AnalysisJob: content_markdown: str | None = None content_html: str | None = None pages_json: str | None = None + document_json: str | None = None + chunks_json: str | None = None error_message: str | None = None started_at: datetime | None = None completed_at: datetime | None = None @@ -55,12 +57,19 @@ class AnalysisJob: self.started_at = _utcnow() def mark_completed( - self, markdown: str, html: str, pages_json: str, + self, + markdown: str, + html: str, + pages_json: str, + document_json: str | None = None, + chunks_json: str | None = None, ) -> None: self.status = AnalysisStatus.COMPLETED self.content_markdown = markdown self.content_html = html self.pages_json = pages_json + self.document_json = document_json + self.chunks_json = chunks_json self.completed_at = _utcnow() def mark_failed(self, error: str) -> None: diff --git a/document-parser/domain/ports.py b/document-parser/domain/ports.py index bc36c9b..f4ea64d 100644 --- a/document-parser/domain/ports.py +++ b/document-parser/domain/ports.py @@ -9,7 +9,12 @@ from __future__ import annotations from typing import TYPE_CHECKING, Protocol if TYPE_CHECKING: - from domain.value_objects import ConversionOptions, ConversionResult + from domain.value_objects import ( + ChunkingOptions, + ChunkResult, + ConversionOptions, + ConversionResult, + ) class DocumentConverter(Protocol): @@ -20,5 +25,20 @@ class DocumentConverter(Protocol): """ async def convert( - self, file_path: str, options: ConversionOptions, + self, + file_path: str, + options: ConversionOptions, ) -> ConversionResult: ... + + +class DocumentChunker(Protocol): + """Port for document chunking. + + Takes a serialized DoclingDocument (JSON) and returns chunks. + """ + + async def chunk( + self, + document_json: str, + options: ChunkingOptions, + ) -> list[ChunkResult]: ... diff --git a/document-parser/domain/value_objects.py b/document-parser/domain/value_objects.py index b4c5959..19af834 100644 --- a/document-parser/domain/value_objects.py +++ b/document-parser/domain/value_objects.py @@ -49,3 +49,23 @@ class ConversionResult: content_html: str pages: list[PageDetail] skipped_items: int = 0 + document_json: str | None = None + + +@dataclass +class ChunkingOptions: + chunker_type: str = "hybrid" # "hybrid", "hierarchical", "page" + max_tokens: int = 512 + merge_peers: bool = True + repeat_table_header: bool = True + + def is_default(self) -> bool: + return self == ChunkingOptions() + + +@dataclass +class ChunkResult: + text: str + headings: list[str] = field(default_factory=list) + source_page: int | None = None + token_count: int = 0 diff --git a/document-parser/infra/local_chunker.py b/document-parser/infra/local_chunker.py new file mode 100644 index 0000000..7024496 --- /dev/null +++ b/document-parser/infra/local_chunker.py @@ -0,0 +1,84 @@ +"""Local Docling chunker — runs chunking in-process using docling-core. + +This adapter implements the DocumentChunker port. It deserializes a +DoclingDocument from JSON, applies the requested chunker, and returns +domain ChunkResult objects. +""" + +from __future__ import annotations + +import asyncio +import json +import logging + +from docling_core.transforms.chunker import HierarchicalChunker +from docling_core.transforms.chunker.hybrid_chunker import HybridChunker +from docling_core.types.doc.document import DoclingDocument + +from domain.value_objects import ChunkingOptions, ChunkResult + +logger = logging.getLogger(__name__) + + +def _chunk_sync(document_json: str, options: ChunkingOptions) -> list[ChunkResult]: + if not document_json or not document_json.strip(): + raise ValueError("Empty document JSON — nothing to chunk") + + try: + doc_data = json.loads(document_json) + except json.JSONDecodeError as e: + raise ValueError(f"Malformed document JSON: {e}") from e + + doc = DoclingDocument.model_validate(doc_data) + + chunker = _build_chunker(options) + results: list[ChunkResult] = [] + + for chunk in chunker.chunk(doc): + source_page = None + token_count = 0 + + if hasattr(chunk, "meta") and chunk.meta and chunk.meta.doc_items: + for doc_item in chunk.meta.doc_items: + if hasattr(doc_item, "prov") and doc_item.prov: + source_page = doc_item.prov[0].page_no + break + + if hasattr(chunker, "tokenizer") and chunker.tokenizer: + token_count = chunker.tokenizer.count_tokens(chunk.text) + + headings = list(chunk.meta.headings) if chunk.meta and chunk.meta.headings else [] + + results.append( + ChunkResult( + text=chunk.text, + headings=headings, + source_page=source_page, + token_count=token_count, + ) + ) + + logger.info("Chunked document into %d chunks (chunker=%s)", len(results), options.chunker_type) + return results + + +def _build_chunker(options: ChunkingOptions) -> HierarchicalChunker | HybridChunker: + if options.chunker_type == "hierarchical": + return HierarchicalChunker() + + return HybridChunker( + max_tokens=options.max_tokens, + merge_peers=options.merge_peers, + repeat_table_header=options.repeat_table_header, + ) + + +class LocalChunker: + """Adapter that runs docling-core chunking locally.""" + + async def chunk( + self, + document_json: str, + options: ChunkingOptions, + ) -> list[ChunkResult]: + return await asyncio.to_thread(_chunk_sync, document_json, options) diff --git a/document-parser/infra/local_converter.py b/document-parser/infra/local_converter.py index efd4a5a..9e0b23a 100644 --- a/document-parser/infra/local_converter.py +++ b/document-parser/infra/local_converter.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio import contextlib +import json import logging import threading @@ -83,6 +84,7 @@ def _get_element_type(item: DocItem) -> str: # Pipeline factory # --------------------------------------------------------------------------- + def _build_docling_converter(options: ConversionOptions) -> DoclingConverter: table_options = TableStructureOptions( do_cell_matching=True, @@ -126,6 +128,7 @@ def _select_converter(options: ConversionOptions) -> DoclingConverter: # Page extraction # --------------------------------------------------------------------------- + def _extract_pages_detail(doc_result) -> tuple[list[PageDetail], int]: pages: dict[int, PageDetail] = {} document = doc_result.document @@ -149,7 +152,9 @@ def _extract_pages_detail(doc_result) -> tuple[list[PageDetail], int]: def _process_content_item( - item: DocItem | GroupItem, level: int, pages: dict[int, PageDetail], + item: DocItem | GroupItem, + level: int, + pages: dict[int, PageDetail], ) -> bool: if isinstance(item, GroupItem): return True @@ -163,9 +168,13 @@ def _process_content_item( if page_no not in pages: logger.warning( "Page %d not found in document metadata — using US Letter fallback (%sx%s pt)", - page_no, _DEFAULT_PAGE_WIDTH, _DEFAULT_PAGE_HEIGHT, + page_no, + _DEFAULT_PAGE_WIDTH, + _DEFAULT_PAGE_HEIGHT, + ) + pages[page_no] = PageDetail( + page_number=page_no, width=_DEFAULT_PAGE_WIDTH, height=_DEFAULT_PAGE_HEIGHT ) - pages[page_no] = PageDetail(page_number=page_no, width=_DEFAULT_PAGE_WIDTH, height=_DEFAULT_PAGE_HEIGHT) page_height = pages[page_no].height @@ -199,6 +208,7 @@ def _process_content_item( # Synchronous conversion (called via asyncio.to_thread) # --------------------------------------------------------------------------- + def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResult: with _converter_lock: conv = _select_converter(options) @@ -213,7 +223,9 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul PageDetail( page_number=i + 1, width=doc.pages[i + 1].size.width if (i + 1) in doc.pages else _DEFAULT_PAGE_WIDTH, - height=doc.pages[i + 1].size.height if (i + 1) in doc.pages else _DEFAULT_PAGE_HEIGHT, + height=doc.pages[i + 1].size.height + if (i + 1) in doc.pages + else _DEFAULT_PAGE_HEIGHT, ) for i in range(page_count) ] @@ -227,6 +239,7 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul content_html=doc.export_to_html(), pages=pages_detail, skipped_items=skipped, + document_json=json.dumps(doc.export_to_dict()), ) @@ -234,10 +247,13 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul # Public adapter class # --------------------------------------------------------------------------- + class LocalConverter: """Adapter that runs Docling locally as a Python library.""" async def convert( - self, file_path: str, options: ConversionOptions, + self, + file_path: str, + options: ConversionOptions, ) -> ConversionResult: return await asyncio.to_thread(_convert_sync, file_path, options) diff --git a/document-parser/infra/serve_converter.py b/document-parser/infra/serve_converter.py index d85e4a5..410e497 100644 --- a/document-parser/infra/serve_converter.py +++ b/document-parser/infra/serve_converter.py @@ -71,7 +71,9 @@ class ServeConverter: return headers async def convert( - self, file_path: str, options: ConversionOptions, + self, + file_path: str, + options: ConversionOptions, ) -> ConversionResult: """Convert a document by uploading it to Docling Serve.""" path = Path(file_path) @@ -202,7 +204,9 @@ def _add_element(item: dict, pages: dict[int, PageDetail]) -> None: page_no = prov.get("page_no", 1) if page_no not in pages: pages[page_no] = PageDetail( - page_number=page_no, width=612.0, height=792.0, + page_number=page_no, + width=612.0, + height=792.0, ) bbox_data = prov.get("bbox", {}) diff --git a/document-parser/infra/settings.py b/document-parser/infra/settings.py index 0b19c54..348f551 100644 --- a/document-parser/infra/settings.py +++ b/document-parser/infra/settings.py @@ -14,7 +14,9 @@ class Settings: conversion_timeout: int = 600 upload_dir: str = "./uploads" db_path: str = "./data/docling_studio.db" - cors_origins: list[str] = field(default_factory=lambda: ["http://localhost:3000", "http://localhost:5173"]) + cors_origins: list[str] = field( + default_factory=lambda: ["http://localhost:3000", "http://localhost:5173"] + ) @classmethod def from_env(cls) -> Settings: diff --git a/document-parser/main.py b/document-parser/main.py index 20a5478..f6d1058 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -40,6 +40,7 @@ def _build_converter(): """Build the converter adapter based on configuration.""" if settings.conversion_engine == "remote": from infra.serve_converter import ServeConverter + logger.info("Using remote Docling Serve at %s", settings.docling_serve_url) return ServeConverter( base_url=settings.docling_serve_url, @@ -47,14 +48,26 @@ def _build_converter(): ) else: from infra.local_converter import LocalConverter + logger.info("Using local Docling converter") return LocalConverter() +def _build_chunker(): + """Build the chunker adapter — only available in local mode.""" + if settings.conversion_engine == "local": + from infra.local_chunker import LocalChunker + + return LocalChunker() + return None + + def _build_analysis_service() -> AnalysisService: converter = _build_converter() + chunker = _build_chunker() return AnalysisService( converter=converter, + chunker=chunker, conversion_timeout=settings.conversion_timeout, ) @@ -63,6 +76,7 @@ def _build_analysis_service() -> AnalysisService: # FastAPI app # --------------------------------------------------------------------------- + @asynccontextmanager async def lifespan(app: FastAPI): await init_db() diff --git a/document-parser/persistence/analysis_repo.py b/document-parser/persistence/analysis_repo.py index f8b3df7..0652e5b 100644 --- a/document-parser/persistence/analysis_repo.py +++ b/document-parser/persistence/analysis_repo.py @@ -16,6 +16,7 @@ def _parse_dt(value: str | None) -> datetime | None: def _row_to_job(row) -> AnalysisJob: + keys = row.keys() return AnalysisJob( id=row["id"], document_id=row["document_id"], @@ -23,11 +24,13 @@ def _row_to_job(row) -> AnalysisJob: content_markdown=row["content_markdown"], content_html=row["content_html"], pages_json=row["pages_json"], + document_json=row["document_json"] if "document_json" in keys else None, + chunks_json=row["chunks_json"] if "chunks_json" in keys else None, error_message=row["error_message"], started_at=_parse_dt(row["started_at"]), completed_at=_parse_dt(row["completed_at"]), created_at=_parse_dt(row["created_at"]) or datetime.now(), - document_filename=row["filename"] if "filename" in row.keys() else None, # noqa: SIM118 + document_filename=row["filename"] if "filename" in keys else None, ) @@ -60,9 +63,7 @@ async def find_all(*, limit: int = 200, offset: int = 0) -> list[AnalysisJob]: async def find_by_id(job_id: str) -> AnalysisJob | None: async with get_connection() as db: - cursor = await db.execute( - f"{_SELECT_WITH_DOC} WHERE aj.id = ?", (job_id,) - ) + cursor = await db.execute(f"{_SELECT_WITH_DOC} WHERE aj.id = ?", (job_id,)) row = await cursor.fetchone() return _row_to_job(row) if row else None @@ -72,17 +73,36 @@ async def update_status(job: AnalysisJob) -> None: await db.execute( """UPDATE analysis_jobs SET status = ?, content_markdown = ?, content_html = ?, - pages_json = ?, error_message = ?, started_at = ?, completed_at = ? + pages_json = ?, document_json = ?, chunks_json = ?, + error_message = ?, started_at = ?, completed_at = ? WHERE id = ?""", - (job.status.value, job.content_markdown, job.content_html, - job.pages_json, job.error_message, - str(job.started_at) if job.started_at else None, - str(job.completed_at) if job.completed_at else None, - job.id), + ( + job.status.value, + job.content_markdown, + job.content_html, + job.pages_json, + job.document_json, + job.chunks_json, + job.error_message, + str(job.started_at) if job.started_at else None, + str(job.completed_at) if job.completed_at else None, + job.id, + ), ) await db.commit() +async def update_chunks(job_id: str, chunks_json: str) -> bool: + """Update only the chunks_json column for a completed analysis.""" + async with get_connection() as db: + cursor = await db.execute( + "UPDATE analysis_jobs SET chunks_json = ? WHERE id = ?", + (chunks_json, job_id), + ) + await db.commit() + return cursor.rowcount > 0 + + async def delete(job_id: str) -> bool: async with get_connection() as db: cursor = await db.execute("DELETE FROM analysis_jobs WHERE id = ?", (job_id,)) @@ -93,8 +113,6 @@ async def delete(job_id: str) -> bool: async def delete_by_document(document_id: str) -> int: """Delete all analysis jobs for a given document. Returns count deleted.""" async with get_connection() as db: - cursor = await db.execute( - "DELETE FROM analysis_jobs WHERE document_id = ?", (document_id,) - ) + cursor = await db.execute("DELETE FROM analysis_jobs WHERE document_id = ?", (document_id,)) await db.commit() return cursor.rowcount diff --git a/document-parser/persistence/database.py b/document-parser/persistence/database.py index 6aaa37b..93c7d94 100644 --- a/document-parser/persistence/database.py +++ b/document-parser/persistence/database.py @@ -30,6 +30,8 @@ CREATE TABLE IF NOT EXISTS analysis_jobs ( content_markdown TEXT, content_html TEXT, pages_json TEXT, + document_json TEXT, + chunks_json TEXT, error_message TEXT, started_at TEXT, completed_at TEXT, @@ -42,11 +44,29 @@ CREATE INDEX IF NOT EXISTS idx_documents_created_at ON documents(created_at); """ +_MIGRATIONS = [ + ("document_json", "ALTER TABLE analysis_jobs ADD COLUMN document_json TEXT"), + ("chunks_json", "ALTER TABLE analysis_jobs ADD COLUMN chunks_json TEXT"), +] + + +async def _run_migrations(db: aiosqlite.Connection) -> None: + """Add columns that may be missing in older databases.""" + cursor = await db.execute("PRAGMA table_info(analysis_jobs)") + existing = {row[1] for row in await cursor.fetchall()} + for col_name, ddl in _MIGRATIONS: + if col_name not in existing: + await db.execute(ddl) + logger.info("Migration: added column %s to analysis_jobs", col_name) + await db.commit() + + async def init_db() -> None: """Create database file and tables if they don't exist.""" os.makedirs(os.path.dirname(DB_PATH) or ".", exist_ok=True) async with aiosqlite.connect(DB_PATH) as db: await db.executescript(_SCHEMA) + await _run_migrations(db) await db.commit() logger.info("Database initialized at %s", DB_PATH) diff --git a/document-parser/persistence/document_repo.py b/document-parser/persistence/document_repo.py index 868b24d..623e36a 100644 --- a/document-parser/persistence/document_repo.py +++ b/document-parser/persistence/document_repo.py @@ -28,8 +28,15 @@ async def insert(doc: Document) -> None: await db.execute( """INSERT INTO documents (id, filename, content_type, file_size, page_count, storage_path, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)""", - (doc.id, doc.filename, doc.content_type, doc.file_size, - doc.page_count, doc.storage_path, str(doc.created_at)), + ( + doc.id, + doc.filename, + doc.content_type, + doc.file_size, + doc.page_count, + doc.storage_path, + str(doc.created_at), + ), ) await db.commit() diff --git a/document-parser/services/analysis_service.py b/document-parser/services/analysis_service.py index 2059506..9c99934 100644 --- a/document-parser/services/analysis_service.py +++ b/document-parser/services/analysis_service.py @@ -13,24 +13,46 @@ import logging from dataclasses import asdict from typing import TYPE_CHECKING -from domain.models import AnalysisJob -from domain.value_objects import ConversionOptions, ConversionResult +from domain.models import AnalysisJob, AnalysisStatus +from domain.value_objects import ChunkingOptions, ChunkResult, ConversionOptions, ConversionResult if TYPE_CHECKING: - from domain.ports import DocumentConverter + from domain.ports import DocumentChunker, DocumentConverter from persistence import analysis_repo, document_repo logger = logging.getLogger(__name__) +def _chunk_to_dict(c: ChunkResult) -> dict: + """Serialize ChunkResult to a camelCase dict matching the frontend API contract.""" + return { + "text": c.text, + "headings": c.headings, + "sourcePage": c.source_page, + "tokenCount": c.token_count, + } + + class AnalysisService: """Orchestrates document analysis using an injected converter.""" - def __init__(self, converter: DocumentConverter, conversion_timeout: int = 600): + def __init__( + self, + converter: DocumentConverter, + chunker: DocumentChunker | None = None, + conversion_timeout: int = 600, + ): self._converter = converter + self._chunker = chunker self._conversion_timeout = conversion_timeout - async def create(self, document_id: str, *, pipeline_options: dict | None = None) -> AnalysisJob: + async def create( + self, + document_id: str, + *, + pipeline_options: dict | None = None, + chunking_options: dict | None = None, + ) -> AnalysisJob: """Create a new analysis job and launch background processing.""" doc = await document_repo.find_by_id(document_id) if not doc: @@ -41,7 +63,13 @@ class AnalysisService: await analysis_repo.insert(job) task = asyncio.create_task( - self._run_analysis(job.id, doc.storage_path, doc.filename, pipeline_options) + self._run_analysis( + job.id, + doc.storage_path, + doc.filename, + pipeline_options, + chunking_options, + ) ) task.add_done_callback(functools.partial(_on_task_done, job_id=job.id)) @@ -56,10 +84,35 @@ class AnalysisService: async def delete(self, job_id: str) -> bool: return await analysis_repo.delete(job_id) + async def rechunk(self, job_id: str, chunking_options: dict) -> list[ChunkResult]: + """Re-chunk an existing completed analysis with new options.""" + job = await analysis_repo.find_by_id(job_id) + if not job: + raise ValueError(f"Analysis not found: {job_id}") + if job.status != AnalysisStatus.COMPLETED: + raise ValueError(f"Analysis is not completed: {job_id}") + if not job.document_json: + raise ValueError(f"No document data available for re-chunking: {job_id}") + if not self._chunker: + raise ValueError("Chunking is not available") + + options = ChunkingOptions(**chunking_options) + chunks = await self._chunker.chunk(job.document_json, options) + + chunks_json = json.dumps([_chunk_to_dict(c) for c in chunks]) + await analysis_repo.update_chunks(job_id, chunks_json) + + return chunks + async def _run_analysis( - self, job_id: str, file_path: str, filename: str, pipeline_options: dict | None = None, + self, + job_id: str, + file_path: str, + filename: str, + pipeline_options: dict | None = None, + chunking_options: dict | None = None, ) -> None: - """Background task: run conversion and update job status.""" + """Background task: run conversion and optionally chunk.""" try: job = await analysis_repo.find_by_id(job_id) if not job: @@ -79,10 +132,19 @@ class AnalysisService: pages_json = json.dumps([asdict(p) for p in result.pages]) + chunks_json = None + if chunking_options and self._chunker and result.document_json: + chunk_opts = ChunkingOptions(**chunking_options) + chunks = await self._chunker.chunk(result.document_json, chunk_opts) + chunks_json = json.dumps([_chunk_to_dict(c) for c in chunks]) + logger.info("Chunking produced %d chunks for job %s", len(chunks), job_id) + job.mark_completed( markdown=result.content_markdown, html=result.content_html, pages_json=pages_json, + document_json=result.document_json, + chunks_json=chunks_json, ) await analysis_repo.update_status(job) diff --git a/document-parser/tests/test_api_endpoints.py b/document-parser/tests/test_api_endpoints.py index 173dd30..2f34518 100644 --- a/document-parser/tests/test_api_endpoints.py +++ b/document-parser/tests/test_api_endpoints.py @@ -53,8 +53,12 @@ class TestDocumentEndpoints: @patch("services.document_service.find_by_id", new_callable=AsyncMock) def test_get_document(self, mock_find, client): mock_find.return_value = Document( - id="d1", filename="test.pdf", content_type="application/pdf", - file_size=2048, page_count=3, storage_path="/tmp/test", + id="d1", + filename="test.pdf", + content_type="application/pdf", + file_size=2048, + page_count=3, + storage_path="/tmp/test", ) resp = client.get("/api/documents/d1") @@ -74,8 +78,10 @@ class TestDocumentEndpoints: @patch("services.document_service.upload", new_callable=AsyncMock) def test_upload_document(self, mock_upload, client): mock_upload.return_value = Document( - id="new-1", filename="uploaded.pdf", - content_type="application/pdf", file_size=512, + id="new-1", + filename="uploaded.pdf", + content_type="application/pdf", + file_size=512, storage_path="/tmp/uploaded", ) @@ -115,9 +121,11 @@ class TestDocumentEndpoints: class TestAnalysisEndpoints: def test_list_analyses(self, client, mock_analysis_service): - mock_analysis_service.find_all = AsyncMock(return_value=[ - AnalysisJob(id="j1", document_id="d1", document_filename="test.pdf"), - ]) + mock_analysis_service.find_all = AsyncMock( + return_value=[ + AnalysisJob(id="j1", document_id="d1", document_filename="test.pdf"), + ] + ) resp = client.get("/api/analyses") assert resp.status_code == 200 @@ -146,37 +154,52 @@ class TestAnalysisEndpoints: assert resp.status_code == 404 def test_create_analysis(self, client, mock_analysis_service): - mock_analysis_service.create = AsyncMock(return_value=AnalysisJob( - id="j1", document_id="d1", document_filename="test.pdf", - )) + mock_analysis_service.create = AsyncMock( + return_value=AnalysisJob( + id="j1", + document_id="d1", + document_filename="test.pdf", + ) + ) resp = client.post("/api/analyses", json={"documentId": "d1"}) assert resp.status_code == 200 data = resp.json() assert data["id"] == "j1" assert data["documentId"] == "d1" - mock_analysis_service.create.assert_called_once_with("d1", pipeline_options=None) + mock_analysis_service.create.assert_called_once_with( + "d1", + pipeline_options=None, + chunking_options=None, + ) def test_create_analysis_with_pipeline_options(self, client, mock_analysis_service): - mock_analysis_service.create = AsyncMock(return_value=AnalysisJob( - id="j2", document_id="d1", document_filename="test.pdf", - )) + mock_analysis_service.create = AsyncMock( + return_value=AnalysisJob( + id="j2", + document_id="d1", + document_filename="test.pdf", + ) + ) - resp = client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": { - "do_ocr": False, - "do_table_structure": True, - "table_mode": "fast", - "do_code_enrichment": True, - "do_formula_enrichment": False, - "do_picture_classification": False, - "do_picture_description": False, - "generate_picture_images": True, - "generate_page_images": False, - "images_scale": 2.0, - } - }) + resp = client.post( + "/api/analyses", + json={ + "documentId": "d1", + "pipelineOptions": { + "do_ocr": False, + "do_table_structure": True, + "table_mode": "fast", + "do_code_enrichment": True, + "do_formula_enrichment": False, + "do_picture_classification": False, + "do_picture_description": False, + "generate_picture_images": True, + "generate_page_images": False, + "images_scale": 2.0, + }, + }, + ) assert resp.status_code == 200 data = resp.json() assert data["id"] == "j2" @@ -191,14 +214,17 @@ class TestAnalysisEndpoints: def test_create_analysis_with_partial_pipeline_options(self, client, mock_analysis_service): """Pipeline options should use defaults for unspecified fields.""" - mock_analysis_service.create = AsyncMock(return_value=AnalysisJob( - id="j3", document_id="d1", document_filename="test.pdf", - )) + mock_analysis_service.create = AsyncMock( + return_value=AnalysisJob( + id="j3", + document_id="d1", + document_filename="test.pdf", + ) + ) - resp = client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": {"do_ocr": False} - }) + resp = client.post( + "/api/analyses", json={"documentId": "d1", "pipelineOptions": {"do_ocr": False}} + ) assert resp.status_code == 200 opts = mock_analysis_service.create.call_args.kwargs["pipeline_options"] diff --git a/document-parser/tests/test_bbox.py b/document-parser/tests/test_bbox.py index 9ffe67c..f201a73 100644 --- a/document-parser/tests/test_bbox.py +++ b/document-parser/tests/test_bbox.py @@ -14,6 +14,7 @@ from domain.bbox import EMPTY_BBOX, to_topleft_list # Standard conversions # --------------------------------------------------------------------------- + class TestToTopleftListStandard: """Normal bbox conversions (happy path).""" @@ -29,10 +30,10 @@ class TestToTopleftListStandard: result = to_topleft_list(bbox, page_height=792.0) # After conversion: new_t = 792 - 700 = 92, new_b = 792 - 600 = 192 - assert result[0] == 50 # l unchanged - assert result[1] == pytest.approx(92.0) # t = page_height - old_t - assert result[2] == 200 # r unchanged - assert result[3] == pytest.approx(192.0) # b = page_height - old_b + assert result[0] == 50 # l unchanged + assert result[1] == pytest.approx(92.0) # t = page_height - old_t + assert result[2] == 200 # r unchanged + assert result[3] == pytest.approx(192.0) # b = page_height - old_b def test_result_has_positive_dimensions(self): """Converted bbox should always have b > t (positive height).""" @@ -60,6 +61,7 @@ class TestToTopleftListStandard: # Page format variations # --------------------------------------------------------------------------- + class TestPageFormats: """Verify correct conversion across different page sizes.""" @@ -105,6 +107,7 @@ class TestPageFormats: # Degenerate / edge-case bboxes # --------------------------------------------------------------------------- + class TestDegenerateBboxes: """Bboxes that are invalid or degenerate should return EMPTY_BBOX.""" @@ -151,6 +154,7 @@ class TestDegenerateBboxes: # Precision and boundary values # --------------------------------------------------------------------------- + class TestPrecision: """Floating-point precision and edge values.""" diff --git a/document-parser/tests/test_chunking.py b/document-parser/tests/test_chunking.py new file mode 100644 index 0000000..b429d35 --- /dev/null +++ b/document-parser/tests/test_chunking.py @@ -0,0 +1,292 @@ +"""Tests for chunking feature — domain, schemas, service, and API endpoints.""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +from api.schemas import ChunkingOptionsRequest, ChunkResponse, RechunkRequest +from domain.models import AnalysisJob, AnalysisStatus +from domain.value_objects import ChunkingOptions, ChunkResult +from main import app + +# --------------------------------------------------------------------------- +# Domain: value objects +# --------------------------------------------------------------------------- + + +class TestChunkingOptions: + def test_defaults(self): + opts = ChunkingOptions() + assert opts.chunker_type == "hybrid" + assert opts.max_tokens == 512 + assert opts.merge_peers is True + assert opts.repeat_table_header is True + + def test_custom_values(self): + opts = ChunkingOptions(chunker_type="hierarchical", max_tokens=256, merge_peers=False) + assert opts.chunker_type == "hierarchical" + assert opts.max_tokens == 256 + assert opts.merge_peers is False + + def test_is_default(self): + assert ChunkingOptions().is_default() + assert not ChunkingOptions(max_tokens=256).is_default() + + +class TestChunkResult: + def test_defaults(self): + chunk = ChunkResult(text="hello") + assert chunk.text == "hello" + assert chunk.headings == [] + assert chunk.source_page is None + assert chunk.token_count == 0 + + def test_full_values(self): + chunk = ChunkResult( + text="content", + headings=["Title", "Section"], + source_page=3, + token_count=42, + ) + assert chunk.headings == ["Title", "Section"] + assert chunk.source_page == 3 + assert chunk.token_count == 42 + + def test_serializable(self): + chunk = ChunkResult(text="x", headings=["h1"], source_page=1, token_count=10) + data = asdict(chunk) + assert data == {"text": "x", "headings": ["h1"], "source_page": 1, "token_count": 10} + + +# --------------------------------------------------------------------------- +# Domain: AnalysisJob with chunking fields +# --------------------------------------------------------------------------- + + +class TestAnalysisJobChunking: + def test_default_chunking_fields(self): + job = AnalysisJob() + assert job.document_json is None + assert job.chunks_json is None + + def test_mark_completed_with_chunks(self): + job = AnalysisJob() + job.mark_running() + job.mark_completed( + markdown="# Title", + html="

Title

", + pages_json="[]", + document_json='{"name": "doc"}', + chunks_json='[{"text": "chunk1"}]', + ) + assert job.status == AnalysisStatus.COMPLETED + assert job.document_json == '{"name": "doc"}' + assert job.chunks_json == '[{"text": "chunk1"}]' + + def test_mark_completed_without_chunks(self): + job = AnalysisJob() + job.mark_running() + job.mark_completed(markdown="md", html="html", pages_json="[]") + assert job.document_json is None + assert job.chunks_json is None + + +# --------------------------------------------------------------------------- +# Schemas +# --------------------------------------------------------------------------- + + +class TestChunkingOptionsRequest: + def test_defaults(self): + opts = ChunkingOptionsRequest() + assert opts.chunker_type == "hybrid" + assert opts.max_tokens == 512 + assert opts.merge_peers is True + assert opts.repeat_table_header is True + + def test_custom_values(self): + opts = ChunkingOptionsRequest(chunker_type="hierarchical", max_tokens=1024) + assert opts.chunker_type == "hierarchical" + assert opts.max_tokens == 1024 + + def test_invalid_chunker_type(self): + with pytest.raises(ValueError, match="chunker_type"): + ChunkingOptionsRequest(chunker_type="invalid") + + def test_max_tokens_too_low(self): + with pytest.raises(ValueError, match="max_tokens"): + ChunkingOptionsRequest(max_tokens=10) + + def test_max_tokens_too_high(self): + with pytest.raises(ValueError, match="max_tokens"): + ChunkingOptionsRequest(max_tokens=10000) + + def test_boundary_max_tokens(self): + opts_low = ChunkingOptionsRequest(max_tokens=64) + assert opts_low.max_tokens == 64 + opts_high = ChunkingOptionsRequest(max_tokens=8192) + assert opts_high.max_tokens == 8192 + + +class TestChunkResponse: + def test_serializes_to_camel_case(self): + resp = ChunkResponse(text="hello", headings=["H1"], source_page=1, token_count=5) + data = resp.model_dump(by_alias=True) + assert "sourcePage" in data + assert "tokenCount" in data + assert data["text"] == "hello" + + +class TestRechunkRequest: + def test_parses(self): + req = RechunkRequest(chunkingOptions=ChunkingOptionsRequest(max_tokens=256)) + assert req.chunkingOptions.max_tokens == 256 + + +# --------------------------------------------------------------------------- +# API endpoints +# --------------------------------------------------------------------------- + + +@pytest.fixture +def client(): + return TestClient(app, raise_server_exceptions=False) + + +@pytest.fixture +def mock_analysis_service(client): + mock_svc = MagicMock() + original = getattr(app.state, "analysis_service", None) + app.state.analysis_service = mock_svc + yield mock_svc + app.state.analysis_service = original + + +class TestCreateAnalysisWithChunking: + def test_create_with_chunking_options(self, client, mock_analysis_service): + mock_analysis_service.create = AsyncMock( + return_value=AnalysisJob( + id="j1", + document_id="d1", + document_filename="test.pdf", + ) + ) + + resp = client.post( + "/api/analyses", + json={ + "documentId": "d1", + "chunkingOptions": { + "chunker_type": "hybrid", + "max_tokens": 256, + "merge_peers": False, + }, + }, + ) + assert resp.status_code == 200 + + call_kwargs = mock_analysis_service.create.call_args + chunking = call_kwargs.kwargs["chunking_options"] + assert chunking["chunker_type"] == "hybrid" + assert chunking["max_tokens"] == 256 + assert chunking["merge_peers"] is False + + def test_create_without_chunking_options(self, client, mock_analysis_service): + mock_analysis_service.create = AsyncMock( + return_value=AnalysisJob( + id="j1", + document_id="d1", + document_filename="test.pdf", + ) + ) + + resp = client.post("/api/analyses", json={"documentId": "d1"}) + assert resp.status_code == 200 + + call_kwargs = mock_analysis_service.create.call_args + assert call_kwargs.kwargs["chunking_options"] is None + + def test_response_includes_chunking_fields(self, client, mock_analysis_service): + job = AnalysisJob(id="j1", document_id="d1", document_filename="test.pdf") + job.mark_running() + job.mark_completed( + markdown="# Title", + html="

Title

", + pages_json="[]", + document_json='{"name": "doc"}', + chunks_json=json.dumps([asdict(ChunkResult(text="chunk1", token_count=5))]), + ) + mock_analysis_service.find_by_id = AsyncMock(return_value=job) + + resp = client.get("/api/analyses/j1") + assert resp.status_code == 200 + data = resp.json() + assert data["hasDocumentJson"] is True + assert data["chunksJson"] is not None + chunks = json.loads(data["chunksJson"]) + assert len(chunks) == 1 + assert chunks[0]["text"] == "chunk1" + + +class TestRechunkEndpoint: + def test_rechunk_success(self, client, mock_analysis_service): + mock_analysis_service.rechunk = AsyncMock( + return_value=[ + ChunkResult(text="chunk1", headings=["H1"], source_page=1, token_count=10), + ChunkResult(text="chunk2", headings=["H1", "H2"], source_page=2, token_count=20), + ] + ) + + resp = client.post( + "/api/analyses/j1/rechunk", + json={ + "chunkingOptions": {"chunker_type": "hybrid", "max_tokens": 128}, + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert len(data) == 2 + assert data[0]["text"] == "chunk1" + assert data[0]["sourcePage"] == 1 + assert data[0]["tokenCount"] == 10 + assert data[1]["headings"] == ["H1", "H2"] + + def test_rechunk_not_completed(self, client, mock_analysis_service): + mock_analysis_service.rechunk = AsyncMock( + side_effect=ValueError("Analysis is not completed: j1"), + ) + + resp = client.post( + "/api/analyses/j1/rechunk", + json={ + "chunkingOptions": {"chunker_type": "hybrid"}, + }, + ) + assert resp.status_code == 400 + + def test_rechunk_no_document_json(self, client, mock_analysis_service): + mock_analysis_service.rechunk = AsyncMock( + side_effect=ValueError("No document data available for re-chunking: j1"), + ) + + resp = client.post( + "/api/analyses/j1/rechunk", + json={ + "chunkingOptions": {"chunker_type": "hierarchical"}, + }, + ) + assert resp.status_code == 400 + + def test_rechunk_invalid_chunker_type(self, client, mock_analysis_service): + resp = client.post( + "/api/analyses/j1/rechunk", + json={ + "chunkingOptions": {"chunker_type": "invalid"}, + }, + ) + assert resp.status_code == 422 diff --git a/document-parser/tests/test_pipeline_options.py b/document-parser/tests/test_pipeline_options.py index f402fdc..19ea577 100644 --- a/document-parser/tests/test_pipeline_options.py +++ b/document-parser/tests/test_pipeline_options.py @@ -23,6 +23,7 @@ from infra.local_converter import ( # build_converter — verifies Docling pipeline options are wired correctly # --------------------------------------------------------------------------- + class TestBuildConverter: """Verify that build_converter produces a DocumentConverter with the right PdfPipelineOptions.""" @@ -101,18 +102,20 @@ class TestBuildConverter: assert opts.images_scale == 2.0 def test_all_options_combined(self): - conv = build_converter(ConversionOptions( - do_ocr=False, - do_table_structure=True, - table_mode="fast", - do_code_enrichment=True, - do_formula_enrichment=True, - do_picture_classification=True, - do_picture_description=True, - generate_picture_images=True, - generate_page_images=True, - images_scale=1.5, - )) + conv = build_converter( + ConversionOptions( + do_ocr=False, + do_table_structure=True, + table_mode="fast", + do_code_enrichment=True, + do_formula_enrichment=True, + do_picture_classification=True, + do_picture_description=True, + generate_picture_images=True, + generate_page_images=True, + images_scale=1.5, + ) + ) opts = self._get_pipeline_options(conv) assert opts.do_ocr is False assert opts.do_table_structure is True @@ -130,6 +133,7 @@ class TestBuildConverter: # convert_document — default vs custom converter routing # --------------------------------------------------------------------------- + class TestConvertDocumentRouting: """Verify convert_document uses default converter for default opts, custom otherwise.""" @@ -142,6 +146,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_get_default.return_value = mock_conv @@ -159,6 +164,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -176,6 +182,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -193,6 +200,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -210,6 +218,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -226,6 +235,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -242,6 +252,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -258,6 +269,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -275,6 +287,7 @@ class TestConvertDocumentRouting: mock_result.document.iterate_items.return_value = [] mock_result.document.export_to_markdown.return_value = "" mock_result.document.export_to_html.return_value = "" + mock_result.document.export_to_dict.return_value = {} mock_conv.convert.return_value = mock_result mock_build.return_value = mock_conv @@ -299,30 +312,37 @@ class TestConvertDocumentRouting: # Service layer — pipeline options forwarding # --------------------------------------------------------------------------- + class TestServiceForwardsPipelineOptions: """Verify analysis_service.create and _run_analysis forward pipeline options.""" @pytest.fixture def mock_doc(self): from domain.models import Document + return Document(id="d1", filename="test.pdf", storage_path="/tmp/test.pdf") @pytest.fixture def mock_job(self): from domain.models import AnalysisJob + return AnalysisJob(id="j1", document_id="d1", document_filename="test.pdf") @patch("services.analysis_service.document_repo") @patch("services.analysis_service.analysis_repo") @pytest.mark.asyncio async def test_create_passes_pipeline_options_to_run( - self, mock_analysis_repo, mock_doc_repo, mock_doc, + self, + mock_analysis_repo, + mock_doc_repo, + mock_doc, ): mock_doc_repo.find_by_id = AsyncMock(return_value=mock_doc) mock_analysis_repo.insert = AsyncMock() mock_converter = AsyncMock() from services.analysis_service import AnalysisService + svc = AnalysisService(converter=mock_converter) opts = {"do_ocr": False, "table_mode": "fast"} @@ -335,13 +355,17 @@ class TestServiceForwardsPipelineOptions: @patch("services.analysis_service.analysis_repo") @pytest.mark.asyncio async def test_create_passes_none_when_no_options( - self, mock_analysis_repo, mock_doc_repo, mock_doc, + self, + mock_analysis_repo, + mock_doc_repo, + mock_doc, ): mock_doc_repo.find_by_id = AsyncMock(return_value=mock_doc) mock_analysis_repo.insert = AsyncMock() mock_converter = AsyncMock() from services.analysis_service import AnalysisService + svc = AnalysisService(converter=mock_converter) with patch("services.analysis_service.asyncio.create_task") as mock_task: @@ -352,7 +376,10 @@ class TestServiceForwardsPipelineOptions: @patch("services.analysis_service.document_repo") @pytest.mark.asyncio async def test_run_analysis_forwards_options_to_convert( - self, mock_doc_repo, mock_analysis_repo, mock_job, + self, + mock_doc_repo, + mock_analysis_repo, + mock_job, ): from domain.value_objects import ConversionResult, PageDetail @@ -369,6 +396,7 @@ class TestServiceForwardsPipelineOptions: ) from services.analysis_service import AnalysisService + svc = AnalysisService(converter=mock_converter) opts = { @@ -399,7 +427,10 @@ class TestServiceForwardsPipelineOptions: @patch("services.analysis_service.document_repo") @pytest.mark.asyncio async def test_run_analysis_uses_defaults_when_no_options( - self, mock_doc_repo, mock_analysis_repo, mock_job, + self, + mock_doc_repo, + mock_analysis_repo, + mock_job, ): from domain.value_objects import ConversionResult, PageDetail @@ -416,6 +447,7 @@ class TestServiceForwardsPipelineOptions: ) from services.analysis_service import AnalysisService + svc = AnalysisService(converter=mock_converter) await svc._run_analysis("j1", "/tmp/test.pdf", "test.pdf", None) @@ -429,7 +461,10 @@ class TestServiceForwardsPipelineOptions: @patch("services.analysis_service.document_repo") @pytest.mark.asyncio async def test_run_analysis_marks_failed_on_error( - self, mock_doc_repo, mock_analysis_repo, mock_job, + self, + mock_doc_repo, + mock_analysis_repo, + mock_job, ): mock_analysis_repo.find_by_id = AsyncMock(return_value=mock_job) mock_analysis_repo.update_status = AsyncMock() @@ -438,6 +473,7 @@ class TestServiceForwardsPipelineOptions: mock_converter.convert.side_effect = RuntimeError("Docling crashed") from services.analysis_service import AnalysisService + svc = AnalysisService(converter=mock_converter) await svc._run_analysis("j1", "/tmp/test.pdf", "test.pdf", {"do_ocr": False}) @@ -453,6 +489,7 @@ class TestServiceForwardsPipelineOptions: # API endpoint — full request/response with pipeline options # --------------------------------------------------------------------------- + class TestAnalysisEndpointPipelineOptions: """Integration-level tests for the analysis creation endpoint with pipeline options.""" @@ -461,6 +498,7 @@ class TestAnalysisEndpointPipelineOptions: from fastapi.testclient import TestClient from main import app + return TestClient(app, raise_server_exceptions=False) @pytest.fixture @@ -468,6 +506,7 @@ class TestAnalysisEndpointPipelineOptions: from unittest.mock import MagicMock from main import app + mock = MagicMock() original = getattr(app.state, "analysis_service", None) app.state.analysis_service = mock @@ -476,20 +515,25 @@ class TestAnalysisEndpointPipelineOptions: def test_no_pipeline_options_sends_none(self, client, mock_svc): from domain.models import AnalysisJob + mock_svc.create = AsyncMock(return_value=AnalysisJob(id="j1", document_id="d1")) client.post("/api/analyses", json={"documentId": "d1"}) - mock_svc.create.assert_called_once_with("d1", pipeline_options=None) + mock_svc.create.assert_called_once_with("d1", pipeline_options=None, chunking_options=None) def test_empty_pipeline_options_object_uses_defaults(self, client, mock_svc): from domain.models import AnalysisJob + mock_svc.create = AsyncMock(return_value=AnalysisJob(id="j1", document_id="d1")) - client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": {}, - }) + client.post( + "/api/analyses", + json={ + "documentId": "d1", + "pipelineOptions": {}, + }, + ) opts = mock_svc.create.call_args.kwargs["pipeline_options"] assert opts["do_ocr"] is True @@ -501,12 +545,16 @@ class TestAnalysisEndpointPipelineOptions: def test_partial_pipeline_options_merges_with_defaults(self, client, mock_svc): from domain.models import AnalysisJob + mock_svc.create = AsyncMock(return_value=AnalysisJob(id="j1", document_id="d1")) - client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": {"do_ocr": False, "images_scale": 1.5}, - }) + client.post( + "/api/analyses", + json={ + "documentId": "d1", + "pipelineOptions": {"do_ocr": False, "images_scale": 1.5}, + }, + ) opts = mock_svc.create.call_args.kwargs["pipeline_options"] assert opts["do_ocr"] is False @@ -522,6 +570,7 @@ class TestAnalysisEndpointPipelineOptions: def test_full_pipeline_options(self, client, mock_svc): from domain.models import AnalysisJob + mock_svc.create = AsyncMock(return_value=AnalysisJob(id="j1", document_id="d1")) payload = { @@ -547,18 +596,25 @@ class TestAnalysisEndpointPipelineOptions: assert opts == payload["pipelineOptions"] def test_invalid_pipeline_option_type_rejected(self, client, mock_svc): - resp = client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": {"do_ocr": "not-a-bool"}, - }) + resp = client.post( + "/api/analyses", + json={ + "documentId": "d1", + "pipelineOptions": {"do_ocr": "not-a-bool"}, + }, + ) assert resp.status_code == 422 def test_unknown_pipeline_option_ignored(self, client, mock_svc): from domain.models import AnalysisJob + mock_svc.create = AsyncMock(return_value=AnalysisJob(id="j1", document_id="d1")) - resp = client.post("/api/analyses", json={ - "documentId": "d1", - "pipelineOptions": {"do_ocr": True, "unknown_field": True}, - }) + resp = client.post( + "/api/analyses", + json={ + "documentId": "d1", + "pipelineOptions": {"do_ocr": True, "unknown_field": True}, + }, + ) assert resp.status_code == 200 diff --git a/document-parser/tests/test_repos.py b/document-parser/tests/test_repos.py index 8e89c54..44ebc17 100644 --- a/document-parser/tests/test_repos.py +++ b/document-parser/tests/test_repos.py @@ -1,6 +1,5 @@ """Tests for persistence repositories using a temporary SQLite database.""" - import pytest from domain.models import AnalysisJob, AnalysisStatus, Document diff --git a/document-parser/tests/test_schemas.py b/document-parser/tests/test_schemas.py index c5bb94b..1352802 100644 --- a/document-parser/tests/test_schemas.py +++ b/document-parser/tests/test_schemas.py @@ -1,6 +1,5 @@ """Tests for API schemas — camelCase serialization and validation.""" - import pytest from api.schemas import ( @@ -97,7 +96,10 @@ class TestPipelineOptionsRequest: def test_custom_values(self): opts = PipelineOptionsRequest( - do_ocr=False, table_mode="fast", do_code_enrichment=True, images_scale=2.0, + do_ocr=False, + table_mode="fast", + do_code_enrichment=True, + images_scale=2.0, ) assert opts.do_ocr is False assert opts.table_mode == "fast" diff --git a/document-parser/tests/test_serve_converter.py b/document-parser/tests/test_serve_converter.py index 89490b3..dd4d4ab 100644 --- a/document-parser/tests/test_serve_converter.py +++ b/document-parser/tests/test_serve_converter.py @@ -20,6 +20,7 @@ from infra.serve_converter import ( # Unit tests — form data building # --------------------------------------------------------------------------- + class TestBuildFormData: def test_default_options(self): data = _build_form_data(ConversionOptions()) @@ -36,7 +37,9 @@ class TestBuildFormData: def test_custom_options(self): opts = ConversionOptions( - do_ocr=False, table_mode="fast", images_scale=2.0, + do_ocr=False, + table_mode="fast", + images_scale=2.0, generate_picture_images=True, ) data = _build_form_data(opts) @@ -50,6 +53,7 @@ class TestBuildFormData: # Unit tests — response parsing # --------------------------------------------------------------------------- + class TestParseResponse: def test_minimal_response(self): data = { @@ -82,12 +86,34 @@ class TestParseResponse: { "label": "title", "text": "Title", - "prov": [{"page_no": 1, "bbox": {"l": 10, "t": 20, "r": 200, "b": 40, "coord_origin": "TOPLEFT"}}], + "prov": [ + { + "page_no": 1, + "bbox": { + "l": 10, + "t": 20, + "r": 200, + "b": 40, + "coord_origin": "TOPLEFT", + }, + } + ], }, { "label": "paragraph", "text": "Text", - "prov": [{"page_no": 1, "bbox": {"l": 10, "t": 50, "r": 200, "b": 70, "coord_origin": "TOPLEFT"}}], + "prov": [ + { + "page_no": 1, + "bbox": { + "l": 10, + "t": 50, + "r": 200, + "b": 70, + "coord_origin": "TOPLEFT", + }, + } + ], }, ], "tables": [], @@ -112,7 +138,9 @@ class TestParseResponse: "1": {"size": {"width": 612.0, "height": 792.0}}, "2": {"size": {"width": 595.0, "height": 842.0}}, }, - "texts": [], "tables": [], "pictures": [], + "texts": [], + "tables": [], + "pictures": [], }, } } @@ -135,7 +163,9 @@ class TestParseResponse: def test_json_content_as_string(self): json_doc = { "pages": {"1": {"size": {"width": 612.0, "height": 792.0}}}, - "texts": [], "tables": [], "pictures": [], + "texts": [], + "tables": [], + "pictures": [], } data = { "document": { @@ -171,10 +201,40 @@ class TestParseResponse: "pages": {"1": {"size": {"width": 612.0, "height": 792.0}}}, "texts": [], "tables": [ - {"label": "table", "text": "", "prov": [{"page_no": 1, "bbox": {"l": 10, "t": 10, "r": 300, "b": 200, "coord_origin": "TOPLEFT"}}]}, + { + "label": "table", + "text": "", + "prov": [ + { + "page_no": 1, + "bbox": { + "l": 10, + "t": 10, + "r": 300, + "b": 200, + "coord_origin": "TOPLEFT", + }, + } + ], + }, ], "pictures": [ - {"label": "picture", "text": "", "prov": [{"page_no": 1, "bbox": {"l": 50, "t": 300, "r": 250, "b": 500, "coord_origin": "TOPLEFT"}}]}, + { + "label": "picture", + "text": "", + "prov": [ + { + "page_no": 1, + "bbox": { + "l": 50, + "t": 300, + "r": 250, + "b": 500, + "coord_origin": "TOPLEFT", + }, + } + ], + }, ], }, } @@ -189,14 +249,19 @@ class TestParseResponse: # Unit tests — bbox extraction # --------------------------------------------------------------------------- + class TestExtractBbox: def test_topleft_passthrough(self): - bbox = _extract_bbox({"l": 10, "t": 20, "r": 100, "b": 50, "coord_origin": "TOPLEFT"}, 792.0) + bbox = _extract_bbox( + {"l": 10, "t": 20, "r": 100, "b": 50, "coord_origin": "TOPLEFT"}, 792.0 + ) assert bbox == [10, 20, 100, 50] def test_bottomleft_conversion(self): # In BOTTOMLEFT: t (top of box) has higher y than b (bottom of box) - bbox = _extract_bbox({"l": 10, "t": 772, "r": 100, "b": 742, "coord_origin": "BOTTOMLEFT"}, 792.0) + bbox = _extract_bbox( + {"l": 10, "t": 772, "r": 100, "b": 742, "coord_origin": "BOTTOMLEFT"}, 792.0 + ) # new_top = 792 - 772 = 20, new_bottom = 792 - 742 = 50 assert bbox == [10, 20, 100, 50] @@ -217,9 +282,11 @@ class TestExtractBbox: # Unit tests — label mapping # --------------------------------------------------------------------------- + class TestLabelMapping: def test_known_labels(self): from infra.serve_converter import _LABEL_MAP + assert _LABEL_MAP["table"] == "table" assert _LABEL_MAP["picture"] == "picture" assert _LABEL_MAP["figure"] == "picture" @@ -232,6 +299,7 @@ class TestLabelMapping: def test_unknown_label_defaults_to_text(self): from infra.serve_converter import _LABEL_MAP + assert _LABEL_MAP.get("unknown_thing", "text") == "text" @@ -239,6 +307,7 @@ class TestLabelMapping: # Unit tests — ServeConverter # --------------------------------------------------------------------------- + class TestServeConverter: def test_headers_with_api_key(self): conv = ServeConverter(base_url="http://localhost:5001", api_key="secret") @@ -257,6 +326,7 @@ class TestServeConverter: # Integration tests — HTTP calls (mocked) # --------------------------------------------------------------------------- + class TestServeConverterConvert: @pytest.mark.asyncio async def test_successful_conversion(self, tmp_path): @@ -270,7 +340,22 @@ class TestServeConverterConvert: "json_content": { "pages": {"1": {"size": {"width": 612.0, "height": 792.0}}}, "texts": [ - {"label": "title", "text": "Converted", "prov": [{"page_no": 1, "bbox": {"l": 10, "t": 20, "r": 200, "b": 40, "coord_origin": "TOPLEFT"}}]}, + { + "label": "title", + "text": "Converted", + "prov": [ + { + "page_no": 1, + "bbox": { + "l": 10, + "t": 20, + "r": 200, + "b": 40, + "coord_origin": "TOPLEFT", + }, + } + ], + }, ], "tables": [], "pictures": [], @@ -312,7 +397,9 @@ class TestServeConverterConvert: mock_response = MagicMock() mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "Server Error", request=MagicMock(), response=MagicMock(status_code=500), + "Server Error", + request=MagicMock(), + response=MagicMock(status_code=500), ) mock_client = AsyncMock() @@ -322,8 +409,10 @@ class TestServeConverterConvert: conv = ServeConverter(base_url="http://localhost:5001") - with patch("infra.serve_converter.httpx.AsyncClient", return_value=mock_client), \ - pytest.raises(httpx.HTTPStatusError): + with ( + patch("infra.serve_converter.httpx.AsyncClient", return_value=mock_client), + pytest.raises(httpx.HTTPStatusError), + ): await conv.convert(str(test_file), ConversionOptions()) @pytest.mark.asyncio @@ -358,6 +447,7 @@ class TestServeConverterConvert: # Integration — converter wiring in main.py # --------------------------------------------------------------------------- + class TestConverterWiring: def test_local_engine_builds_local_converter(self): from infra.local_converter import LocalConverter @@ -365,14 +455,19 @@ class TestConverterWiring: with patch("main.settings", Settings(conversion_engine="local")): from main import _build_converter + converter = _build_converter() assert isinstance(converter, LocalConverter) def test_remote_engine_builds_serve_converter(self): from infra.settings import Settings - with patch("main.settings", Settings(conversion_engine="remote", docling_serve_url="http://serve:5001")): + with patch( + "main.settings", + Settings(conversion_engine="remote", docling_serve_url="http://serve:5001"), + ): from main import _build_converter + converter = _build_converter() assert isinstance(converter, ServeConverter) assert converter._base_url == "http://serve:5001" @@ -380,8 +475,16 @@ class TestConverterWiring: def test_remote_engine_passes_api_key(self): from infra.settings import Settings - with patch("main.settings", Settings(conversion_engine="remote", docling_serve_url="http://serve:5001", docling_serve_api_key="my-key")): + with patch( + "main.settings", + Settings( + conversion_engine="remote", + docling_serve_url="http://serve:5001", + docling_serve_api_key="my-key", + ), + ): from main import _build_converter + converter = _build_converter() assert isinstance(converter, ServeConverter) assert converter._api_key == "my-key" diff --git a/frontend/src/features/analysis/api.ts b/frontend/src/features/analysis/api.ts index 85b6da0..902c744 100644 --- a/frontend/src/features/analysis/api.ts +++ b/frontend/src/features/analysis/api.ts @@ -1,20 +1,34 @@ -import type { Analysis, PipelineOptions } from '../../shared/types' +import type { Analysis, Chunk, ChunkingOptions, PipelineOptions } from '../../shared/types' import { apiFetch } from '../../shared/api/http' export function createAnalysis( documentId: string, pipelineOptions: PipelineOptions | null = null, + chunkingOptions: ChunkingOptions | null = null, ): Promise { const body: Record = { documentId } if (pipelineOptions) { body.pipelineOptions = pipelineOptions } + if (chunkingOptions) { + body.chunkingOptions = chunkingOptions + } return apiFetch('/api/analyses', { method: 'POST', body: JSON.stringify(body), }) } +export function rechunkAnalysis( + jobId: string, + chunkingOptions: ChunkingOptions, +): Promise { + return apiFetch(`/api/analyses/${jobId}/rechunk`, { + method: 'POST', + body: JSON.stringify({ chunkingOptions }), + }) +} + export function fetchAnalyses(): Promise { return apiFetch('/api/analyses') } diff --git a/frontend/src/features/analysis/pipelineOptions.test.ts b/frontend/src/features/analysis/pipelineOptions.test.ts index 422ee50..a646e64 100644 --- a/frontend/src/features/analysis/pipelineOptions.test.ts +++ b/frontend/src/features/analysis/pipelineOptions.test.ts @@ -131,7 +131,7 @@ describe('useAnalysisStore — pipeline options forwarding', () => { const store = useAnalysisStore() await store.run('d1') - expect(api.createAnalysis).toHaveBeenCalledWith('d1', null) + expect(api.createAnalysis).toHaveBeenCalledWith('d1', null, null) store.stopPolling() }) @@ -155,7 +155,7 @@ describe('useAnalysisStore — pipeline options forwarding', () => { } await store.run('d1', opts) - expect(api.createAnalysis).toHaveBeenCalledWith('d1', opts) + expect(api.createAnalysis).toHaveBeenCalledWith('d1', opts, null) store.stopPolling() }) @@ -168,7 +168,7 @@ describe('useAnalysisStore — pipeline options forwarding', () => { const opts = { do_ocr: false } await store.run('d1', opts) - expect(api.createAnalysis).toHaveBeenCalledWith('d1', { do_ocr: false }) + expect(api.createAnalysis).toHaveBeenCalledWith('d1', { do_ocr: false }, null) store.stopPolling() }) diff --git a/frontend/src/features/analysis/store.test.ts b/frontend/src/features/analysis/store.test.ts index 6b96672..6c6fb15 100644 --- a/frontend/src/features/analysis/store.test.ts +++ b/frontend/src/features/analysis/store.test.ts @@ -70,7 +70,7 @@ describe('useAnalysisStore', () => { expect(store.currentAnalysis).toEqual(job) expect(store.analyses[0]).toEqual(job) expect(store.running).toBe(true) - expect(api.createAnalysis).toHaveBeenCalledWith('d1', null) + expect(api.createAnalysis).toHaveBeenCalledWith('d1', null, null) // Advance timer to trigger polling await vi.advanceTimersByTimeAsync(2000) @@ -90,7 +90,7 @@ describe('useAnalysisStore', () => { const options = { do_ocr: false, table_mode: 'fast' } await store.run('d1', options) - expect(api.createAnalysis).toHaveBeenCalledWith('d1', options) + expect(api.createAnalysis).toHaveBeenCalledWith('d1', options, null) store.stopPolling() }) diff --git a/frontend/src/features/analysis/store.ts b/frontend/src/features/analysis/store.ts index d66dac5..8891cb7 100644 --- a/frontend/src/features/analysis/store.ts +++ b/frontend/src/features/analysis/store.ts @@ -1,6 +1,6 @@ import { defineStore } from 'pinia' import { ref, computed } from 'vue' -import type { Analysis, Page, PipelineOptions } from '../../shared/types' +import type { Analysis, Chunk, ChunkingOptions, Page, PipelineOptions } from '../../shared/types' import * as api from './api' export const useAnalysisStore = defineStore('analysis', () => { @@ -35,14 +35,44 @@ export const useAnalysisStore = defineStore('analysis', () => { } } + const currentChunks = computed(() => { + if (!currentAnalysis.value?.chunksJson) return [] + try { + return JSON.parse(currentAnalysis.value.chunksJson) as Chunk[] + } catch { + return [] + } + }) + + const rechunking = ref(false) + + async function rechunk(jobId: string, chunkingOptions: ChunkingOptions): Promise { + rechunking.value = true + error.value = null + try { + const chunks = await api.rechunkAnalysis(jobId, chunkingOptions) + if (currentAnalysis.value?.id === jobId) { + currentAnalysis.value = await api.fetchAnalysis(jobId) + } + return chunks + } catch (e) { + error.value = (e as Error).message || 'Failed to rechunk' + console.error('Failed to rechunk', e) + throw e + } finally { + rechunking.value = false + } + } + async function run( documentId: string, pipelineOptions: PipelineOptions | null = null, + chunkingOptions: ChunkingOptions | null = null, ): Promise { running.value = true error.value = null try { - const analysis = await api.createAnalysis(documentId, pipelineOptions) + const analysis = await api.createAnalysis(documentId, pipelineOptions, chunkingOptions) currentAnalysis.value = analysis analyses.value.unshift(analysis) startPolling(analysis.id) @@ -118,11 +148,14 @@ export const useAnalysisStore = defineStore('analysis', () => { analyses, currentAnalysis, currentPages, + currentChunks, running, + rechunking, error, clearError, load, run, + rechunk, select, remove, stopPolling, diff --git a/frontend/src/features/chunking/api.test.ts b/frontend/src/features/chunking/api.test.ts new file mode 100644 index 0000000..d258010 --- /dev/null +++ b/frontend/src/features/chunking/api.test.ts @@ -0,0 +1,52 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { rechunkAnalysis, createAnalysis } from '../analysis/api' + +vi.mock('../../shared/api/http', () => ({ + apiFetch: vi.fn(), +})) + +import { apiFetch } from '../../shared/api/http' + +describe('chunking API', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it('createAnalysis sends chunkingOptions when provided', async () => { + const job = { id: '1', documentId: 'doc-1', status: 'PENDING' } + apiFetch.mockResolvedValue(job) + + const chunkingOpts = { chunker_type: 'hybrid' as const, max_tokens: 256 } + await createAnalysis('doc-1', null, chunkingOpts) + + expect(apiFetch).toHaveBeenCalledWith('/api/analyses', { + method: 'POST', + body: JSON.stringify({ documentId: 'doc-1', chunkingOptions: chunkingOpts }), + }) + }) + + it('createAnalysis omits chunkingOptions when null', async () => { + apiFetch.mockResolvedValue({ id: '1' }) + + await createAnalysis('doc-1', null, null) + + expect(apiFetch).toHaveBeenCalledWith('/api/analyses', { + method: 'POST', + body: JSON.stringify({ documentId: 'doc-1' }), + }) + }) + + it('rechunkAnalysis sends POST to rechunk endpoint', async () => { + const chunks = [{ text: 'chunk1', headings: [], sourcePage: 1, tokenCount: 10 }] + apiFetch.mockResolvedValue(chunks) + + const opts = { chunker_type: 'hybrid' as const, max_tokens: 512 } + const result = await rechunkAnalysis('job-1', opts) + + expect(apiFetch).toHaveBeenCalledWith('/api/analyses/job-1/rechunk', { + method: 'POST', + body: JSON.stringify({ chunkingOptions: opts }), + }) + expect(result).toEqual(chunks) + }) +}) diff --git a/frontend/src/features/chunking/index.ts b/frontend/src/features/chunking/index.ts new file mode 100644 index 0000000..a8d8572 --- /dev/null +++ b/frontend/src/features/chunking/index.ts @@ -0,0 +1 @@ +export { default as ChunkPanel } from './ui/ChunkPanel.vue' diff --git a/frontend/src/features/chunking/store.test.ts b/frontend/src/features/chunking/store.test.ts new file mode 100644 index 0000000..80b3fa3 --- /dev/null +++ b/frontend/src/features/chunking/store.test.ts @@ -0,0 +1,135 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { setActivePinia, createPinia } from 'pinia' +import { useAnalysisStore } from '../analysis/store' + +vi.mock('../analysis/api', () => ({ + createAnalysis: vi.fn(), + fetchAnalyses: vi.fn().mockResolvedValue([]), + fetchAnalysis: vi.fn(), + deleteAnalysis: vi.fn(), + rechunkAnalysis: vi.fn(), +})) + +import * as api from '../analysis/api' + +describe('analysis store — chunking', () => { + beforeEach(() => { + setActivePinia(createPinia()) + vi.clearAllMocks() + }) + + it('currentChunks parses chunksJson from current analysis', () => { + const store = useAnalysisStore() + const chunks = [ + { text: 'chunk1', headings: ['H1'], sourcePage: 1, tokenCount: 10 }, + { text: 'chunk2', headings: [], sourcePage: 2, tokenCount: 20 }, + ] + store.currentAnalysis = { + id: 'j1', + documentId: 'd1', + documentFilename: null, + status: 'COMPLETED', + contentMarkdown: null, + contentHtml: null, + pagesJson: null, + chunksJson: JSON.stringify(chunks), + hasDocumentJson: true, + errorMessage: null, + startedAt: null, + completedAt: null, + createdAt: '2024-01-01', + } + expect(store.currentChunks).toEqual(chunks) + }) + + it('currentChunks returns empty array when no chunksJson', () => { + const store = useAnalysisStore() + store.currentAnalysis = { + id: 'j1', + documentId: 'd1', + documentFilename: null, + status: 'COMPLETED', + contentMarkdown: null, + contentHtml: null, + pagesJson: null, + chunksJson: null, + hasDocumentJson: false, + errorMessage: null, + startedAt: null, + completedAt: null, + createdAt: '2024-01-01', + } + expect(store.currentChunks).toEqual([]) + }) + + it('rechunk calls API and refreshes analysis', async () => { + const store = useAnalysisStore() + const chunks = [{ text: 'c1', headings: [], sourcePage: 1, tokenCount: 5 }] + vi.mocked(api.rechunkAnalysis).mockResolvedValue(chunks) + vi.mocked(api.fetchAnalysis).mockResolvedValue({ + id: 'j1', + documentId: 'd1', + documentFilename: null, + status: 'COMPLETED', + contentMarkdown: null, + contentHtml: null, + pagesJson: null, + chunksJson: JSON.stringify(chunks), + hasDocumentJson: true, + errorMessage: null, + startedAt: null, + completedAt: null, + createdAt: '2024-01-01', + }) + + store.currentAnalysis = { + id: 'j1', + documentId: 'd1', + documentFilename: null, + status: 'COMPLETED', + contentMarkdown: null, + contentHtml: null, + pagesJson: null, + chunksJson: null, + hasDocumentJson: true, + errorMessage: null, + startedAt: null, + completedAt: null, + createdAt: '2024-01-01', + } + + const result = await store.rechunk('j1', { chunker_type: 'hybrid', max_tokens: 256 }) + + expect(api.rechunkAnalysis).toHaveBeenCalledWith('j1', { + chunker_type: 'hybrid', + max_tokens: 256, + }) + expect(result).toEqual(chunks) + expect(store.rechunking).toBe(false) + }) + + it('run passes chunkingOptions to API', async () => { + const store = useAnalysisStore() + vi.mocked(api.createAnalysis).mockResolvedValue({ + id: 'j1', + documentId: 'd1', + documentFilename: null, + status: 'PENDING', + contentMarkdown: null, + contentHtml: null, + pagesJson: null, + chunksJson: null, + hasDocumentJson: false, + errorMessage: null, + startedAt: null, + completedAt: null, + createdAt: '2024-01-01', + }) + + await store.run('d1', null, { chunker_type: 'hierarchical' }) + + expect(api.createAnalysis).toHaveBeenCalledWith('d1', null, { + chunker_type: 'hierarchical', + }) + }) +}) diff --git a/frontend/src/features/chunking/ui/ChunkPanel.vue b/frontend/src/features/chunking/ui/ChunkPanel.vue new file mode 100644 index 0000000..b93bcde --- /dev/null +++ b/frontend/src/features/chunking/ui/ChunkPanel.vue @@ -0,0 +1,336 @@ + + + + + diff --git a/frontend/src/pages/StudioPage.vue b/frontend/src/pages/StudioPage.vue index aedc9a4..fd802d0 100644 --- a/frontend/src/pages/StudioPage.vue +++ b/frontend/src/pages/StudioPage.vue @@ -37,6 +37,15 @@ > {{ t('studio.verify') }} +
@@ -280,6 +289,11 @@ @highlight-element="highlightedElementIndex = $event" />
+ + +
+ +
@@ -293,6 +307,8 @@ import { useAnalysisStore } from '../features/analysis/store' import { DocumentUpload, DocumentList } from '../features/document/index' import { ResultTabs } from '../features/analysis/index' import BboxOverlay from '../features/analysis/ui/BboxOverlay.vue' +import { ChunkPanel } from '../features/chunking' +import { useFeatureFlag } from '../features/feature-flags' import { getPreviewUrl } from '../features/document/api' import { useI18n } from '../shared/i18n' import type { PipelineOptions } from '../shared/types' @@ -302,6 +318,7 @@ const router = useRouter() const documentStore = useDocumentStore() const analysisStore = useAnalysisStore() const { t } = useI18n() +const chunkingEnabled = useFeatureFlag('chunking') const mode = ref('configurer') const currentPage = ref(1) diff --git a/frontend/src/shared/i18n.ts b/frontend/src/shared/i18n.ts index 92b314b..6c4b759 100644 --- a/frontend/src/shared/i18n.ts +++ b/frontend/src/shared/i18n.ts @@ -95,6 +95,18 @@ const messages: Messages = { 'history.emptyDocs': 'Aucun document. Importez un document depuis le Studio.', 'history.open': 'Ouvrir', + // Chunking + 'studio.prepare': 'Préparer', + 'chunking.settings': 'Chunking', + 'chunking.chunkerType': 'Type de chunker', + 'chunking.maxTokens': 'Tokens max', + 'chunking.mergePeers': 'Fusionner les pairs', + 'chunking.repeatTableHeader': 'Répéter en-têtes tableaux', + 'chunking.run': 'Chunker', + 'chunking.chunking': 'Chunking...', + 'chunking.chunks': 'chunks', + 'chunking.noChunks': 'Lancez le chunking pour préparer les segments.', + // Settings 'settings.title': 'Paramètres', 'settings.apiUrl': 'API URL', @@ -185,6 +197,17 @@ const messages: Messages = { 'history.emptyDocs': 'No documents yet. Upload a document from the Studio.', 'history.open': 'Open', + 'studio.prepare': 'Prepare', + 'chunking.settings': 'Chunking', + 'chunking.chunkerType': 'Chunker type', + 'chunking.maxTokens': 'Max tokens', + 'chunking.mergePeers': 'Merge peers', + 'chunking.repeatTableHeader': 'Repeat table headers', + 'chunking.run': 'Chunk', + 'chunking.chunking': 'Chunking...', + 'chunking.chunks': 'chunks', + 'chunking.noChunks': 'Run chunking to prepare segments.', + 'settings.title': 'Settings', 'settings.apiUrl': 'API URL', 'settings.version': 'Version', diff --git a/frontend/src/shared/types.ts b/frontend/src/shared/types.ts index 8035836..5aa7968 100644 --- a/frontend/src/shared/types.ts +++ b/frontend/src/shared/types.ts @@ -30,12 +30,28 @@ export interface Analysis { contentMarkdown: string | null contentHtml: string | null pagesJson: string | null + chunksJson: string | null + hasDocumentJson: boolean errorMessage: string | null startedAt: string | null completedAt: string | null createdAt: string } +export interface ChunkingOptions { + chunker_type?: 'hybrid' | 'hierarchical' + max_tokens?: number + merge_peers?: boolean + repeat_table_header?: boolean +} + +export interface Chunk { + text: string + headings: string[] + sourcePage: number | null + tokenCount: number +} + export interface PageElement { type: string bbox: [number, number, number, number]