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).
This commit is contained in:
parent
4b1ec364f4
commit
a9517d38eb
4 changed files with 153 additions and 10 deletions
|
|
@ -7,7 +7,7 @@ from typing import Annotated
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
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
|
from services.analysis_service import AnalysisService
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
@ -30,6 +30,8 @@ def _to_response(job) -> AnalysisResponse:
|
||||||
content_markdown=job.content_markdown,
|
content_markdown=job.content_markdown,
|
||||||
content_html=job.content_html,
|
content_html=job.content_html,
|
||||||
pages_json=job.pages_json,
|
pages_json=job.pages_json,
|
||||||
|
chunks_json=job.chunks_json,
|
||||||
|
has_document_json=job.document_json is not None,
|
||||||
error_message=job.error_message,
|
error_message=job.error_message,
|
||||||
started_at=str(job.started_at) if job.started_at else None,
|
started_at=str(job.started_at) if job.started_at else None,
|
||||||
completed_at=str(job.completed_at) if job.completed_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:
|
if body.pipelineOptions:
|
||||||
pipeline_opts = body.pipelineOptions.model_dump()
|
pipeline_opts = body.pipelineOptions.model_dump()
|
||||||
|
|
||||||
|
chunking_opts = None
|
||||||
|
if body.chunkingOptions:
|
||||||
|
chunking_opts = body.chunkingOptions.model_dump()
|
||||||
|
|
||||||
try:
|
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:
|
except ValueError as e:
|
||||||
raise HTTPException(status_code=404, detail=str(e)) from 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)
|
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)
|
@router.delete("/{job_id}", status_code=204)
|
||||||
async def delete_analysis(job_id: str, service: ServiceDep):
|
async def delete_analysis(job_id: str, service: ServiceDep):
|
||||||
"""Delete an analysis job."""
|
"""Delete an analysis job."""
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,7 @@ def _to_camel(name: str) -> str:
|
||||||
|
|
||||||
class _CamelModel(BaseModel):
|
class _CamelModel(BaseModel):
|
||||||
"""Base model that serializes field names to camelCase."""
|
"""Base model that serializes field names to camelCase."""
|
||||||
|
|
||||||
model_config = ConfigDict(
|
model_config = ConfigDict(
|
||||||
alias_generator=_to_camel,
|
alias_generator=_to_camel,
|
||||||
populate_by_name=True,
|
populate_by_name=True,
|
||||||
|
|
@ -42,6 +43,8 @@ class AnalysisResponse(_CamelModel):
|
||||||
content_markdown: str | None = None
|
content_markdown: str | None = None
|
||||||
content_html: str | None = None
|
content_html: str | None = None
|
||||||
pages_json: str | None = None
|
pages_json: str | None = None
|
||||||
|
chunks_json: str | None = None
|
||||||
|
has_document_json: bool = False
|
||||||
error_message: str | None = None
|
error_message: str | None = None
|
||||||
started_at: str | datetime | None = None
|
started_at: str | datetime | None = None
|
||||||
completed_at: str | datetime | None = None
|
completed_at: str | datetime | None = None
|
||||||
|
|
@ -50,6 +53,7 @@ class AnalysisResponse(_CamelModel):
|
||||||
|
|
||||||
class PipelineOptionsRequest(BaseModel):
|
class PipelineOptionsRequest(BaseModel):
|
||||||
"""Docling pipeline configuration options."""
|
"""Docling pipeline configuration options."""
|
||||||
|
|
||||||
do_ocr: bool = True
|
do_ocr: bool = True
|
||||||
do_table_structure: bool = True
|
do_table_structure: bool = True
|
||||||
table_mode: str = "accurate" # "accurate" or "fast"
|
table_mode: str = "accurate" # "accurate" or "fast"
|
||||||
|
|
@ -76,6 +80,41 @@ class PipelineOptionsRequest(BaseModel):
|
||||||
return v
|
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):
|
class CreateAnalysisRequest(BaseModel):
|
||||||
documentId: str # camelCase to match existing frontend contract
|
documentId: str # camelCase to match existing frontend contract
|
||||||
pipelineOptions: PipelineOptionsRequest | None = None
|
pipelineOptions: PipelineOptionsRequest | None = None
|
||||||
|
chunkingOptions: ChunkingOptionsRequest | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class RechunkRequest(BaseModel):
|
||||||
|
chunkingOptions: ChunkingOptionsRequest
|
||||||
|
|
|
||||||
|
|
@ -40,6 +40,7 @@ def _build_converter():
|
||||||
"""Build the converter adapter based on configuration."""
|
"""Build the converter adapter based on configuration."""
|
||||||
if settings.conversion_engine == "remote":
|
if settings.conversion_engine == "remote":
|
||||||
from infra.serve_converter import ServeConverter
|
from infra.serve_converter import ServeConverter
|
||||||
|
|
||||||
logger.info("Using remote Docling Serve at %s", settings.docling_serve_url)
|
logger.info("Using remote Docling Serve at %s", settings.docling_serve_url)
|
||||||
return ServeConverter(
|
return ServeConverter(
|
||||||
base_url=settings.docling_serve_url,
|
base_url=settings.docling_serve_url,
|
||||||
|
|
@ -47,14 +48,26 @@ def _build_converter():
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
from infra.local_converter import LocalConverter
|
from infra.local_converter import LocalConverter
|
||||||
|
|
||||||
logger.info("Using local Docling converter")
|
logger.info("Using local Docling converter")
|
||||||
return LocalConverter()
|
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:
|
def _build_analysis_service() -> AnalysisService:
|
||||||
converter = _build_converter()
|
converter = _build_converter()
|
||||||
|
chunker = _build_chunker()
|
||||||
return AnalysisService(
|
return AnalysisService(
|
||||||
converter=converter,
|
converter=converter,
|
||||||
|
chunker=chunker,
|
||||||
conversion_timeout=settings.conversion_timeout,
|
conversion_timeout=settings.conversion_timeout,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -63,6 +76,7 @@ def _build_analysis_service() -> AnalysisService:
|
||||||
# FastAPI app
|
# FastAPI app
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
await init_db()
|
await init_db()
|
||||||
|
|
|
||||||
|
|
@ -13,24 +13,46 @@ import logging
|
||||||
from dataclasses import asdict
|
from dataclasses import asdict
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from domain.models import AnalysisJob
|
from domain.models import AnalysisJob, AnalysisStatus
|
||||||
from domain.value_objects import ConversionOptions, ConversionResult
|
from domain.value_objects import ChunkingOptions, ChunkResult, ConversionOptions, ConversionResult
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from domain.ports import DocumentConverter
|
from domain.ports import DocumentChunker, DocumentConverter
|
||||||
from persistence import analysis_repo, document_repo
|
from persistence import analysis_repo, document_repo
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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:
|
class AnalysisService:
|
||||||
"""Orchestrates document analysis using an injected converter."""
|
"""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._converter = converter
|
||||||
|
self._chunker = chunker
|
||||||
self._conversion_timeout = conversion_timeout
|
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."""
|
"""Create a new analysis job and launch background processing."""
|
||||||
doc = await document_repo.find_by_id(document_id)
|
doc = await document_repo.find_by_id(document_id)
|
||||||
if not doc:
|
if not doc:
|
||||||
|
|
@ -41,7 +63,13 @@ class AnalysisService:
|
||||||
await analysis_repo.insert(job)
|
await analysis_repo.insert(job)
|
||||||
|
|
||||||
task = asyncio.create_task(
|
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))
|
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:
|
async def delete(self, job_id: str) -> bool:
|
||||||
return await analysis_repo.delete(job_id)
|
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(
|
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:
|
) -> None:
|
||||||
"""Background task: run conversion and update job status."""
|
"""Background task: run conversion and optionally chunk."""
|
||||||
try:
|
try:
|
||||||
job = await analysis_repo.find_by_id(job_id)
|
job = await analysis_repo.find_by_id(job_id)
|
||||||
if not job:
|
if not job:
|
||||||
|
|
@ -79,10 +132,19 @@ class AnalysisService:
|
||||||
|
|
||||||
pages_json = json.dumps([asdict(p) for p in result.pages])
|
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(
|
job.mark_completed(
|
||||||
markdown=result.content_markdown,
|
markdown=result.content_markdown,
|
||||||
html=result.content_html,
|
html=result.content_html,
|
||||||
pages_json=pages_json,
|
pages_json=pages_json,
|
||||||
|
document_json=result.document_json,
|
||||||
|
chunks_json=chunks_json,
|
||||||
)
|
)
|
||||||
await analysis_repo.update_status(job)
|
await analysis_repo.update_status(job)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue