From a9517d38eb368ab512f881adc023ee21c2ecb2f2 Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Thu, 2 Apr 2026 12:05:28 +0200 Subject: [PATCH] Add chunking service orchestration, API endpoints, and wiring AnalysisService gains rechunk() and inline chunking during conversion. ChunkingOptionsRequest/ChunkResponse schemas, POST rechunk endpoint, and conditional chunker injection in main.py (local engine only). --- document-parser/api/analyses.py | 32 +++++++- document-parser/api/schemas.py | 39 ++++++++++ document-parser/main.py | 14 ++++ document-parser/services/analysis_service.py | 78 ++++++++++++++++++-- 4 files changed, 153 insertions(+), 10 deletions(-) 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/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/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)